diff --git a/core/trino-main/src/main/java/io/trino/operator/ChangeOnlyUpdatedColumnsMergeProcessor.java b/core/trino-main/src/main/java/io/trino/operator/ChangeOnlyUpdatedColumnsMergeProcessor.java index ee290ea6981c..5a600bf63674 100644 --- a/core/trino-main/src/main/java/io/trino/operator/ChangeOnlyUpdatedColumnsMergeProcessor.java +++ b/core/trino-main/src/main/java/io/trino/operator/ChangeOnlyUpdatedColumnsMergeProcessor.java @@ -14,15 +14,15 @@ package io.trino.operator; import io.trino.spi.Page; +import io.trino.spi.PageBlockUtil; import io.trino.spi.block.Block; -import io.trino.spi.block.ColumnarRow; import io.trino.spi.block.RunLengthEncodedBlock; import java.util.ArrayList; import java.util.List; import static com.google.common.base.Preconditions.checkArgument; -import static io.trino.spi.block.ColumnarRow.toColumnarRow; +import static io.trino.spi.block.RowBlock.getRowFieldsFromBlock; import static io.trino.spi.predicate.Utils.nativeValueToBlock; import static io.trino.spi.type.TinyintType.TINYINT; import static java.util.Objects.requireNonNull; @@ -67,18 +67,16 @@ public Page transformPage(Page inputPage) // TODO: Check with Karol to see if we can get empty pages checkArgument(positionCount > 0, "positionCount should be > 0, but is %s", positionCount); - ColumnarRow mergeRow = toColumnarRow(inputPage.getBlock(mergeRowChannel)); - checkArgument(!mergeRow.mayHaveNull(), "The mergeRow may not have null rows"); - - // We've verified that the mergeRow block has no null rows, so it's okay to get the field blocks - - List builder = new ArrayList<>(dataColumnChannels.size() + 3); - + Block mergeRow = inputPage.getBlock(mergeRowChannel).getLoadedBlock(); + List fields = getRowFieldsFromBlock(mergeRow); + List builder = new ArrayList<>(dataColumnChannels.size() + 4); for (int channel : dataColumnChannels) { - builder.add(mergeRow.getField(channel)); + builder.add(fields.get(channel)); } - Block operationChannelBlock = mergeRow.getField(mergeRow.getFieldCount() - 2); + Block operationChannelBlock = fields.get(fields.size() - 2); builder.add(operationChannelBlock); + Block caseNumberChannelBlock = fields.get(fields.size() - 1); + builder.add(caseNumberChannelBlock); builder.add(inputPage.getBlock(rowIdChannel)); builder.add(RunLengthEncodedBlock.create(INSERT_FROM_UPDATE_BLOCK, positionCount)); @@ -86,7 +84,7 @@ public Page transformPage(Page inputPage) int defaultCaseCount = 0; for (int position = 0; position < positionCount; position++) { - if (TINYINT.getByte(operationChannelBlock, position) == DEFAULT_CASE_OPERATION_NUMBER) { + if (mergeRow.isNull(position)) { defaultCaseCount++; } } @@ -97,7 +95,7 @@ public Page transformPage(Page inputPage) int usedCases = 0; int[] positions = new int[positionCount - defaultCaseCount]; for (int position = 0; position < positionCount; position++) { - if (TINYINT.getByte(operationChannelBlock, position) != DEFAULT_CASE_OPERATION_NUMBER) { + if (!mergeRow.isNull(position)) { positions[usedCases] = position; usedCases++; } @@ -105,6 +103,6 @@ public Page transformPage(Page inputPage) checkArgument(usedCases + defaultCaseCount == positionCount, "usedCases (%s) + defaultCaseCount (%s) != positionCount (%s)", usedCases, defaultCaseCount, positionCount); - return result.getPositions(positions, 0, usedCases); + return PageBlockUtil.getPositions(result, positions, 0, usedCases); } } diff --git a/core/trino-main/src/main/java/io/trino/operator/DeleteAndInsertMergeProcessor.java b/core/trino-main/src/main/java/io/trino/operator/DeleteAndInsertMergeProcessor.java index 700338d5a8c2..3072163b9329 100644 --- a/core/trino-main/src/main/java/io/trino/operator/DeleteAndInsertMergeProcessor.java +++ b/core/trino-main/src/main/java/io/trino/operator/DeleteAndInsertMergeProcessor.java @@ -19,19 +19,19 @@ import io.trino.spi.PageBuilder; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; -import io.trino.spi.block.ColumnarRow; import io.trino.spi.type.Type; import java.util.List; import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.base.Verify.verify; -import static io.trino.spi.block.ColumnarRow.toColumnarRow; +import static io.trino.spi.block.RowBlock.getRowFieldsFromBlock; import static io.trino.spi.connector.ConnectorMergeSink.DELETE_OPERATION_NUMBER; import static io.trino.spi.connector.ConnectorMergeSink.INSERT_OPERATION_NUMBER; import static io.trino.spi.connector.ConnectorMergeSink.UPDATE_DELETE_OPERATION_NUMBER; import static io.trino.spi.connector.ConnectorMergeSink.UPDATE_INSERT_OPERATION_NUMBER; import static io.trino.spi.connector.ConnectorMergeSink.UPDATE_OPERATION_NUMBER; +import static io.trino.spi.type.IntegerType.INTEGER; import static io.trino.spi.type.TinyintType.TINYINT; import static java.util.Objects.requireNonNull; @@ -100,8 +100,9 @@ public Page transformPage(Page inputPage) int originalPositionCount = inputPage.getPositionCount(); checkArgument(originalPositionCount > 0, "originalPositionCount should be > 0, but is %s", originalPositionCount); - ColumnarRow mergeRow = toColumnarRow(inputPage.getBlock(mergeRowChannel)); - Block operationChannelBlock = mergeRow.getField(mergeRow.getFieldCount() - 2); + Block mergeRow = inputPage.getBlock(mergeRowChannel); + List fields = getRowFieldsFromBlock(mergeRow); + Block operationChannelBlock = fields.get(fields.size() - 2); int updatePositions = 0; int insertPositions = 0; @@ -123,6 +124,7 @@ public Page transformPage(Page inputPage) List pageTypes = ImmutableList.builder() .addAll(dataColumnTypes) .add(TINYINT) + .add(INTEGER) .add(rowIdType) .add(TINYINT) .build(); @@ -137,7 +139,7 @@ public Page transformPage(Page inputPage) } // Insert and update because both create an insert row if (operation == INSERT_OPERATION_NUMBER || operation == UPDATE_OPERATION_NUMBER) { - addInsertRow(pageBuilder, mergeRow, position, operation != INSERT_OPERATION_NUMBER); + addInsertRow(pageBuilder, fields, position, operation != INSERT_OPERATION_NUMBER); } } } @@ -170,33 +172,39 @@ private void addDeleteRow(PageBuilder pageBuilder, Page originalPage, int positi // Add the operation column == deleted TINYINT.writeLong(pageBuilder.getBlockBuilder(dataColumnChannels.size()), causedByUpdate ? UPDATE_DELETE_OPERATION_NUMBER : DELETE_OPERATION_NUMBER); + // Add the dummy case number, delete and insert won't use it, use -1 to mark it shouldn't be used + INTEGER.writeLong(pageBuilder.getBlockBuilder(dataColumnChannels.size() + 1), -1); + // Copy row ID column - rowIdType.appendTo(originalPage.getBlock(rowIdChannel), position, pageBuilder.getBlockBuilder(dataColumnChannels.size() + 1)); + rowIdType.appendTo(originalPage.getBlock(rowIdChannel), position, pageBuilder.getBlockBuilder(dataColumnChannels.size() + 2)); // Write 0, meaning this row is not an insert derived from an update - TINYINT.writeLong(pageBuilder.getBlockBuilder(dataColumnChannels.size() + 2), 0); + TINYINT.writeLong(pageBuilder.getBlockBuilder(dataColumnChannels.size() + 3), 0); pageBuilder.declarePosition(); } - private void addInsertRow(PageBuilder pageBuilder, ColumnarRow mergeCaseBlock, int position, boolean causedByUpdate) + private void addInsertRow(PageBuilder pageBuilder, List fields, int position, boolean causedByUpdate) { // Copy the values from the merge block for (int targetChannel : dataColumnChannels) { Type columnType = dataColumnTypes.get(targetChannel); BlockBuilder targetBlock = pageBuilder.getBlockBuilder(targetChannel); // The value comes from that column of the page - columnType.appendTo(mergeCaseBlock.getField(targetChannel), position, targetBlock); + columnType.appendTo(fields.get(targetChannel), position, targetBlock); } // Add the operation column == insert TINYINT.writeLong(pageBuilder.getBlockBuilder(dataColumnChannels.size()), causedByUpdate ? UPDATE_INSERT_OPERATION_NUMBER : INSERT_OPERATION_NUMBER); + // Add the dummy case number, delete and insert won't use it + INTEGER.writeLong(pageBuilder.getBlockBuilder(dataColumnChannels.size() + 1), 0); + // Add null row ID column - pageBuilder.getBlockBuilder(dataColumnChannels.size() + 1).appendNull(); + pageBuilder.getBlockBuilder(dataColumnChannels.size() + 2).appendNull(); // Write 1 if this row is an insert derived from an update, 0 otherwise - TINYINT.writeLong(pageBuilder.getBlockBuilder(dataColumnChannels.size() + 2), causedByUpdate ? 1 : 0); + TINYINT.writeLong(pageBuilder.getBlockBuilder(dataColumnChannels.size() + 3), causedByUpdate ? 1 : 0); pageBuilder.declarePosition(); } diff --git a/core/trino-main/src/main/java/io/trino/operator/output/RowPositionsAppender.java b/core/trino-main/src/main/java/io/trino/operator/output/RowPositionsAppender.java index 00b501177df1..ecb986e6ec25 100644 --- a/core/trino-main/src/main/java/io/trino/operator/output/RowPositionsAppender.java +++ b/core/trino-main/src/main/java/io/trino/operator/output/RowPositionsAppender.java @@ -27,13 +27,14 @@ import static io.airlift.slice.SizeOf.sizeOf; import static io.trino.operator.output.PositionsAppenderUtil.calculateBlockResetSize; import static io.trino.operator.output.PositionsAppenderUtil.calculateNewArraySize; -import static io.trino.spi.block.RowBlock.fromFieldBlocks; +import static io.trino.spi.block.RowBlock.fromNotNullSuppressedFieldBlocks; import static java.util.Objects.requireNonNull; public class RowPositionsAppender implements PositionsAppender { private static final int INSTANCE_SIZE = instanceSize(RowPositionsAppender.class); + private final RowType type; private final PositionsAppender[] fieldAppenders; private int initialEntryCount; private boolean initialized; @@ -55,11 +56,12 @@ public static RowPositionsAppender createRowAppender( for (int i = 0; i < fields.length; i++) { fields[i] = positionsAppenderFactory.create(type.getFields().get(i).getType(), expectedPositions, maxPageSizeInBytes); } - return new RowPositionsAppender(fields, expectedPositions); + return new RowPositionsAppender(type, fields, expectedPositions); } - private RowPositionsAppender(PositionsAppender[] fieldAppenders, int expectedPositions) + private RowPositionsAppender(RowType type, PositionsAppender[] fieldAppenders, int expectedPositions) { + this.type = type; this.fieldAppenders = requireNonNull(fieldAppenders, "fields is null"); this.initialEntryCount = expectedPositions; resetSize(); @@ -88,7 +90,7 @@ public void append(IntArrayList positions, Block block) List fieldBlocks = sourceRowBlock.getChildren(); for (int i = 0; i < fieldAppenders.length; i++) { - fieldAppenders[i].append(nonNullPositions, fieldBlocks.get(i)); + fieldAppenders[i].append(positions, fieldBlocks.get(i)); } } else if (allPositionsNull(positions, block)) { @@ -96,6 +98,9 @@ else if (allPositionsNull(positions, block)) { // append positions.size() nulls Arrays.fill(rowIsNull, positionCount, positionCount + positions.size(), true); hasNullRow = true; + for (int i = 0; i < fieldAppenders.length; i++) { + fieldAppenders[i].append(positions, block); + } } else { throw new IllegalArgumentException("unsupported block type: " + block); @@ -113,6 +118,9 @@ public void appendRle(Block value, int rlePositionCount) if (sourceRowBlock.isNull(0)) { // append rlePositionCount nulls Arrays.fill(rowIsNull, positionCount, positionCount + rlePositionCount, true); + for (int i = 0; i < fieldAppenders.length; i++) { + fieldAppenders[i].appendRle(value.getSingleValueBlock(0), rlePositionCount); + } hasNullRow = true; } else { @@ -128,6 +136,9 @@ public void appendRle(Block value, int rlePositionCount) else if (value.isNull(0)) { // append rlePositionCount nulls Arrays.fill(rowIsNull, positionCount, positionCount + rlePositionCount, true); + for (int i = 0; i < fieldAppenders.length; i++) { + fieldAppenders[i].appendRle(value.getSingleValueBlock(0), rlePositionCount); + } hasNullRow = true; } else { @@ -145,6 +156,9 @@ public void append(int position, Block value) if (sourceRowBlock.isNull(position)) { rowIsNull[positionCount] = true; hasNullRow = true; + for (int i = 0; i < fieldAppenders.length; i++) { + fieldAppenders[i].append(position, value); + } } else { // append not null row value @@ -159,6 +173,9 @@ public void append(int position, Block value) else if (value.isNull(position)) { rowIsNull[positionCount] = true; hasNullRow = true; + for (int i = 0; i < fieldAppenders.length; i++) { + fieldAppenders[i].append(position, value); + } } else { throw new IllegalArgumentException("unsupported block type: " + value); @@ -176,10 +193,10 @@ public Block build() } Block result; if (hasNonNullRow) { - result = fromFieldBlocks(positionCount, hasNullRow ? Optional.of(rowIsNull) : Optional.empty(), fieldBlocks); + result = fromNotNullSuppressedFieldBlocks(positionCount, hasNullRow ? Optional.of(rowIsNull) : Optional.empty(), fieldBlocks); } else { - Block nullRowBlock = fromFieldBlocks(1, Optional.of(new boolean[] {true}), fieldBlocks); + Block nullRowBlock = type.createNullBlock(); result = RunLengthEncodedBlock.create(nullRowBlock, positionCount); } diff --git a/core/trino-main/src/main/java/io/trino/sql/planner/QueryPlanner.java b/core/trino-main/src/main/java/io/trino/sql/planner/QueryPlanner.java index d00f50e6dbb8..ca124cf8e2ee 100644 --- a/core/trino-main/src/main/java/io/trino/sql/planner/QueryPlanner.java +++ b/core/trino-main/src/main/java/io/trino/sql/planner/QueryPlanner.java @@ -556,6 +556,8 @@ public PlanNode plan(Delete node) List columnSymbols = columnSymbolsBuilder.build(); Symbol operationSymbol = symbolAllocator.newSymbol("operation", TINYINT); assignmentsBuilder.put(operationSymbol, new GenericLiteral("TINYINT", String.valueOf(DELETE_OPERATION_NUMBER))); + Symbol caseNumberSymbol = symbolAllocator.newSymbol("case_number", INTEGER); + assignmentsBuilder.put(caseNumberSymbol, new GenericLiteral("INTEGER", String.valueOf(0))); Symbol projectedRowIdSymbol = symbolAllocator.newSymbol(rowIdSymbol.getName(), rowIdType); assignmentsBuilder.put(projectedRowIdSymbol, rowIdSymbol.toSymbolReference()); assignmentsBuilder.put(symbolAllocator.newSymbol("insert_from_update", TINYINT), new GenericLiteral("TINYINT", "0")); @@ -929,11 +931,13 @@ private MergeWriterNode createMergePipeline(Table table, RelationPlan relationPl } Symbol operationSymbol = symbolAllocator.newSymbol("operation", TINYINT); + Symbol caseNumberSymbol = symbolAllocator.newSymbol("case_number", INTEGER); Symbol insertFromUpdateSymbol = symbolAllocator.newSymbol("insert_from_update", TINYINT); List projectedSymbols = ImmutableList.builder() .addAll(columnSymbols) .add(operationSymbol) + .add(caseNumberSymbol) .add(rowIdSymbol) .add(insertFromUpdateSymbol) .build(); diff --git a/core/trino-main/src/main/java/io/trino/testing/PlanTester.java b/core/trino-main/src/main/java/io/trino/testing/PlanTester.java new file mode 100644 index 000000000000..b09a3cdad213 --- /dev/null +++ b/core/trino-main/src/main/java/io/trino/testing/PlanTester.java @@ -0,0 +1,916 @@ +/* + * 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.testing; + +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableSet; +import com.google.common.io.Closer; +import io.airlift.node.NodeInfo; +import io.airlift.units.Duration; +import io.opentelemetry.api.trace.Span; +import io.opentelemetry.api.trace.Tracer; +import io.trino.FeaturesConfig; +import io.trino.Session; +import io.trino.SystemSessionProperties; +import io.trino.client.NodeVersion; +import io.trino.connector.CatalogFactory; +import io.trino.connector.CatalogServiceProviderModule; +import io.trino.connector.ConnectorName; +import io.trino.connector.ConnectorServicesProvider; +import io.trino.connector.CoordinatorDynamicCatalogManager; +import io.trino.connector.DefaultCatalogFactory; +import io.trino.connector.InMemoryCatalogStore; +import io.trino.connector.LazyCatalogFactory; +import io.trino.connector.system.AnalyzePropertiesSystemTable; +import io.trino.connector.system.CatalogSystemTable; +import io.trino.connector.system.ColumnPropertiesSystemTable; +import io.trino.connector.system.GlobalSystemConnector; +import io.trino.connector.system.MaterializedViewPropertiesSystemTable; +import io.trino.connector.system.MaterializedViewSystemTable; +import io.trino.connector.system.NodeSystemTable; +import io.trino.connector.system.SchemaPropertiesSystemTable; +import io.trino.connector.system.TableCommentSystemTable; +import io.trino.connector.system.TablePropertiesSystemTable; +import io.trino.connector.system.TransactionsSystemTable; +import io.trino.cost.ComposableStatsCalculator; +import io.trino.cost.CostCalculator; +import io.trino.cost.CostCalculatorUsingExchanges; +import io.trino.cost.CostCalculatorWithEstimatedExchanges; +import io.trino.cost.CostComparator; +import io.trino.cost.FilterStatsCalculator; +import io.trino.cost.ScalarStatsCalculator; +import io.trino.cost.StatsCalculator; +import io.trino.cost.StatsCalculatorModule.StatsRulesProvider; +import io.trino.cost.StatsNormalizer; +import io.trino.cost.TaskCountEstimator; +import io.trino.eventlistener.EventListenerConfig; +import io.trino.eventlistener.EventListenerManager; +import io.trino.exchange.ExchangeManagerRegistry; +import io.trino.execution.DynamicFilterConfig; +import io.trino.execution.NodeTaskMap; +import io.trino.execution.QueryManagerConfig; +import io.trino.execution.QueryPreparer; +import io.trino.execution.QueryPreparer.PreparedQuery; +import io.trino.execution.ScheduledSplit; +import io.trino.execution.SplitAssignment; +import io.trino.execution.TableExecuteContextManager; +import io.trino.execution.TaskManagerConfig; +import io.trino.execution.querystats.PlanOptimizersStatsCollector; +import io.trino.execution.resourcegroups.NoOpResourceGroupManager; +import io.trino.execution.scheduler.NodeScheduler; +import io.trino.execution.scheduler.NodeSchedulerConfig; +import io.trino.execution.scheduler.UniformNodeSelectorFactory; +import io.trino.execution.warnings.WarningCollector; +import io.trino.index.IndexManager; +import io.trino.memory.MemoryManagerConfig; +import io.trino.memory.NodeMemoryConfig; +import io.trino.metadata.AnalyzePropertyManager; +import io.trino.metadata.BlockEncodingManager; +import io.trino.metadata.CatalogManager; +import io.trino.metadata.ColumnPropertyManager; +import io.trino.metadata.DisabledSystemSecurityMetadata; +import io.trino.metadata.FunctionBundle; +import io.trino.metadata.FunctionManager; +import io.trino.metadata.GlobalFunctionCatalog; +import io.trino.metadata.HandleResolver; +import io.trino.metadata.InMemoryNodeManager; +import io.trino.metadata.InternalBlockEncodingSerde; +import io.trino.metadata.InternalFunctionBundle; +import io.trino.metadata.InternalNodeManager; +import io.trino.metadata.MaterializedViewPropertyManager; +import io.trino.metadata.Metadata; +import io.trino.metadata.MetadataManager; +import io.trino.metadata.QualifiedObjectName; +import io.trino.metadata.SchemaPropertyManager; +import io.trino.metadata.SessionPropertyManager; +import io.trino.metadata.Split; +import io.trino.metadata.SystemFunctionBundle; +import io.trino.metadata.TableFunctionRegistry; +import io.trino.metadata.TableHandle; +import io.trino.metadata.TableProceduresPropertyManager; +import io.trino.metadata.TableProceduresRegistry; +import io.trino.metadata.TablePropertyManager; +import io.trino.metadata.TypeRegistry; +import io.trino.operator.Driver; +import io.trino.operator.DriverContext; +import io.trino.operator.DriverFactory; +import io.trino.operator.GroupByHashPageIndexerFactory; +import io.trino.operator.PagesIndex; +import io.trino.operator.PagesIndexPageSorter; +import io.trino.operator.TaskContext; +import io.trino.operator.index.IndexJoinLookupStats; +import io.trino.operator.scalar.json.JsonExistsFunction; +import io.trino.operator.scalar.json.JsonQueryFunction; +import io.trino.operator.scalar.json.JsonValueFunction; +import io.trino.operator.table.ExcludeColumns; +import io.trino.plugin.base.security.AllowAllSystemAccessControl; +import io.trino.security.GroupProviderManager; +import io.trino.server.PluginManager; +import io.trino.server.SessionPropertyDefaults; +import io.trino.server.security.CertificateAuthenticatorManager; +import io.trino.server.security.HeaderAuthenticatorConfig; +import io.trino.server.security.HeaderAuthenticatorManager; +import io.trino.server.security.PasswordAuthenticatorConfig; +import io.trino.server.security.PasswordAuthenticatorManager; +import io.trino.spi.PageIndexerFactory; +import io.trino.spi.PageSorter; +import io.trino.spi.Plugin; +import io.trino.spi.connector.CatalogHandle; +import io.trino.spi.connector.Connector; +import io.trino.spi.connector.ConnectorFactory; +import io.trino.spi.type.TypeManager; +import io.trino.spi.type.TypeOperators; +import io.trino.spiller.GenericSpillerFactory; +import io.trino.split.PageSinkManager; +import io.trino.split.PageSourceManager; +import io.trino.split.SplitManager; +import io.trino.split.SplitSource; +import io.trino.sql.PlannerContext; +import io.trino.sql.analyzer.Analysis; +import io.trino.sql.analyzer.Analyzer; +import io.trino.sql.analyzer.AnalyzerFactory; +import io.trino.sql.analyzer.QueryExplainerFactory; +import io.trino.sql.analyzer.SessionTimeProvider; +import io.trino.sql.analyzer.StatementAnalyzerFactory; +import io.trino.sql.gen.ExpressionCompiler; +import io.trino.sql.gen.JoinCompiler; +import io.trino.sql.gen.JoinFilterFunctionCompiler; +import io.trino.sql.gen.OrderingCompiler; +import io.trino.sql.gen.PageFunctionCompiler; +import io.trino.sql.parser.SqlParser; +import io.trino.sql.planner.LocalExecutionPlanner; +import io.trino.sql.planner.LocalExecutionPlanner.LocalExecutionPlan; +import io.trino.sql.planner.LogicalPlanner; +import io.trino.sql.planner.NodePartitioningManager; +import io.trino.sql.planner.OptimizerConfig; +import io.trino.sql.planner.Plan; +import io.trino.sql.planner.PlanFragmenter; +import io.trino.sql.planner.PlanNodeIdAllocator; +import io.trino.sql.planner.PlanOptimizers; +import io.trino.sql.planner.PlanOptimizersFactory; +import io.trino.sql.planner.RuleStatsRecorder; +import io.trino.sql.planner.SubPlan; +import io.trino.sql.planner.TypeAnalyzer; +import io.trino.sql.planner.optimizations.PlanOptimizer; +import io.trino.sql.planner.plan.PlanNode; +import io.trino.sql.planner.plan.PlanNodeId; +import io.trino.sql.planner.plan.TableScanNode; +import io.trino.sql.planner.planprinter.PlanPrinter; +import io.trino.sql.planner.sanity.PlanSanityChecker; +import io.trino.sql.rewrite.DescribeInputRewrite; +import io.trino.sql.rewrite.DescribeOutputRewrite; +import io.trino.sql.rewrite.ExplainRewrite; +import io.trino.sql.rewrite.ShowQueriesRewrite; +import io.trino.sql.rewrite.ShowStatsRewrite; +import io.trino.sql.rewrite.StatementRewrite; +import io.trino.testing.NullOutputOperator.NullOutputFactory; +import io.trino.transaction.InMemoryTransactionManager; +import io.trino.transaction.TransactionManager; +import io.trino.transaction.TransactionManagerConfig; +import io.trino.type.BlockTypeOperators; +import io.trino.type.InternalTypeManager; +import io.trino.type.JsonPath2016Type; +import io.trino.type.TypeDeserializer; +import io.trino.util.FinalizerService; +import org.intellij.lang.annotations.Language; + +import java.io.Closeable; +import java.io.IOException; +import java.io.UncheckedIOException; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.TimeUnit; +import java.util.function.Function; + +import static com.google.common.base.Preconditions.checkState; +import static com.google.common.collect.ImmutableList.toImmutableList; +import static com.google.common.util.concurrent.MoreExecutors.directExecutor; +import static io.airlift.concurrent.MoreFutures.getFutureValue; +import static io.airlift.concurrent.Threads.daemonThreadsNamed; +import static io.airlift.tracing.Tracing.noopTracer; +import static io.opentelemetry.api.OpenTelemetry.noop; +import static io.trino.connector.CatalogServiceProviderModule.createAccessControlProvider; +import static io.trino.connector.CatalogServiceProviderModule.createAnalyzePropertyManager; +import static io.trino.connector.CatalogServiceProviderModule.createColumnPropertyManager; +import static io.trino.connector.CatalogServiceProviderModule.createFunctionProvider; +import static io.trino.connector.CatalogServiceProviderModule.createIndexProvider; +import static io.trino.connector.CatalogServiceProviderModule.createMaterializedViewPropertyManager; +import static io.trino.connector.CatalogServiceProviderModule.createNodePartitioningProvider; +import static io.trino.connector.CatalogServiceProviderModule.createPageSinkProvider; +import static io.trino.connector.CatalogServiceProviderModule.createPageSourceProvider; +import static io.trino.connector.CatalogServiceProviderModule.createSchemaPropertyManager; +import static io.trino.connector.CatalogServiceProviderModule.createSplitManagerProvider; +import static io.trino.connector.CatalogServiceProviderModule.createTableFunctionProvider; +import static io.trino.connector.CatalogServiceProviderModule.createTableProceduresPropertyManager; +import static io.trino.connector.CatalogServiceProviderModule.createTableProceduresProvider; +import static io.trino.connector.CatalogServiceProviderModule.createTablePropertyManager; +import static io.trino.execution.ParameterExtractor.bindParameters; +import static io.trino.execution.querystats.PlanOptimizersStatsCollector.createPlanOptimizersStatsCollector; +import static io.trino.execution.warnings.WarningCollector.NOOP; +import static io.trino.spi.connector.Constraint.alwaysTrue; +import static io.trino.spi.connector.DynamicFilter.EMPTY; +import static io.trino.spiller.PartitioningSpillerFactory.unsupportedPartitioningSpillerFactory; +import static io.trino.spiller.SingleStreamSpillerFactory.unsupportedSingleStreamSpillerFactory; +import static io.trino.sql.planner.LogicalPlanner.Stage.OPTIMIZED_AND_VALIDATED; +import static io.trino.sql.planner.TypeAnalyzer.createTestingTypeAnalyzer; +import static io.trino.sql.planner.optimizations.PlanNodeSearcher.searchFrom; +import static io.trino.sql.testing.TreeAssertions.assertFormattedSql; +import static io.trino.testing.TestingTaskContext.createTaskContext; +import static io.trino.transaction.TransactionBuilder.transaction; +import static io.trino.version.EmbedVersion.testingVersionEmbedder; +import static java.util.Objects.requireNonNull; +import static java.util.concurrent.Executors.newCachedThreadPool; +import static java.util.concurrent.Executors.newScheduledThreadPool; + +public class PlanTester + implements Closeable +{ + private final Session defaultSession; + private final ExecutorService notificationExecutor; + private final ScheduledExecutorService yieldExecutor; + private final FinalizerService finalizerService; + + private final SqlParser sqlParser; + private final PlanFragmenter planFragmenter; + private final InternalNodeManager nodeManager; + private final TypeOperators typeOperators; + private final BlockTypeOperators blockTypeOperators; + private final PlannerContext plannerContext; + private final GlobalFunctionCatalog globalFunctionCatalog; + private final StatsCalculator statsCalculator; + private final ScalarStatsCalculator scalarStatsCalculator; + private final CostCalculator costCalculator; + private final CostCalculator estimatedExchangesCostCalculator; + private final TaskCountEstimator taskCountEstimator; + private final TestingAccessControlManager accessControl; + private final SplitManager splitManager; + private final PageSourceManager pageSourceManager; + private final IndexManager indexManager; + private final NodePartitioningManager nodePartitioningManager; + private final PageSinkManager pageSinkManager; + private final TransactionManager transactionManager; + private final SessionPropertyManager sessionPropertyManager; + private final SchemaPropertyManager schemaPropertyManager; + private final ColumnPropertyManager columnPropertyManager; + private final TablePropertyManager tablePropertyManager; + private final MaterializedViewPropertyManager materializedViewPropertyManager; + private final AnalyzePropertyManager analyzePropertyManager; + + private final PageFunctionCompiler pageFunctionCompiler; + private final ExpressionCompiler expressionCompiler; + private final JoinFilterFunctionCompiler joinFilterFunctionCompiler; + private final JoinCompiler joinCompiler; + private final CatalogFactory catalogFactory; + private final CoordinatorDynamicCatalogManager catalogManager; + private final PluginManager pluginManager; + private final ExchangeManagerRegistry exchangeManagerRegistry; + private final TaskManagerConfig taskManagerConfig; + private final OptimizerConfig optimizerConfig; + private final StatementAnalyzerFactory statementAnalyzerFactory; + private final TypeAnalyzer typeAnalyzer; + private boolean printPlan; + + public static PlanTester create(Session defaultSession) + { + return new PlanTester(defaultSession, 1); + } + + public static PlanTester create(Session defaultSession, int nodeCountForStats) + { + return new PlanTester(defaultSession, nodeCountForStats); + } + + private PlanTester(Session defaultSession, int nodeCountForStats) + { + requireNonNull(defaultSession, "defaultSession is null"); + + Tracer tracer = noopTracer(); + this.taskManagerConfig = new TaskManagerConfig().setTaskConcurrency(4); + this.notificationExecutor = newCachedThreadPool(daemonThreadsNamed("local-query-runner-executor-%s")); + this.yieldExecutor = newScheduledThreadPool(2, daemonThreadsNamed("local-query-runner-scheduler-%s")); + this.finalizerService = new FinalizerService(); + finalizerService.start(); + + this.typeOperators = new TypeOperators(); + this.blockTypeOperators = new BlockTypeOperators(typeOperators); + this.sqlParser = new SqlParser(); + this.nodeManager = new InMemoryNodeManager(); + PageSorter pageSorter = new PagesIndexPageSorter(new PagesIndex.TestingFactory(false)); + NodeSchedulerConfig nodeSchedulerConfig = new NodeSchedulerConfig().setIncludeCoordinator(true); + this.optimizerConfig = new OptimizerConfig(); + LazyCatalogFactory catalogFactory = new LazyCatalogFactory(); + this.catalogFactory = catalogFactory; + this.catalogManager = new CoordinatorDynamicCatalogManager(new InMemoryCatalogStore(), catalogFactory, directExecutor()); + this.transactionManager = InMemoryTransactionManager.create( + new TransactionManagerConfig().setIdleTimeout(new Duration(1, TimeUnit.DAYS)), + yieldExecutor, + catalogManager, + notificationExecutor); + HandleResolver handleResolver = new HandleResolver(); + + BlockEncodingManager blockEncodingManager = new BlockEncodingManager(); + TypeRegistry typeRegistry = new TypeRegistry(typeOperators, new FeaturesConfig()); + TypeManager typeManager = new InternalTypeManager(typeRegistry); + InternalBlockEncodingSerde blockEncodingSerde = new InternalBlockEncodingSerde(blockEncodingManager, typeManager); + + this.globalFunctionCatalog = new GlobalFunctionCatalog(); + globalFunctionCatalog.addFunctions(SystemFunctionBundle.create(new FeaturesConfig(), typeOperators, blockTypeOperators, nodeManager.getCurrentNode().getNodeVersion())); + TestingGroupProviderManager groupProvider = new TestingGroupProviderManager(); + TableFunctionRegistry tableFunctionRegistry = new TableFunctionRegistry(createTableFunctionProvider(catalogManager)); + Metadata metadata = new MetadataManager( + new DisabledSystemSecurityMetadata(), + transactionManager, + globalFunctionCatalog, + typeManager); + typeRegistry.addType(new JsonPath2016Type(new TypeDeserializer(typeManager), blockEncodingSerde)); + this.joinCompiler = new JoinCompiler(typeOperators); + PageIndexerFactory pageIndexerFactory = new GroupByHashPageIndexerFactory(joinCompiler, blockTypeOperators); + EventListenerManager eventListenerManager = new EventListenerManager(new EventListenerConfig()); + this.accessControl = new TestingAccessControlManager(transactionManager, eventListenerManager); + accessControl.loadSystemAccessControl(AllowAllSystemAccessControl.NAME, ImmutableMap.of()); + + NodeInfo nodeInfo = new NodeInfo("test"); + catalogFactory.setCatalogFactory(new DefaultCatalogFactory( + metadata, + accessControl, + handleResolver, + nodeManager, + pageSorter, + pageIndexerFactory, + nodeInfo, + testingVersionEmbedder(), + noop(), + transactionManager, + typeManager, + nodeSchedulerConfig, + optimizerConfig)); + this.splitManager = new SplitManager(createSplitManagerProvider(catalogManager), tracer, new QueryManagerConfig()); + this.pageSourceManager = new PageSourceManager(createPageSourceProvider(catalogManager)); + this.pageSinkManager = new PageSinkManager(createPageSinkProvider(catalogManager)); + this.indexManager = new IndexManager(createIndexProvider(catalogManager)); + NodeScheduler nodeScheduler = new NodeScheduler(new UniformNodeSelectorFactory(nodeManager, nodeSchedulerConfig, new NodeTaskMap(finalizerService))); + this.sessionPropertyManager = createSessionPropertyManager(catalogManager, taskManagerConfig, optimizerConfig); + this.nodePartitioningManager = new NodePartitioningManager(nodeScheduler, blockTypeOperators, createNodePartitioningProvider(catalogManager)); + TableProceduresRegistry tableProceduresRegistry = new TableProceduresRegistry(createTableProceduresProvider(catalogManager)); + FunctionManager functionManager = new FunctionManager(createFunctionProvider(catalogManager), globalFunctionCatalog); + this.schemaPropertyManager = createSchemaPropertyManager(catalogManager); + this.columnPropertyManager = createColumnPropertyManager(catalogManager); + this.tablePropertyManager = createTablePropertyManager(catalogManager); + this.materializedViewPropertyManager = createMaterializedViewPropertyManager(catalogManager); + this.analyzePropertyManager = createAnalyzePropertyManager(catalogManager); + TableProceduresPropertyManager tableProceduresPropertyManager = createTableProceduresPropertyManager(catalogManager); + + accessControl.setConnectorAccessControlProvider(createAccessControlProvider(catalogManager)); + + globalFunctionCatalog.addFunctions(new InternalFunctionBundle( + new JsonExistsFunction(functionManager, metadata, typeManager), + new JsonValueFunction(functionManager, metadata, typeManager), + new JsonQueryFunction(functionManager, metadata, typeManager))); + + this.plannerContext = new PlannerContext(metadata, typeOperators, blockEncodingSerde, typeManager, functionManager, tracer); + this.typeAnalyzer = createTestingTypeAnalyzer(plannerContext); + this.pageFunctionCompiler = new PageFunctionCompiler(functionManager, 0); + this.expressionCompiler = new ExpressionCompiler(functionManager, pageFunctionCompiler); + this.joinFilterFunctionCompiler = new JoinFilterFunctionCompiler(functionManager); + + this.statementAnalyzerFactory = new StatementAnalyzerFactory( + plannerContext, + sqlParser, + SessionTimeProvider.DEFAULT, + accessControl, + transactionManager, + groupProvider, + tableProceduresRegistry, + tableFunctionRegistry, + sessionPropertyManager, + tablePropertyManager, + analyzePropertyManager, + tableProceduresPropertyManager); + this.statsCalculator = createNewStatsCalculator(plannerContext, typeAnalyzer); + this.scalarStatsCalculator = new ScalarStatsCalculator(plannerContext, typeAnalyzer); + this.taskCountEstimator = new TaskCountEstimator(() -> nodeCountForStats); + this.costCalculator = new CostCalculatorUsingExchanges(taskCountEstimator); + this.estimatedExchangesCostCalculator = new CostCalculatorWithEstimatedExchanges(costCalculator, taskCountEstimator); + + this.planFragmenter = new PlanFragmenter(metadata, functionManager, transactionManager, catalogManager, new QueryManagerConfig()); + + GlobalSystemConnector globalSystemConnector = new GlobalSystemConnector(ImmutableSet.of( + new NodeSystemTable(nodeManager), + new CatalogSystemTable(metadata, accessControl), + new TableCommentSystemTable(metadata, accessControl), + new MaterializedViewSystemTable(metadata, accessControl), + new SchemaPropertiesSystemTable(metadata, accessControl, schemaPropertyManager), + new TablePropertiesSystemTable(metadata, accessControl, tablePropertyManager), + new MaterializedViewPropertiesSystemTable(metadata, accessControl, materializedViewPropertyManager), + new ColumnPropertiesSystemTable(metadata, accessControl, columnPropertyManager), + new AnalyzePropertiesSystemTable(metadata, accessControl, analyzePropertyManager), + new TransactionsSystemTable(typeManager, transactionManager)), + ImmutableSet.of(), + ImmutableSet.of(new ExcludeColumns.ExcludeColumnsFunction())); + + exchangeManagerRegistry = new ExchangeManagerRegistry(); + this.pluginManager = new PluginManager( + (loader, createClassLoader) -> {}, + catalogFactory, + globalFunctionCatalog, + new NoOpResourceGroupManager(), + accessControl, + Optional.of(new PasswordAuthenticatorManager(new PasswordAuthenticatorConfig())), + new CertificateAuthenticatorManager(), + Optional.of(new HeaderAuthenticatorManager(new HeaderAuthenticatorConfig())), + eventListenerManager, + new GroupProviderManager(), + new SessionPropertyDefaults(nodeInfo, accessControl), + typeRegistry, + blockEncodingManager, + handleResolver, + exchangeManagerRegistry); + + catalogManager.registerGlobalSystemConnector(globalSystemConnector); + + // rewrite session to use managed SessionPropertyMetadata + this.defaultSession = new Session( + defaultSession.getQueryId(), + Span.getInvalid(), + defaultSession.getTransactionId(), + defaultSession.isClientTransactionSupport(), + defaultSession.getIdentity(), + defaultSession.getSource(), + defaultSession.getCatalog(), + defaultSession.getSchema(), + defaultSession.getPath(), + defaultSession.getTraceToken(), + defaultSession.getTimeZoneKey(), + defaultSession.getLocale(), + defaultSession.getRemoteUserAddress(), + defaultSession.getUserAgent(), + defaultSession.getClientInfo(), + defaultSession.getClientTags(), + defaultSession.getClientCapabilities(), + defaultSession.getResourceEstimates(), + defaultSession.getStart(), + defaultSession.getSystemProperties(), + defaultSession.getCatalogProperties(), + sessionPropertyManager, + defaultSession.getPreparedStatements(), + defaultSession.getProtocolHeaders(), + defaultSession.getExchangeEncryptionKey()); + } + + private static SessionPropertyManager createSessionPropertyManager( + ConnectorServicesProvider connectorServicesProvider, + TaskManagerConfig taskManagerConfig, + OptimizerConfig optimizerConfig) + { + SystemSessionProperties sessionProperties = new SystemSessionProperties( + new QueryManagerConfig(), + taskManagerConfig, + new MemoryManagerConfig(), + new FeaturesConfig(), + optimizerConfig, + new NodeMemoryConfig(), + new DynamicFilterConfig(), + new NodeSchedulerConfig()); + return CatalogServiceProviderModule.createSessionPropertyManager(ImmutableSet.of(sessionProperties), connectorServicesProvider); + } + + private static StatsCalculator createNewStatsCalculator(PlannerContext plannerContext, TypeAnalyzer typeAnalyzer) + { + StatsNormalizer normalizer = new StatsNormalizer(); + ScalarStatsCalculator scalarStatsCalculator = new ScalarStatsCalculator(plannerContext, typeAnalyzer); + FilterStatsCalculator filterStatsCalculator = new FilterStatsCalculator(plannerContext, scalarStatsCalculator, normalizer); + return new ComposableStatsCalculator(new StatsRulesProvider(plannerContext, scalarStatsCalculator, filterStatsCalculator, normalizer).get()); + } + + @Override + public void close() + { + notificationExecutor.shutdownNow(); + yieldExecutor.shutdownNow(); + catalogManager.stop(); + finalizerService.destroy(); + } + + public TransactionManager getTransactionManager() + { + return transactionManager; + } + + public PlannerContext getPlannerContext() + { + return plannerContext; + } + + public TablePropertyManager getTablePropertyManager() + { + return tablePropertyManager; + } + + public AnalyzePropertyManager getAnalyzePropertyManager() + { + return analyzePropertyManager; + } + + public NodePartitioningManager getNodePartitioningManager() + { + return nodePartitioningManager; + } + + public PageSourceManager getPageSourceManager() + { + return pageSourceManager; + } + + public SplitManager getSplitManager() + { + return splitManager; + } + + public StatsCalculator getStatsCalculator() + { + return statsCalculator; + } + + public CostCalculator getCostCalculator() + { + return costCalculator; + } + + public CostCalculator getEstimatedExchangesCostCalculator() + { + return estimatedExchangesCostCalculator; + } + + public TaskCountEstimator getTaskCountEstimator() + { + return taskCountEstimator; + } + + public TestingAccessControlManager getAccessControl() + { + return accessControl; + } + + public Session getDefaultSession() + { + return defaultSession; + } + + public void createCatalog(String catalogName, ConnectorFactory connectorFactory, Map properties) + { + catalogFactory.addConnectorFactory(connectorFactory, ignored -> connectorFactory.getClass().getClassLoader()); + catalogManager.createCatalog(catalogName, new ConnectorName(connectorFactory.getName()), properties, false); + } + + public void installPlugin(Plugin plugin) + { + pluginManager.installPlugin(plugin, ignored -> plugin.getClass().getClassLoader()); + } + + public void addFunctions(FunctionBundle functionBundle) + { + globalFunctionCatalog.addFunctions(functionBundle); + } + + public void createCatalog(String catalogName, String connectorName, Map properties) + { + catalogManager.createCatalog(catalogName, new ConnectorName(connectorName), properties, false); + } + + public CatalogManager getCatalogManager() + { + return catalogManager; + } + + public Connector getConnector(String catalogName) + { + return catalogManager + .getConnectorServices(getCatalogHandle(catalogName)) + .getConnector(); + } + + public PlanTester printPlan() + { + printPlan = true; + return this; + } + + public T inTransaction(Function transactionSessionConsumer) + { + return inTransaction(getDefaultSession(), transactionSessionConsumer); + } + + public T inTransaction(Session session, Function transactionSessionConsumer) + { + return transaction(getTransactionManager(), getAccessControl()) + .singleStatement() + .execute(session, transactionSessionConsumer); + } + + public CatalogHandle getCatalogHandle(String catalogName) + { + return inTransaction(transactionSession -> getPlannerContext().getMetadata().getCatalogHandle(transactionSession, catalogName)).orElseThrow(); + } + + public TableHandle getTableHandle(String catalogName, String schemaName, String tableName) + { + return inTransaction(transactionSession -> + getPlannerContext().getMetadata().getTableHandle( + transactionSession, + new QualifiedObjectName(catalogName, schemaName, tableName)) + .orElseThrow()); + } + + public void executeStatement(@Language("SQL") String sql) + { + accessControl.checkCanExecuteQuery(defaultSession.getIdentity()); + + inTransaction(defaultSession, transactionSession -> { + try (Closer closer = Closer.create()) { + List drivers = createDrivers(transactionSession, sql); + drivers.forEach(closer::register); + + boolean done = false; + while (!done) { + boolean processed = false; + for (Driver driver : drivers) { + if (!driver.isFinished()) { + driver.processForNumberOfIterations(1); + processed = true; + } + } + done = !processed; + } + return null; + } + catch (IOException e) { + throw new UncheckedIOException(e); + } + }); + } + + public SubPlan createSubPlans(Session session, Plan plan, boolean forceSingleNode) + { + return planFragmenter.createSubPlans(session, plan, forceSingleNode, NOOP); + } + + private List createDrivers(Session session, @Language("SQL") String sql) + { + Plan plan = createPlan(session, sql); + if (printPlan) { + System.out.println(PlanPrinter.textLogicalPlan( + plan.getRoot(), + plan.getTypes(), + plannerContext.getMetadata(), + plannerContext.getFunctionManager(), + plan.getStatsAndCosts(), + session, + 0, + false)); + } + + SubPlan subplan = createSubPlans(session, plan, true); + if (!subplan.getChildren().isEmpty()) { + throw new AssertionError("Expected sub-plan to have no children"); + } + + TaskContext taskContext = createTaskContext(notificationExecutor, yieldExecutor, session); + TableExecuteContextManager tableExecuteContextManager = new TableExecuteContextManager(); + tableExecuteContextManager.registerTableExecuteContextForQuery(taskContext.getQueryContext().getQueryId()); + LocalExecutionPlanner executionPlanner = new LocalExecutionPlanner( + plannerContext, + typeAnalyzer, + Optional.empty(), + pageSourceManager, + indexManager, + nodePartitioningManager, + pageSinkManager, + (ignore1, ignore2, ignore3, ignore4, ignore5) -> { + throw new UnsupportedOperationException(); + }, + expressionCompiler, + pageFunctionCompiler, + joinFilterFunctionCompiler, + new IndexJoinLookupStats(), + this.taskManagerConfig, + new GenericSpillerFactory(unsupportedSingleStreamSpillerFactory()), + unsupportedSingleStreamSpillerFactory(), + unsupportedPartitioningSpillerFactory(), + new PagesIndex.TestingFactory(false), + joinCompiler, + new OrderingCompiler(plannerContext.getTypeOperators()), + new DynamicFilterConfig(), + blockTypeOperators, + tableExecuteContextManager, + exchangeManagerRegistry, + nodeManager.getCurrentNode().getNodeVersion()); + + // plan query + LocalExecutionPlan localExecutionPlan = executionPlanner.plan( + taskContext, + subplan.getFragment().getRoot(), + subplan.getFragment().getOutputPartitioningScheme().getOutputLayout(), + plan.getTypes(), + subplan.getFragment().getPartitionedSources(), + new NullOutputFactory()); + + // generate splitAssignments + List splitAssignments = new ArrayList<>(); + long sequenceId = 0; + for (TableScanNode tableScan : findTableScanNodes(subplan.getFragment().getRoot())) { + TableHandle table = tableScan.getTable(); + + SplitSource splitSource = splitManager.getSplits( + session, + Span.getInvalid(), + table, + EMPTY, + alwaysTrue()); + + ImmutableSet.Builder scheduledSplits = ImmutableSet.builder(); + while (!splitSource.isFinished()) { + for (Split split : getNextBatch(splitSource)) { + scheduledSplits.add(new ScheduledSplit(sequenceId++, tableScan.getId(), split)); + } + } + + splitAssignments.add(new SplitAssignment(tableScan.getId(), scheduledSplits.build(), true)); + } + + // create drivers + List drivers = new ArrayList<>(); + Map driverFactoriesBySource = new HashMap<>(); + for (DriverFactory driverFactory : localExecutionPlan.getDriverFactories()) { + for (int i = 0; i < driverFactory.getDriverInstances().orElse(1); i++) { + if (driverFactory.getSourceId().isPresent()) { + checkState(driverFactoriesBySource.put(driverFactory.getSourceId().get(), driverFactory) == null); + } + else { + DriverContext driverContext = taskContext.addPipelineContext(driverFactory.getPipelineId(), driverFactory.isInputDriver(), driverFactory.isOutputDriver(), false).addDriverContext(); + Driver driver = driverFactory.createDriver(driverContext); + drivers.add(driver); + } + } + } + + // add split assignments to the drivers + ImmutableSet partitionedSources = ImmutableSet.copyOf(subplan.getFragment().getPartitionedSources()); + for (SplitAssignment splitAssignment : splitAssignments) { + DriverFactory driverFactory = driverFactoriesBySource.get(splitAssignment.getPlanNodeId()); + checkState(driverFactory != null); + boolean partitioned = partitionedSources.contains(driverFactory.getSourceId().orElseThrow()); + for (ScheduledSplit split : splitAssignment.getSplits()) { + DriverContext driverContext = taskContext.addPipelineContext(driverFactory.getPipelineId(), driverFactory.isInputDriver(), driverFactory.isOutputDriver(), partitioned).addDriverContext(); + Driver driver = driverFactory.createDriver(driverContext); + driver.updateSplitAssignment(new SplitAssignment(split.getPlanNodeId(), ImmutableSet.of(split), true)); + drivers.add(driver); + } + } + + for (DriverFactory driverFactory : localExecutionPlan.getDriverFactories()) { + driverFactory.noMoreDrivers(); + } + + return ImmutableList.copyOf(drivers); + } + + public Plan createPlan(Session session, @Language("SQL") String sql) + { + return createPlan(session, sql, getPlanOptimizers(true), OPTIMIZED_AND_VALIDATED, NOOP, createPlanOptimizersStatsCollector()); + } + + public List getPlanOptimizers(boolean forceSingleNode) + { + return getPlanOptimizersFactory(forceSingleNode).get(); + } + + public PlanOptimizersFactory getPlanOptimizersFactory(boolean forceSingleNode) + { + return new PlanOptimizers( + plannerContext, + typeAnalyzer, + taskManagerConfig, + splitManager, + pageSourceManager, + statsCalculator, + scalarStatsCalculator, + costCalculator, + estimatedExchangesCostCalculator, + new CostComparator(optimizerConfig), + taskCountEstimator, + nodePartitioningManager, + new RuleStatsRecorder()); + } + + public Plan createPlan(Session session, @Language("SQL") String sql, List optimizers, LogicalPlanner.Stage stage, WarningCollector warningCollector, PlanOptimizersStatsCollector planOptimizersStatsCollector) + { + // session must be in a transaction registered with the transaction manager in this query runner + transactionManager.getTransactionInfo(session.getRequiredTransactionId()); + + PreparedQuery preparedQuery = new QueryPreparer(sqlParser).prepareQuery(session, sql); + + assertFormattedSql(sqlParser, preparedQuery.getStatement()); + + PlanNodeIdAllocator idAllocator = new PlanNodeIdAllocator(); + + AnalyzerFactory analyzerFactory = createAnalyzerFactory(createQueryExplainerFactory(optimizers)); + Analyzer analyzer = analyzerFactory.createAnalyzer( + session, + preparedQuery.getParameters(), + bindParameters(preparedQuery.getStatement(), preparedQuery.getParameters()), + warningCollector, + planOptimizersStatsCollector); + + LogicalPlanner logicalPlanner = new LogicalPlanner( + session, + optimizers, + new PlanSanityChecker(true), + idAllocator, + getPlannerContext(), + typeAnalyzer, + statsCalculator, + costCalculator, + warningCollector, + planOptimizersStatsCollector); + + Analysis analysis = analyzer.analyze(preparedQuery.getStatement()); + // make PlanTester always compute plan statistics for test purposes + return logicalPlanner.plan(analysis, stage); + } + + private QueryExplainerFactory createQueryExplainerFactory(List optimizers) + { + return new QueryExplainerFactory( + createPlanOptimizersFactory(optimizers), + planFragmenter, + plannerContext, + statementAnalyzerFactory, + statsCalculator, + costCalculator, + new NodeVersion("test")); + } + + private PlanOptimizersFactory createPlanOptimizersFactory(List optimizers) + { + return new PlanOptimizersFactory() + { + @Override + public List get() + { + return optimizers; + } + }; + } + + private AnalyzerFactory createAnalyzerFactory(QueryExplainerFactory queryExplainerFactory) + { + return new AnalyzerFactory( + statementAnalyzerFactory, + new StatementRewrite(ImmutableSet.of( + new DescribeInputRewrite(sqlParser), + new DescribeOutputRewrite(sqlParser), + new ShowQueriesRewrite( + plannerContext.getMetadata(), + sqlParser, + accessControl, + sessionPropertyManager, + schemaPropertyManager, + columnPropertyManager, + tablePropertyManager, + materializedViewPropertyManager), + new ShowStatsRewrite(plannerContext.getMetadata(), queryExplainerFactory, statsCalculator), + new ExplainRewrite(queryExplainerFactory, new QueryPreparer(sqlParser)))), + plannerContext.getTracer()); + } + + private static List getNextBatch(SplitSource splitSource) + { + return getFutureValue(splitSource.getNextBatch(1000)).getSplits(); + } + + private static List findTableScanNodes(PlanNode node) + { + return searchFrom(node) + .where(TableScanNode.class::isInstance) + .findAll().stream() + .map(TableScanNode.class::cast) + .collect(toImmutableList()); + } +} diff --git a/core/trino-main/src/test/java/io/trino/block/AbstractTestBlock.java b/core/trino-main/src/test/java/io/trino/block/AbstractTestBlock.java index a56b02549d1f..9ec068f16763 100644 --- a/core/trino-main/src/test/java/io/trino/block/AbstractTestBlock.java +++ b/core/trino-main/src/test/java/io/trino/block/AbstractTestBlock.java @@ -160,6 +160,10 @@ else if (field.getName().equals("fieldBlockBuildersList")) { // RowBlockBuilder fieldBlockBuildersList is a simple wrapper around the // array already accounted for in the instance } + else if (field.getName().equals("fieldBlocksList")) { + // RowBlockBuilder fieldBlockBuildersList is a simple wrapper around the + // array already accounted for in the instance + } else { throw new IllegalArgumentException(format("Unknown type encountered: %s", type)); } diff --git a/core/trino-main/src/test/java/io/trino/block/TestBlockBuilder.java b/core/trino-main/src/test/java/io/trino/block/TestBlockBuilder.java index fba62ca72040..da32e186f4e6 100644 --- a/core/trino-main/src/test/java/io/trino/block/TestBlockBuilder.java +++ b/core/trino-main/src/test/java/io/trino/block/TestBlockBuilder.java @@ -152,7 +152,7 @@ private static void assertInvalidPosition(Block block, int[] positions, int offs { assertThatThrownBy(() -> block.getPositions(positions, offset, length).getLong(0, 0)) .isInstanceOfAny(IllegalArgumentException.class, IndexOutOfBoundsException.class) - .hasMessage("Invalid position %d in block with %d positions", positions[0], block.getPositionCount()); + .hasMessage("Invalid position %d and length %d in block with %d positions", positions[0], length, block.getPositionCount()); } private static void assertInvalidOffset(Block block, int[] positions, int offset, int length) diff --git a/core/trino-main/src/test/java/io/trino/block/TestRowBlock.java b/core/trino-main/src/test/java/io/trino/block/TestRowBlock.java index a8504538660d..584ad6103909 100644 --- a/core/trino-main/src/test/java/io/trino/block/TestRowBlock.java +++ b/core/trino-main/src/test/java/io/trino/block/TestRowBlock.java @@ -25,17 +25,18 @@ import org.testng.annotations.Test; import java.util.ArrayList; -import java.util.Arrays; import java.util.Collections; import java.util.List; import java.util.Optional; import static io.airlift.slice.Slices.utf8Slice; import static io.trino.spi.block.RowBlock.fromFieldBlocks; +import static io.trino.spi.block.RowBlock.fromNotNullSuppressedFieldBlocks; import static io.trino.spi.type.BigintType.BIGINT; import static io.trino.spi.type.VarcharType.VARCHAR; import static java.lang.String.format; import static java.util.Objects.requireNonNull; +import static org.assertj.core.api.Assertions.assertThat; import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertFalse; import static org.testng.Assert.assertTrue; @@ -71,24 +72,21 @@ public void testEstimatedDataSizeForStats() @Test public void testFromFieldBlocksNoNullsDetection() { - Block emptyBlock = new ByteArrayBlock(0, Optional.empty(), new byte[0]); - Block fieldBlock = new ByteArrayBlock(5, Optional.empty(), createExpectedValue(5).getBytes()); - - boolean[] rowIsNull = new boolean[fieldBlock.getPositionCount()]; - Arrays.fill(rowIsNull, false); - - // Blocks may discard the null mask during creation if no values are null - assertFalse(fromFieldBlocks(5, Optional.of(rowIsNull), new Block[]{fieldBlock}).mayHaveNull()); - // Last position is null must retain the nulls mask + // Blocks does not discard the null mask during creation if no values are null + boolean[] rowIsNull = new boolean[5]; + assertThat(fromNotNullSuppressedFieldBlocks(5, Optional.of(rowIsNull), new Block[] { + new ByteArrayBlock(5, Optional.empty(), createExpectedValue(5).getBytes())}).mayHaveNull()).isTrue(); rowIsNull[rowIsNull.length - 1] = true; - assertTrue(fromFieldBlocks(5, Optional.of(rowIsNull), new Block[]{fieldBlock}).mayHaveNull()); + assertThat(fromNotNullSuppressedFieldBlocks(5, Optional.of(rowIsNull), new Block[] { + new ByteArrayBlock(5, Optional.of(rowIsNull), createExpectedValue(5).getBytes())}).mayHaveNull()).isTrue(); + // Empty blocks have no nulls and can also discard their null mask - assertFalse(fromFieldBlocks(0, Optional.of(new boolean[0]), new Block[]{emptyBlock}).mayHaveNull()); + assertThat(fromNotNullSuppressedFieldBlocks(0, Optional.of(new boolean[0]), new Block[] {new ByteArrayBlock(0, Optional.empty(), new byte[0])}).mayHaveNull()).isFalse(); // Normal blocks should have null masks preserved List fieldTypes = ImmutableList.of(VARCHAR, BIGINT); Block hasNullsBlock = createBlockBuilderWithValues(fieldTypes, alternatingNullValues(generateTestRows(fieldTypes, 100))).build(); - assertTrue(hasNullsBlock.mayHaveNull()); + assertThat(hasNullsBlock.mayHaveNull()).isTrue(); } private int getExpectedEstimatedDataSize(List row) @@ -106,19 +104,13 @@ private int getExpectedEstimatedDataSize(List row) public void testCompactBlock() { Block emptyBlock = new ByteArrayBlock(0, Optional.empty(), new byte[0]); - Block compactFieldBlock1 = new ByteArrayBlock(5, Optional.empty(), createExpectedValue(5).getBytes()); - Block compactFieldBlock2 = new ByteArrayBlock(5, Optional.empty(), createExpectedValue(5).getBytes()); - Block incompactFieldBlock1 = new ByteArrayBlock(5, Optional.empty(), createExpectedValue(6).getBytes()); - Block incompactFieldBlock2 = new ByteArrayBlock(5, Optional.empty(), createExpectedValue(6).getBytes()); boolean[] rowIsNull = {false, true, false, false, false, false}; - assertCompact(fromFieldBlocks(0, Optional.empty(), new Block[] {emptyBlock, emptyBlock})); - assertCompact(fromFieldBlocks(rowIsNull.length, Optional.of(rowIsNull), new Block[] {compactFieldBlock1, compactFieldBlock2})); - // TODO: add test case for a sliced RowBlock - - // underlying field blocks are not compact - testIncompactBlock(fromFieldBlocks(rowIsNull.length, Optional.of(rowIsNull), new Block[] {incompactFieldBlock1, incompactFieldBlock2})); - testIncompactBlock(fromFieldBlocks(rowIsNull.length, Optional.of(rowIsNull), new Block[] {incompactFieldBlock1, incompactFieldBlock2})); + // NOTE: nested row blocks are required to have the exact same size so they are always compact + assertCompact(fromFieldBlocks(0, new Block[] {emptyBlock, emptyBlock})); + assertCompact(fromNotNullSuppressedFieldBlocks(rowIsNull.length, Optional.of(rowIsNull), new Block[] { + new ByteArrayBlock(6, Optional.of(rowIsNull), createExpectedValue(6).getBytes()), + new ByteArrayBlock(6, Optional.of(rowIsNull), createExpectedValue(6).getBytes())})); } private void testWith(List fieldTypes, List[] expectedValues) diff --git a/core/trino-main/src/test/java/io/trino/operator/TestPageUtils.java b/core/trino-main/src/test/java/io/trino/operator/TestPageUtils.java index 698b7aa0c18c..d380eaf7316f 100644 --- a/core/trino-main/src/test/java/io/trino/operator/TestPageUtils.java +++ b/core/trino-main/src/test/java/io/trino/operator/TestPageUtils.java @@ -14,11 +14,12 @@ package io.trino.operator; import io.trino.spi.Page; +import io.trino.spi.block.ArrayBlock; import io.trino.spi.block.Block; -import io.trino.spi.block.DictionaryBlock; import io.trino.spi.block.LazyBlock; import org.testng.annotations.Test; +import java.util.Optional; import java.util.concurrent.atomic.AtomicLong; import static io.trino.block.BlockAssertions.createIntsBlock; @@ -50,17 +51,20 @@ public void testRecordMaterializedBytes() public void testNestedBlocks() { Block elements = lazyWrapper(createIntsBlock(1, 2, 3)); - Block dictBlock = DictionaryBlock.create(2, elements, new int[] {0, 0}); - Page page = new Page(2, dictBlock); + Block arrayBlock = ArrayBlock.fromElementBlock(2, Optional.empty(), new int[] {0, 1, 3}, elements); + long initialArraySize = arrayBlock.getSizeInBytes(); + Page page = new Page(2, arrayBlock); AtomicLong sizeInBytes = new AtomicLong(); recordMaterializedBytes(page, sizeInBytes::getAndAdd); - assertEquals(sizeInBytes.get(), dictBlock.getSizeInBytes()); + assertEquals(arrayBlock.getSizeInBytes(), initialArraySize); + assertEquals(sizeInBytes.get(), arrayBlock.getSizeInBytes()); // dictionary block caches size in bytes - dictBlock.getLoadedBlock(); - assertEquals(sizeInBytes.get(), dictBlock.getSizeInBytes() + elements.getSizeInBytes()); + arrayBlock.getLoadedBlock(); + assertEquals(sizeInBytes.get(), arrayBlock.getSizeInBytes()); + assertEquals(sizeInBytes.get(), initialArraySize + elements.getSizeInBytes()); } private static LazyBlock lazyWrapper(Block block) diff --git a/core/trino-main/src/test/java/io/trino/operator/join/TestLookupJoinPageBuilder.java b/core/trino-main/src/test/java/io/trino/operator/join/TestLookupJoinPageBuilder.java index c2ed75c593b6..d8e4345ddc2a 100644 --- a/core/trino-main/src/test/java/io/trino/operator/join/TestLookupJoinPageBuilder.java +++ b/core/trino-main/src/test/java/io/trino/operator/join/TestLookupJoinPageBuilder.java @@ -20,6 +20,7 @@ import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; import io.trino.spi.block.DictionaryBlock; +import io.trino.spi.block.LongArrayBlock; import io.trino.spi.type.Type; import org.testng.annotations.Test; @@ -100,7 +101,7 @@ public void testDifferentPositions() JoinProbe probe = joinProbeFactory.createJoinProbe(page); Page output = lookupJoinPageBuilder.build(probe); assertEquals(output.getChannelCount(), 2); - assertTrue(output.getBlock(0) instanceof DictionaryBlock); + assertTrue(output.getBlock(0) instanceof LongArrayBlock); assertEquals(output.getPositionCount(), 0); lookupJoinPageBuilder.reset(); diff --git a/core/trino-main/src/test/java/io/trino/operator/join/unspilled/TestLookupJoinPageBuilder.java b/core/trino-main/src/test/java/io/trino/operator/join/unspilled/TestLookupJoinPageBuilder.java index 3d9d254ceaa5..fd68a289dcd0 100644 --- a/core/trino-main/src/test/java/io/trino/operator/join/unspilled/TestLookupJoinPageBuilder.java +++ b/core/trino-main/src/test/java/io/trino/operator/join/unspilled/TestLookupJoinPageBuilder.java @@ -21,6 +21,7 @@ import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; import io.trino.spi.block.DictionaryBlock; +import io.trino.spi.block.LongArrayBlock; import io.trino.spi.type.Type; import org.testng.annotations.Test; @@ -101,7 +102,7 @@ public void testDifferentPositions() JoinProbe probe = joinProbeFactory.createJoinProbe(page, lookupSource); Page output = lookupJoinPageBuilder.build(probe); assertEquals(output.getChannelCount(), 2); - assertTrue(output.getBlock(0) instanceof DictionaryBlock); + assertTrue(output.getBlock(0) instanceof LongArrayBlock); assertEquals(output.getPositionCount(), 0); lookupJoinPageBuilder.reset(); diff --git a/core/trino-main/src/test/java/io/trino/server/security/oauth2/BaseOAuth2WebUiAuthenticationFilterTest.java b/core/trino-main/src/test/java/io/trino/server/security/oauth2/BaseOAuth2WebUiAuthenticationFilterTest.java index ecc60a8dcfc9..b687c299c548 100644 --- a/core/trino-main/src/test/java/io/trino/server/security/oauth2/BaseOAuth2WebUiAuthenticationFilterTest.java +++ b/core/trino-main/src/test/java/io/trino/server/security/oauth2/BaseOAuth2WebUiAuthenticationFilterTest.java @@ -145,8 +145,10 @@ public void tearDown() server, hydraIdP, () -> { - httpClient.dispatcher().executorService().shutdown(); - httpClient.connectionPool().evictAll(); + if (httpClient != null) { + httpClient.dispatcher().executorService().shutdown(); + httpClient.connectionPool().evictAll(); + } }); server = null; hydraIdP = null; diff --git a/core/trino-main/src/test/java/io/trino/server/security/oauth2/TestOAuth2WebUiAuthenticationFilterWithOpaque.java b/core/trino-main/src/test/java/io/trino/server/security/oauth2/TestOAuth2WebUiAuthenticationFilterWithOpaque.java deleted file mode 100644 index 1576c0a2a22c..000000000000 --- a/core/trino-main/src/test/java/io/trino/server/security/oauth2/TestOAuth2WebUiAuthenticationFilterWithOpaque.java +++ /dev/null @@ -1,84 +0,0 @@ -/* - * 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.server.security.oauth2; - -import com.google.common.collect.ImmutableMap; -import com.google.common.io.Resources; -import io.airlift.json.JsonCodec; -import io.jsonwebtoken.impl.DefaultClaims; -import okhttp3.Request; -import okhttp3.Response; - -import java.io.IOException; -import java.util.Map; - -import static jakarta.ws.rs.core.HttpHeaders.AUTHORIZATION; -import static org.assertj.core.api.Assertions.assertThat; -import static org.assertj.core.api.Assertions.fail; -import static org.assertj.core.api.InstanceOfAssertFactories.list; - -public class TestOAuth2WebUiAuthenticationFilterWithOpaque - extends BaseOAuth2WebUiAuthenticationFilterTest -{ - @Override - protected Map getOAuth2Config(String idpUrl) - { - return ImmutableMap.builder() - .put("web-ui.enabled", "true") - .put("web-ui.authentication.type", "oauth2") - .put("http-server.https.enabled", "true") - .put("http-server.https.keystore.path", Resources.getResource("cert/localhost.pem").getPath()) - .put("http-server.https.keystore.key", "") - .put("http-server.authentication.oauth2.issuer", "https://localhost:4444/") - .put("http-server.authentication.oauth2.auth-url", idpUrl + "/oauth2/auth") - .put("http-server.authentication.oauth2.token-url", idpUrl + "/oauth2/token") - .put("http-server.authentication.oauth2.jwks-url", idpUrl + "/.well-known/jwks.json") - .put("http-server.authentication.oauth2.userinfo-url", idpUrl + "/userinfo") - .put("http-server.authentication.oauth2.client-id", TRINO_CLIENT_ID) - .put("http-server.authentication.oauth2.client-secret", TRINO_CLIENT_SECRET) - // This is necessary as Hydra does not return `sub` from `/userinfo` for client credential grants. - .put("http-server.authentication.oauth2.principal-field", "iss") - .put("http-server.authentication.oauth2.additional-audiences", TRUSTED_CLIENT_ID) - .put("http-server.authentication.oauth2.max-clock-skew", "0s") - .put("http-server.authentication.oauth2.user-mapping.pattern", "(.*)(@.*)?") - .put("http-server.authentication.oauth2.oidc.discovery", "false") - .put("oauth2-jwk.http-client.trust-store-path", Resources.getResource("cert/localhost.pem").getPath()) - .buildOrThrow(); - } - - @Override - protected TestingHydraIdentityProvider getHydraIdp() - throws Exception - { - TestingHydraIdentityProvider hydraIdP = new TestingHydraIdentityProvider(TTL_ACCESS_TOKEN_IN_SECONDS, false, false); - hydraIdP.start(); - - return hydraIdP; - } - - @Override - protected void validateAccessToken(String cookieValue) - { - Request request = new Request.Builder().url("https://localhost:" + hydraIdP.getAuthPort() + "/userinfo").addHeader(AUTHORIZATION, "Bearer " + cookieValue).build(); - try (Response response = httpClient.newCall(request).execute()) { - assertThat(response.body()).isNotNull(); - DefaultClaims claims = new DefaultClaims(JsonCodec.mapJsonCodec(String.class, Object.class).fromJson(response.body().bytes())); - assertThat(claims.getSubject()).isEqualTo("foo@bar.com"); - assertThat(claims.get("aud")).asInstanceOf(list(String.class)).contains(TRINO_CLIENT_ID); - } - catch (IOException e) { - fail("Exception while calling /userinfo", e); - } - } -} diff --git a/core/trino-main/src/test/java/io/trino/server/security/oauth2/TestingHydraIdentityProvider.java b/core/trino-main/src/test/java/io/trino/server/security/oauth2/TestingHydraIdentityProvider.java index 0c3bad2912c5..311a1cbebe5c 100644 --- a/core/trino-main/src/test/java/io/trino/server/security/oauth2/TestingHydraIdentityProvider.java +++ b/core/trino-main/src/test/java/io/trino/server/security/oauth2/TestingHydraIdentityProvider.java @@ -19,6 +19,7 @@ import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.common.io.Resources; +import com.google.common.net.InetAddresses; import com.google.inject.Key; import com.nimbusds.oauth2.sdk.GrantType; import io.airlift.http.server.HttpServerConfig; @@ -26,6 +27,7 @@ import io.airlift.http.server.testing.TestingHttpServer; import io.airlift.log.Level; import io.airlift.log.Logging; +import io.airlift.node.NodeConfig; import io.airlift.node.NodeInfo; import io.trino.server.testing.TestingTrinoServer; import io.trino.server.ui.OAuth2WebUiAuthenticationFilter; @@ -54,6 +56,7 @@ import org.testcontainers.utility.MountableFile; import java.io.IOException; +import java.net.InetAddress; import java.net.URI; import java.time.Duration; import java.util.List; @@ -69,7 +72,7 @@ public class TestingHydraIdentityProvider implements AutoCloseable { - private static final String HYDRA_IMAGE = "oryd/hydra:v1.10.6"; + private static final String HYDRA_IMAGE = "oryd/hydra:v1.11.10"; private static final String ISSUER = "https://localhost:4444/"; private static final String DSN = "postgres://hydra:mysecretpassword@database:5432/hydra?sslmode=disable"; @@ -133,6 +136,7 @@ public void start() .withEnv("SERVE_TLS_KEY_PATH", "/tmp/certs/localhost.pem") .withEnv("SERVE_TLS_CERT_PATH", "/tmp/certs/localhost.pem") .withEnv("TTL_ACCESS_TOKEN", ttlAccessToken.getSeconds() + "s") + .withEnv("TTL_ID_TOKEN", ttlAccessToken.getSeconds() + "s") .withEnv("STRATEGIES_ACCESS_TOKEN", useJwt ? "jwt" : null) .withEnv("LOG_LEAK_SENSITIVE_VALUES", "true") .withCommand("serve", "all") @@ -219,7 +223,9 @@ public void close() private TestingHttpServer createTestingLoginAndConsentServer() throws IOException { - NodeInfo nodeInfo = new NodeInfo("test"); + NodeInfo nodeInfo = new NodeInfo(new NodeConfig() + .setEnvironment("test") + .setNodeInternalAddress(InetAddresses.toAddrString(InetAddress.getLocalHost()))); HttpServerConfig config = new HttpServerConfig().setHttpPort(0); HttpServerInfo httpServerInfo = new HttpServerInfo(config, nodeInfo); return new TestingHttpServer(httpServerInfo, nodeInfo, config, new AcceptAllLoginsAndConsentsServlet(), ImmutableMap.of()); diff --git a/core/trino-main/src/test/java/io/trino/sql/gen/TestExpressionCompiler.java b/core/trino-main/src/test/java/io/trino/sql/gen/TestExpressionCompiler.java index 61bbe8c2dfdc..d645daa3ac15 100644 --- a/core/trino-main/src/test/java/io/trino/sql/gen/TestExpressionCompiler.java +++ b/core/trino-main/src/test/java/io/trino/sql/gen/TestExpressionCompiler.java @@ -2131,7 +2131,7 @@ public void testIn() ('mBzobeIVrA','QLPvGtYprB','vBOLRnlaMN'), ('yznCLJqmVK','WgoKkNdSLR','mwAoQsMjrt') )""") - .binding("a", toLiteral(null, DOUBLE))) + .binding("a", toLiteral(null, VarcharType.createVarcharType(10)))) .isNull(BOOLEAN); } diff --git a/core/trino-main/src/test/java/io/trino/sql/planner/TestDeleteAndInsertMergeProcessor.java b/core/trino-main/src/test/java/io/trino/sql/planner/TestDeleteAndInsertMergeProcessor.java index e1dde20ba1be..e3cab4c09489 100644 --- a/core/trino-main/src/test/java/io/trino/sql/planner/TestDeleteAndInsertMergeProcessor.java +++ b/core/trino-main/src/test/java/io/trino/sql/planner/TestDeleteAndInsertMergeProcessor.java @@ -57,20 +57,21 @@ public void testSimpleDeletedRowMerge() // THEN DELETE // expected: ('Dave', 11, 'Darbyshire') DeleteAndInsertMergeProcessor processor = makeMergeProcessor(); - Page inputPage = makePageFromBlocks( - 2, - Optional.empty(), - new Block[] { + Block[] rowIdBlocks = new Block[] { makeLongArrayBlock(1, 1), // TransactionId makeLongArrayBlock(1, 0), // rowId - makeIntArrayBlock(536870912, 536870912)}, // bucket - new Block[] { + makeIntArrayBlock(536870912, 536870912)}; // bucket + Block[] mergeCaseBlocks = new Block[] { makeVarcharArrayBlock("", "Dave"), // customer makeIntArrayBlock(0, 11), // purchases makeVarcharArrayBlock("", "Devon"), // address - makeByteArrayBlock(1, 1), // "present" boolean + makeByteArrayBlock(0, 1), // "present" boolean makeByteArrayBlock(DEFAULT_CASE_OPERATION_NUMBER, DELETE_OPERATION_NUMBER), - makeIntArrayBlock(-1, 0)}); + makeIntArrayBlock(0, 0)}; + + Page inputPage = new Page( + RowBlock.fromNotNullSuppressedFieldBlocks(2, Optional.empty(), rowIdBlocks), + RowBlock.fromNotNullSuppressedFieldBlocks(2, Optional.of(new boolean[] {false, false}), mergeCaseBlocks)); Page outputPage = processor.transformPage(inputPage); assertThat(outputPage.getPositionCount()).isEqualTo(1); @@ -79,7 +80,7 @@ public void testSimpleDeletedRowMerge() assertThat((int) TINYINT.getByte(outputPage.getBlock(3), 0)).isEqualTo(DELETE_OPERATION_NUMBER); // Show that the row to be deleted is rowId 0, e.g. ('Dave', 11, 'Devon') - Block rowIdRow = outputPage.getBlock(4).getObject(0, Block.class); + Block rowIdRow = outputPage.getBlock(5).getObject(0, Block.class); assertThat(INTEGER.getInt(rowIdRow, 1)).isEqualTo(0); } @@ -104,9 +105,9 @@ public void testUpdateAndDeletedMerge() 5, Optional.of(rowIdNulls), new Block[] { - makeLongArrayBlockWithNulls(rowIdNulls, 5, 2, 1, 2, 2), // TransactionId - makeLongArrayBlockWithNulls(rowIdNulls, 5, 0, 3, 1, 2), // rowId - makeIntArrayBlockWithNulls(rowIdNulls, 5, 536870912, 536870912, 536870912, 536870912)}, // bucket + makeLongArrayBlockWithNulls(rowIdNulls, 5, 2, 0, 1, 2, 2), // TransactionId + makeLongArrayBlockWithNulls(rowIdNulls, 5, 0, 0, 3, 1, 2), // rowId + makeIntArrayBlockWithNulls(rowIdNulls, 5, 536870912, 0, 536870912, 536870912, 536870912)}, // bucket new Block[] { // customer makeVarcharArrayBlock("Aaron", "Carol", "Dave", "Dave", "Ed"), @@ -122,7 +123,7 @@ public void testUpdateAndDeletedMerge() Page outputPage = processor.transformPage(inputPage); assertThat(outputPage.getPositionCount()).isEqualTo(8); - RowBlock rowIdBlock = (RowBlock) outputPage.getBlock(4); + RowBlock rowIdBlock = (RowBlock) outputPage.getBlock(5); assertThat(rowIdBlock.getPositionCount()).isEqualTo(8); // Show that the first row has address "Arches" assertThat(getString(outputPage.getBlock(2), 1)).isEqualTo("Arches/Arches"); @@ -144,9 +145,9 @@ public void testAnotherMergeCase() 5, Optional.of(rowIdNulls), new Block[] { - makeLongArrayBlockWithNulls(rowIdNulls, 5, 2, 1, 2, 2), // TransactionId - makeLongArrayBlockWithNulls(rowIdNulls, 5, 0, 3, 1, 2), // rowId - makeIntArrayBlockWithNulls(rowIdNulls, 5, 536870912, 536870912, 536870912, 536870912)}, // bucket + makeLongArrayBlockWithNulls(rowIdNulls, 5, 2, 0, 1, 2, 2), // TransactionId + makeLongArrayBlockWithNulls(rowIdNulls, 5, 0, 0, 3, 1, 2), // rowId + makeIntArrayBlockWithNulls(rowIdNulls, 5, 536870912, 0, 536870912, 536870912, 536870912)}, // bucket new Block[] { // customer makeVarcharArrayBlock("Aaron", "Carol", "Dave", "Dave", "Ed"), @@ -162,7 +163,7 @@ public void testAnotherMergeCase() Page outputPage = processor.transformPage(inputPage); assertThat(outputPage.getPositionCount()).isEqualTo(8); - RowBlock rowIdBlock = (RowBlock) outputPage.getBlock(4); + RowBlock rowIdBlock = (RowBlock) outputPage.getBlock(5); assertThat(rowIdBlock.getPositionCount()).isEqualTo(8); // Show that the first row has address "Arches/Arches" assertThat(getString(outputPage.getBlock(2), 1)).isEqualTo("Arches/Arches"); @@ -171,7 +172,7 @@ public void testAnotherMergeCase() private Page makePageFromBlocks(int positionCount, Optional rowIdNulls, Block[] rowIdBlocks, Block[] mergeCaseBlocks) { Block[] pageBlocks = new Block[] { - RowBlock.fromFieldBlocks(positionCount, rowIdNulls, rowIdBlocks), + RowBlock.fromNotNullSuppressedFieldBlocks(positionCount, rowIdNulls, rowIdBlocks), RowBlock.fromFieldBlocks(positionCount, Optional.empty(), mergeCaseBlocks) }; return new Page(pageBlocks); @@ -198,7 +199,7 @@ private LongArrayBlock makeLongArrayBlock(long... elements) private LongArrayBlock makeLongArrayBlockWithNulls(boolean[] nulls, int positionCount, long... elements) { - assertThat(countNonNull(nulls) + elements.length).isEqualTo(positionCount); + assertThat(elements.length).isEqualTo(positionCount); return new LongArrayBlock(elements.length, Optional.of(nulls), elements); } @@ -209,21 +210,10 @@ private IntArrayBlock makeIntArrayBlock(int... elements) private IntArrayBlock makeIntArrayBlockWithNulls(boolean[] nulls, int positionCount, int... elements) { - assertThat(countNonNull(nulls) + elements.length).isEqualTo(positionCount); + assertThat(elements.length).isEqualTo(positionCount); return new IntArrayBlock(elements.length, Optional.of(nulls), elements); } - private int countNonNull(boolean[] nulls) - { - int count = 0; - for (int position = 0; position < nulls.length; position++) { - if (nulls[position]) { - count++; - } - } - return count; - } - private ByteArrayBlock makeByteArrayBlock(int... elements) { byte[] bytes = new byte[elements.length]; @@ -237,7 +227,12 @@ private Block makeVarcharArrayBlock(String... elements) { BlockBuilder builder = VARCHAR.createBlockBuilder(new PageBuilderStatus().createBlockBuilderStatus(), elements.length); for (String element : elements) { - VARCHAR.writeSlice(builder, Slices.utf8Slice(element)); + if (element == null) { + builder.appendNull(); + } + else { + VARCHAR.writeSlice(builder, Slices.utf8Slice(element)); + } } return builder.build(); } diff --git a/core/trino-spi/src/main/java/io/trino/spi/PageBlockUtil.java b/core/trino-spi/src/main/java/io/trino/spi/PageBlockUtil.java new file mode 100644 index 000000000000..c259274a03be --- /dev/null +++ b/core/trino-spi/src/main/java/io/trino/spi/PageBlockUtil.java @@ -0,0 +1,83 @@ +/* + * 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.spi; + +import io.trino.spi.block.ArrayBlock; +import io.trino.spi.block.Block; +import io.trino.spi.block.ByteArrayBlock; +import io.trino.spi.block.DictionaryBlock; +import io.trino.spi.block.Fixed12Block; +import io.trino.spi.block.Int128ArrayBlock; +import io.trino.spi.block.IntArrayBlock; +import io.trino.spi.block.LongArrayBlock; +import io.trino.spi.block.MapBlock; +import io.trino.spi.block.RowBlock; +import io.trino.spi.block.RunLengthEncodedBlock; +import io.trino.spi.block.ShortArrayBlock; +import io.trino.spi.block.VariableWidthBlock; + +import static java.util.Objects.requireNonNull; + +public class PageBlockUtil +{ + private PageBlockUtil() {} + + public static boolean isValueBlock(Block block) + { + return block instanceof ArrayBlock + || block instanceof ByteArrayBlock + || block instanceof ShortArrayBlock + || block instanceof IntArrayBlock + || block instanceof LongArrayBlock + || block instanceof Int128ArrayBlock + || block instanceof Fixed12Block + || block instanceof MapBlock + || block instanceof RowBlock + || block instanceof VariableWidthBlock; + } + + public static Block getUnderlyingValueBlock(Block block) + { + if (block instanceof RunLengthEncodedBlock runLengthEncodedBlock) { + return runLengthEncodedBlock.getValue(); + } + if (block instanceof DictionaryBlock dictionaryBlock) { + return dictionaryBlock.getDictionary(); + } + return block; + } + + public static int getUnderlyingValuePosition(Block block, int position) + { + if (block instanceof RunLengthEncodedBlock) { + return 0; + } + if (block instanceof DictionaryBlock dictionaryBlock) { + return dictionaryBlock.getId(position); + } + return position; + } + + public static Page getPositions(Page page, int[] retainedPositions, int offset, int length) + { + requireNonNull(retainedPositions, "retainedPositions is null"); + + Block[] blocks = new Block[page.getChannelCount()]; + for (int i = 0; i < blocks.length; i++) { + blocks[i] = page.getBlock(i).getPositions(retainedPositions, offset, length); + } + return Page.wrapBlocksWithoutCopy(length, blocks); + } +} diff --git a/core/trino-spi/src/main/java/io/trino/spi/block/AbstractRowBlock.java b/core/trino-spi/src/main/java/io/trino/spi/block/AbstractRowBlock.java index ad3c6f773470..f7b2f885d3c2 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/block/AbstractRowBlock.java +++ b/core/trino-spi/src/main/java/io/trino/spi/block/AbstractRowBlock.java @@ -17,6 +17,7 @@ import java.util.List; import java.util.OptionalInt; +import java.util.stream.IntStream; import static io.trino.spi.block.BlockUtil.arraySame; import static io.trino.spi.block.BlockUtil.checkArrayRange; @@ -24,7 +25,6 @@ import static io.trino.spi.block.BlockUtil.checkValidPositions; import static io.trino.spi.block.BlockUtil.checkValidRegion; import static io.trino.spi.block.BlockUtil.compactArray; -import static io.trino.spi.block.BlockUtil.compactOffsets; import static io.trino.spi.block.RowBlock.createRowBlockInternal; public abstract class AbstractRowBlock @@ -41,9 +41,15 @@ public final List getChildren() protected abstract Block[] getRawFieldBlocks(); @Nullable - protected abstract int[] getFieldBlockOffsets(); + protected int[] getFieldBlockOffsets() + { + return IntStream.range(0, getPositionCount() + 1).toArray(); + } - protected abstract int getOffsetBase(); + protected int getOffsetBase() + { + return 0; + } /** * @return the underlying rowIsNull array, or null when all rows are guaranteed to be non-null @@ -54,8 +60,7 @@ public final List getChildren() // the offset in each field block, it can also be viewed as the "entry-based" offset in the RowBlock public final int getFieldBlockOffset(int position) { - int[] offsets = getFieldBlockOffsets(); - return offsets != null ? offsets[position + getOffsetBase()] : position + getOffsetBase(); + return position; } protected AbstractRowBlock(int numFields) @@ -75,59 +80,45 @@ public String getEncodingName() @Override public Block copyPositions(int[] positions, int offset, int length) { + boolean[] rowIsNull = getRowIsNull(); + Block[] fieldBlocks = getRawFieldBlocks(); + checkArrayRange(positions, offset, length); - int[] newOffsets = null; - - int[] fieldBlockPositions = new int[length]; - int fieldBlockPositionCount; - boolean[] newRowIsNull; - if (getRowIsNull() == null) { - // No nulls are present - newRowIsNull = null; - for (int i = 0; i < fieldBlockPositions.length; i++) { - int position = positions[offset + i]; - checkReadablePosition(this, position); - fieldBlockPositions[i] = getFieldBlockOffset(position); - } - fieldBlockPositionCount = fieldBlockPositions.length; + Block[] newBlocks = new Block[numFields]; + for (int i = 0; i < fieldBlocks.length; i++) { + newBlocks[i] = fieldBlocks[i].copyPositions(positions, offset, length); } - else { + + boolean[] newRowIsNull = null; + if (rowIsNull != null) { newRowIsNull = new boolean[length]; - newOffsets = new int[length + 1]; - fieldBlockPositionCount = 0; for (int i = 0; i < length; i++) { - newOffsets[i] = fieldBlockPositionCount; - int position = positions[offset + i]; - boolean positionIsNull = isNull(position); - newRowIsNull[i] = positionIsNull; - fieldBlockPositions[fieldBlockPositionCount] = getFieldBlockOffset(position); - fieldBlockPositionCount += positionIsNull ? 0 : 1; - } - // Record last offset position - newOffsets[length] = fieldBlockPositionCount; - if (fieldBlockPositionCount == length) { - // No nulls encountered, discard the null mask and offsets - newRowIsNull = null; - newOffsets = null; + newRowIsNull[i] = rowIsNull[positions[offset + i]]; } } - Block[] newBlocks = new Block[numFields]; - Block[] rawBlocks = getRawFieldBlocks(); - for (int i = 0; i < newBlocks.length; i++) { - newBlocks[i] = rawBlocks[i].copyPositions(fieldBlockPositions, 0, fieldBlockPositionCount); - } - return createRowBlockInternal(0, length, newRowIsNull, newOffsets, newBlocks); + return createRowBlockInternal(length, newRowIsNull, newBlocks); } @Override - public Block getRegion(int position, int length) + public Block getRegion(int positionOffset, int length) { int positionCount = getPositionCount(); - checkValidRegion(positionCount, position, length); + boolean[] rowIsNull = getRowIsNull(); + Block[] fieldBlocks = getRawFieldBlocks(); - return createRowBlockInternal(position + getOffsetBase(), length, getRowIsNull(), getFieldBlockOffsets(), getRawFieldBlocks()); + checkValidRegion(positionCount, positionOffset, length); + + // This copies the null array, but this dramatically simplifies this class. + // Without a copy here, we would need a null array offset, and that would mean that the + // null array would be offset while the field blocks are not offset, which is confusing. + boolean[] newRowIsNull = rowIsNull == null ? null : compactArray(rowIsNull, positionOffset, length); + Block[] newBlocks = new Block[fieldBlocks.length]; + for (int i = 0; i < newBlocks.length; i++) { + newBlocks[i] = fieldBlocks[i].getRegion(positionOffset, length); + } + return createRowBlockInternal(length, newRowIsNull, newBlocks); } @Override @@ -201,7 +192,7 @@ public final long getPositionsSizeInBytes(boolean[] positions, int selectedRowPo boolean[] rowIsNull = getRowIsNull(); if (rowIsNull != null) { // Some positions in usedPositions may be null which must be removed from the selectedFieldPositionCount - int offsetBase = getOffsetBase(); + int offsetBase = 0; for (int i = 0; i < positions.length; i++) { if (positions[i] && rowIsNull[i + offsetBase]) { selectedFieldPositionCount--; // selected row is null, don't include it in the selected field positions @@ -220,67 +211,34 @@ public final long getPositionsSizeInBytes(boolean[] positions, int selectedRowPo private long getSpecificPositionsSizeInBytes(boolean[] positions, int selectedRowPositions) { - int positionCount = getPositionCount(); - int offsetBase = getOffsetBase(); - boolean[] rowIsNull = getRowIsNull(); - // No fixed width size per row, specific positions used must be tracked - int totalFieldPositions = getRawFieldBlocks()[0].getPositionCount(); - boolean[] fieldPositions; - int selectedFieldPositionCount; - if (rowIsNull == null) { - // No nulls, so the same number of positions are used - selectedFieldPositionCount = selectedRowPositions; - if (offsetBase == 0 && positionCount == totalFieldPositions) { - // No need to adapt the positions array at all, reuse it directly - fieldPositions = positions; - } - else { - // no nulls present, so we can just shift the positions array into alignment with the elements block with other positions unused - fieldPositions = new boolean[totalFieldPositions]; - System.arraycopy(positions, 0, fieldPositions, offsetBase, positions.length); - } - } - else { - fieldPositions = new boolean[totalFieldPositions]; - selectedFieldPositionCount = 0; - for (int i = 0; i < positions.length; i++) { - if (positions[i] && !rowIsNull[offsetBase + i]) { - selectedFieldPositionCount++; - fieldPositions[getFieldBlockOffset(i)] = true; - } - } - } - Block[] rawFieldBlocks = getRawFieldBlocks(); long sizeInBytes = ((Integer.BYTES + Byte.BYTES) * (long) selectedRowPositions); // offsets + rowIsNull for (int j = 0; j < numFields; j++) { - sizeInBytes += rawFieldBlocks[j].getPositionsSizeInBytes(fieldPositions, selectedFieldPositionCount); + sizeInBytes += rawFieldBlocks[j].getPositionsSizeInBytes(positions, selectedRowPositions); } return sizeInBytes; } @Override - public Block copyRegion(int position, int length) + public Block copyRegion(int positionOffset, int length) { int positionCount = getPositionCount(); - checkValidRegion(positionCount, position, length); + boolean[] rowIsNull = getRowIsNull(); + Block[] fieldBlocks = getRawFieldBlocks(); - int startFieldBlockOffset = getFieldBlockOffset(position); - int endFieldBlockOffset = getFieldBlockOffset(position + length); - int fieldBlockLength = endFieldBlockOffset - startFieldBlockOffset; - Block[] newBlocks = new Block[numFields]; - for (int i = 0; i < numFields; i++) { - newBlocks[i] = getRawFieldBlocks()[i].copyRegion(startFieldBlockOffset, fieldBlockLength); + checkValidRegion(positionCount, positionOffset, length); + + Block[] newBlocks = new Block[fieldBlocks.length]; + for (int i = 0; i < fieldBlocks.length; i++) { + newBlocks[i] = fieldBlocks[i].copyRegion(positionOffset, length); } - int[] fieldBlockOffsets = getFieldBlockOffsets(); - int[] newOffsets = fieldBlockOffsets == null ? null : compactOffsets(fieldBlockOffsets, position + getOffsetBase(), length); - boolean[] rowIsNull = getRowIsNull(); - boolean[] newRowIsNull = rowIsNull == null ? null : compactArray(rowIsNull, position + getOffsetBase(), length); - if (arraySame(newBlocks, getRawFieldBlocks()) && newOffsets == fieldBlockOffsets && newRowIsNull == rowIsNull) { + boolean[] newRowIsNull = rowIsNull == null ? null : compactArray(rowIsNull, positionOffset, length); + if (newRowIsNull == rowIsNull && arraySame(newBlocks, fieldBlocks)) { return this; } - return createRowBlockInternal(0, length, newRowIsNull, newOffsets, newBlocks); + + return createRowBlockInternal(length, newRowIsNull, newBlocks); } @Override @@ -299,17 +257,13 @@ public Block getSingleValueBlock(int position) { checkReadablePosition(this, position); - int startFieldBlockOffset = getFieldBlockOffset(position); - int endFieldBlockOffset = getFieldBlockOffset(position + 1); - int fieldBlockLength = endFieldBlockOffset - startFieldBlockOffset; Block[] newBlocks = new Block[numFields]; for (int i = 0; i < numFields; i++) { - newBlocks[i] = getRawFieldBlocks()[i].copyRegion(startFieldBlockOffset, fieldBlockLength); + newBlocks[i] = getRawFieldBlocks()[i].getSingleValueBlock(position); } boolean[] newRowIsNull = isNull(position) ? new boolean[] {true} : null; - int[] newOffsets = isNull(position) ? new int[] {0, fieldBlockLength} : null; - return createRowBlockInternal(0, 1, newRowIsNull, newOffsets, newBlocks); + return createRowBlockInternal(1, newRowIsNull, newBlocks); } @Override @@ -334,6 +288,6 @@ public boolean isNull(int position) { checkReadablePosition(this, position); boolean[] rowIsNull = getRowIsNull(); - return rowIsNull != null && rowIsNull[position + getOffsetBase()]; + return rowIsNull != null && rowIsNull[position]; } } diff --git a/core/trino-spi/src/main/java/io/trino/spi/block/Block.java b/core/trino-spi/src/main/java/io/trino/spi/block/Block.java index 3281bb7e6727..651f93fb3b3a 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/block/Block.java +++ b/core/trino-spi/src/main/java/io/trino/spi/block/Block.java @@ -22,6 +22,7 @@ import java.util.function.ObjLongConsumer; import static io.trino.spi.block.BlockUtil.checkArrayRange; +import static io.trino.spi.block.DictionaryId.randomDictionaryId; public interface Block { @@ -243,7 +244,7 @@ default Block getPositions(int[] positions, int offset, int length) { checkArrayRange(positions, offset, length); - return new DictionaryBlock(offset, length, this, positions); + return DictionaryBlock.createInternal(offset, length, this, positions, randomDictionaryId()); } /** diff --git a/core/trino-spi/src/main/java/io/trino/spi/block/BlockUtil.java b/core/trino-spi/src/main/java/io/trino/spi/block/BlockUtil.java index 69dd522d696c..4f66ffa79f79 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/block/BlockUtil.java +++ b/core/trino-spi/src/main/java/io/trino/spi/block/BlockUtil.java @@ -79,22 +79,36 @@ static void checkReadablePosition(Block block, int position) static int calculateNewArraySize(int currentSize) { - // grow array by 50% - long newSize = (long) currentSize + (currentSize >> 1); + return calculateNewArraySize(currentSize, DEFAULT_CAPACITY); + } - // verify new size is within reasonable bounds - if (newSize < DEFAULT_CAPACITY) { - newSize = DEFAULT_CAPACITY; + static int calculateNewArraySize(int currentSize, int minimumSize) + { + if (currentSize < 0 || currentSize > MAX_ARRAY_SIZE || minimumSize < 0 || minimumSize > MAX_ARRAY_SIZE) { + throw new IllegalArgumentException("Invalid currentSize or minimumSize"); } - else if (newSize > MAX_ARRAY_SIZE) { - newSize = MAX_ARRAY_SIZE; - if (newSize == currentSize) { - throw new IllegalArgumentException(format("Cannot grow array beyond '%s'", MAX_ARRAY_SIZE)); - } + if (currentSize == MAX_ARRAY_SIZE) { + throw new IllegalArgumentException("Cannot grow array beyond size " + MAX_ARRAY_SIZE); } + + minimumSize = Math.max(minimumSize, DEFAULT_CAPACITY); + + // grow the array by 50% if possible + long newSize = (long) currentSize + (currentSize >> 1); + + // ensure new size is within bounds + newSize = clamp(newSize, minimumSize, MAX_ARRAY_SIZE); return (int) newSize; } + public static int clamp(long value, int min, int max) + { + if (min > max) { + throw new IllegalArgumentException(min + " > " + max); + } + return (int) Math.min(max, Math.max(value, min)); + } + static int calculateBlockResetSize(int currentSize) { long newSize = (long) ceil(currentSize * BLOCK_RESET_SKEW); diff --git a/core/trino-spi/src/main/java/io/trino/spi/block/ColumnarRow.java b/core/trino-spi/src/main/java/io/trino/spi/block/ColumnarRow.java index 6eb174484b8c..ff5f65be21f2 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/block/ColumnarRow.java +++ b/core/trino-spi/src/main/java/io/trino/spi/block/ColumnarRow.java @@ -15,6 +15,8 @@ import jakarta.annotation.Nullable; +import java.util.stream.Stream; + import static java.util.Objects.requireNonNull; public final class ColumnarRow @@ -42,13 +44,18 @@ public static ColumnarRow toColumnarRow(Block block) throw new IllegalArgumentException("Invalid row block: " + block.getClass().getName()); } - // get fields for visible region - int firstRowPosition = rowBlock.getFieldBlockOffset(0); - int totalRowCount = rowBlock.getFieldBlockOffset(block.getPositionCount()) - firstRowPosition; - Block[] fieldBlocks = new Block[rowBlock.numFields]; - for (int i = 0; i < fieldBlocks.length; i++) { - fieldBlocks[i] = rowBlock.getRawFieldBlocks()[i].getRegion(firstRowPosition, totalRowCount); + int[] nonNullPositions = new int[rowBlock.getPositionCount()]; + int idCount = 0; + for (int position = 0; position < nonNullPositions.length; position++) { + if (!rowBlock.isNull(position)) { + nonNullPositions[idCount] = position; + idCount++; + } } + int nonNullPositionCount = idCount; + Block[] fieldBlocks = Stream.of(rowBlock.getRawFieldBlocks()) + .map(field -> DictionaryBlock.create(nonNullPositionCount, field, nonNullPositions)) + .toArray(Block[]::new); return new ColumnarRow(block.getPositionCount(), block, fieldBlocks); } diff --git a/core/trino-spi/src/main/java/io/trino/spi/block/DictionaryBlock.java b/core/trino-spi/src/main/java/io/trino/spi/block/DictionaryBlock.java index f48fa8b1f16b..0df8c17eb201 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/block/DictionaryBlock.java +++ b/core/trino-spi/src/main/java/io/trino/spi/block/DictionaryBlock.java @@ -25,6 +25,9 @@ import static io.airlift.slice.SizeOf.instanceSize; import static io.airlift.slice.SizeOf.sizeOf; +import static io.trino.spi.PageBlockUtil.getUnderlyingValueBlock; +import static io.trino.spi.PageBlockUtil.getUnderlyingValuePosition; +import static io.trino.spi.PageBlockUtil.isValueBlock; import static io.trino.spi.block.BlockUtil.checkArrayRange; import static io.trino.spi.block.BlockUtil.checkValidPosition; import static io.trino.spi.block.BlockUtil.checkValidPositions; @@ -55,7 +58,7 @@ public class DictionaryBlock public static Block create(int positionCount, Block dictionary, int[] ids) { - return createInternal(positionCount, dictionary, ids, randomDictionaryId()); + return createInternal(0, positionCount, dictionary, ids, randomDictionaryId()); } /** @@ -63,16 +66,16 @@ public static Block create(int positionCount, Block dictionary, int[] ids) */ public static Block createProjectedDictionaryBlock(int positionCount, Block dictionary, int[] ids, DictionaryId dictionarySourceId) { - return createInternal(positionCount, dictionary, ids, dictionarySourceId); + return createInternal(0, positionCount, dictionary, ids, dictionarySourceId); } - private static Block createInternal(int positionCount, Block dictionary, int[] ids, DictionaryId dictionarySourceId) + static Block createInternal(int idsOffset, int positionCount, Block dictionary, int[] ids, DictionaryId dictionarySourceId) { if (positionCount == 0) { return dictionary.copyRegion(0, 0); } if (positionCount == 1) { - return dictionary.getRegion(ids[0], 1); + return dictionary.getRegion(ids[idsOffset], 1); } // if dictionary is an RLE then this can just be a new RLE @@ -80,17 +83,16 @@ private static Block createInternal(int positionCount, Block dictionary, int[] i return RunLengthEncodedBlock.create(rle.getValue(), positionCount); } + if (isValueBlock(dictionary)) { + return new DictionaryBlock(idsOffset, positionCount, dictionary, ids, false, false, dictionarySourceId); + } + // unwrap dictionary in dictionary - if (dictionary instanceof DictionaryBlock dictionaryBlock) { - int[] newIds = new int[positionCount]; - for (int position = 0; position < positionCount; position++) { - newIds[position] = dictionaryBlock.getId(ids[position]); - } - dictionary = dictionaryBlock.getDictionary(); - dictionarySourceId = randomDictionaryId(); - ids = newIds; + int[] newIds = new int[positionCount]; + for (int position = 0; position < positionCount; position++) { + newIds[position] = getUnderlyingValuePosition(dictionary, ids[idsOffset + position]); } - return new DictionaryBlock(0, positionCount, dictionary, ids, false, false, dictionarySourceId); + return new DictionaryBlock(0, positionCount, getUnderlyingValueBlock(dictionary), newIds, false, false, randomDictionaryId()); } DictionaryBlock(int idsOffset, int positionCount, Block dictionary, int[] ids) @@ -597,6 +599,34 @@ public Block getDictionary() return dictionary; } + public Block createProjection(Block newDictionary) + { + if (newDictionary.getPositionCount() != dictionary.getPositionCount()) { + throw new IllegalArgumentException("newDictionary must have the same position count"); + } + + // if the new dictionary is lazy be careful to not materialize it + if (newDictionary instanceof LazyBlock lazyBlock) { + return new LazyBlock(positionCount, () -> { + Block newDictionaryBlock = lazyBlock.getBlock(); + return createProjection(newDictionaryBlock); + }); + } + if (isValueBlock(newDictionary)) { + return new DictionaryBlock(idsOffset, positionCount, newDictionary, ids, isCompact(), false, dictionarySourceId); + } + if (newDictionary instanceof RunLengthEncodedBlock rle) { + return RunLengthEncodedBlock.create(rle.getValue(), positionCount); + } + + // unwrap dictionary in dictionary + int[] newIds = new int[positionCount]; + for (int position = 0; position < positionCount; position++) { + newIds[position] = getUnderlyingValuePosition(newDictionary, getIdUnchecked(position)); + } + return new DictionaryBlock(0, positionCount, getUnderlyingValueBlock(newDictionary), newIds, false, false, randomDictionaryId()); + } + Slice getIds() { return Slices.wrappedIntArray(ids, idsOffset, positionCount); diff --git a/core/trino-spi/src/main/java/io/trino/spi/block/RowBlock.java b/core/trino-spi/src/main/java/io/trino/spi/block/RowBlock.java index fe01877938af..76b33f118134 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/block/RowBlock.java +++ b/core/trino-spi/src/main/java/io/trino/spi/block/RowBlock.java @@ -15,13 +15,15 @@ import jakarta.annotation.Nullable; +import java.util.Arrays; +import java.util.List; import java.util.Optional; import java.util.function.ObjLongConsumer; +import java.util.stream.IntStream; import static io.airlift.slice.SizeOf.instanceSize; import static io.airlift.slice.SizeOf.sizeOf; -import static io.trino.spi.block.BlockUtil.copyIsNullAndAppendNull; -import static io.trino.spi.block.BlockUtil.copyOffsetsAndAppendNull; +import static io.trino.spi.block.BlockUtil.checkArrayRange; import static io.trino.spi.block.BlockUtil.ensureBlocksAreLoaded; import static java.lang.String.format; import static java.util.Objects.requireNonNull; @@ -31,79 +33,92 @@ public class RowBlock { private static final int INSTANCE_SIZE = instanceSize(RowBlock.class); - private final int startOffset; private final int positionCount; private final boolean[] rowIsNull; private final int[] fieldBlockOffsets; private final Block[] fieldBlocks; + private final List fieldBlocksList; private volatile long sizeInBytes = -1; private final long retainedSizeInBytes; + /** + * Create a row block directly from field blocks. The returned RowBlock will not contain any null rows, although the fields may contain null values. + */ + public static RowBlock fromFieldBlocks(int positionCount, Block[] fieldBlocks) + { + return createRowBlockInternal(positionCount, null, fieldBlocks); + } + /** * Create a row block directly from columnar nulls and field blocks. */ public static Block fromFieldBlocks(int positionCount, Optional rowIsNullOptional, Block[] fieldBlocks) { boolean[] rowIsNull = rowIsNullOptional.orElse(null); - int[] fieldBlockOffsets = null; + validateConstructorArguments(positionCount, rowIsNull, fieldBlocks); + return new RowBlock(positionCount, rowIsNull, fieldBlocks); + } + + /** + * Create a row block directly from field blocks that are not null-suppressed. The field value of a null row must be null. + */ + public static RowBlock fromNotNullSuppressedFieldBlocks(int positionCount, Optional rowIsNullOptional, Block[] fieldBlocks) + { + // verify that field values for null rows are null + boolean[] rowIsNull = rowIsNullOptional.orElse(null); if (rowIsNull != null) { - // Check for nulls when computing field block offsets - fieldBlockOffsets = new int[positionCount + 1]; - fieldBlockOffsets[0] = 0; - for (int position = 0; position < positionCount; position++) { - fieldBlockOffsets[position + 1] = fieldBlockOffsets[position] + (rowIsNull[position] ? 0 : 1); - } - // fieldBlockOffsets is positionCount + 1 in length - if (fieldBlockOffsets[positionCount] == positionCount) { - // No nulls encountered, discard the null mask - rowIsNull = null; - fieldBlockOffsets = null; + checkArrayRange(rowIsNull, 0, positionCount); + + for (int fieldIndex = 0; fieldIndex < fieldBlocks.length; fieldIndex++) { + Block field = fieldBlocks[fieldIndex]; + // LazyBlock may not have loaded the field yet + if (!(field instanceof LazyBlock lazyBlock) || lazyBlock.isLoaded()) { + for (int position = 0; position < positionCount; position++) { + if (rowIsNull[position] && !field.isNull(position)) { + throw new IllegalArgumentException(format("Field value for null row must be null: field %s, position %s", fieldIndex, position)); + } + } + } } } - - validateConstructorArguments(0, positionCount, rowIsNull, fieldBlockOffsets, fieldBlocks); - return new RowBlock(0, positionCount, rowIsNull, fieldBlockOffsets, fieldBlocks); + return createRowBlockInternal(positionCount, rowIsNull, fieldBlocks); } /** * Create a row block directly without per element validations. */ - static RowBlock createRowBlockInternal(int startOffset, int positionCount, @Nullable boolean[] rowIsNull, @Nullable int[] fieldBlockOffsets, Block[] fieldBlocks) + static RowBlock createRowBlockInternal(int positionCount, @Nullable boolean[] rowIsNull, Block[] fieldBlocks) { - validateConstructorArguments(startOffset, positionCount, rowIsNull, fieldBlockOffsets, fieldBlocks); - return new RowBlock(startOffset, positionCount, rowIsNull, fieldBlockOffsets, fieldBlocks); + validateConstructorArguments(positionCount, rowIsNull, fieldBlocks); + return new RowBlock(positionCount, rowIsNull, fieldBlocks); } - private static void validateConstructorArguments(int startOffset, int positionCount, @Nullable boolean[] rowIsNull, @Nullable int[] fieldBlockOffsets, Block[] fieldBlocks) + private static void validateConstructorArguments(int positionCount, @Nullable boolean[] rowIsNull, Block[] fieldBlocks) { - if (startOffset < 0) { - throw new IllegalArgumentException("arrayOffset is negative"); - } - if (positionCount < 0) { throw new IllegalArgumentException("positionCount is negative"); } - if (rowIsNull != null && rowIsNull.length - startOffset < positionCount) { + if (rowIsNull != null && rowIsNull.length < positionCount) { throw new IllegalArgumentException("rowIsNull length is less than positionCount"); } - if ((rowIsNull == null) != (fieldBlockOffsets == null)) { - throw new IllegalArgumentException("When rowIsNull is (non) null then fieldBlockOffsets should be (non) null as well"); - } - - if (fieldBlockOffsets != null && fieldBlockOffsets.length - startOffset < positionCount + 1) { - throw new IllegalArgumentException("fieldBlockOffsets length is less than positionCount"); - } - requireNonNull(fieldBlocks, "fieldBlocks is null"); if (fieldBlocks.length <= 0) { throw new IllegalArgumentException("Number of fields in RowBlock must be positive"); } + if (rowIsNull != null) { + for (Block field : fieldBlocks) { + if (field.getPositionCount() < positionCount) { + throw new IllegalArgumentException("Sparse RowBlock is not supported"); + } + } + } + int firstFieldBlockPositionCount = fieldBlocks[0].getPositionCount(); for (int i = 1; i < fieldBlocks.length; i++) { if (firstFieldBlockPositionCount != fieldBlocks[i].getPositionCount()) { @@ -116,15 +131,15 @@ private static void validateConstructorArguments(int startOffset, int positionCo * Use createRowBlockInternal or fromFieldBlocks instead of this method. The caller of this method is assumed to have * validated the arguments with validateConstructorArguments. */ - private RowBlock(int startOffset, int positionCount, @Nullable boolean[] rowIsNull, @Nullable int[] fieldBlockOffsets, Block[] fieldBlocks) + private RowBlock(int positionCount, @Nullable boolean[] rowIsNull, Block[] fieldBlocks) { super(fieldBlocks.length); - this.startOffset = startOffset; this.positionCount = positionCount; - this.rowIsNull = rowIsNull; - this.fieldBlockOffsets = fieldBlockOffsets; + this.rowIsNull = positionCount == 0 ? null : rowIsNull; + this.fieldBlockOffsets = IntStream.range(0, positionCount + 1).toArray(); this.fieldBlocks = fieldBlocks; + this.fieldBlocksList = List.of(fieldBlocks); this.retainedSizeInBytes = INSTANCE_SIZE + sizeOf(fieldBlockOffsets) + sizeOf(rowIsNull); } @@ -135,17 +150,9 @@ protected Block[] getRawFieldBlocks() return fieldBlocks; } - @Override - @Nullable - protected int[] getFieldBlockOffsets() - { - return fieldBlockOffsets; - } - - @Override - protected int getOffsetBase() + public Block getFieldBlock(int fieldIndex) { - return startOffset; + return fieldBlocks[fieldIndex]; } @Override @@ -177,12 +184,8 @@ public long getSizeInBytes() long sizeInBytes = getBaseSizeInBytes(); boolean hasUnloadedBlocks = false; - int startFieldBlockOffset = fieldBlockOffsets != null ? fieldBlockOffsets[startOffset] : startOffset; - int endFieldBlockOffset = fieldBlockOffsets != null ? fieldBlockOffsets[startOffset + positionCount] : startOffset + positionCount; - int fieldBlockLength = endFieldBlockOffset - startFieldBlockOffset; - for (Block fieldBlock : fieldBlocks) { - sizeInBytes += fieldBlock.getRegionSizeInBytes(startFieldBlockOffset, fieldBlockLength); + sizeInBytes += fieldBlock.getSizeInBytes(); hasUnloadedBlocks = hasUnloadedBlocks || !fieldBlock.isLoaded(); } @@ -207,6 +210,12 @@ public long getRetainedSizeInBytes() return retainedSizeInBytes; } + @Override + public Block getPositions(int[] retainedPositions, int offset, int length) + { + return copyPositions(retainedPositions, offset, length); + } + @Override public void retainedBytesForEachPart(ObjLongConsumer consumer) { @@ -248,35 +257,116 @@ public Block getLoadedBlock() return this; } return createRowBlockInternal( - startOffset, positionCount, rowIsNull, - fieldBlockOffsets, loadedFieldBlocks); } @Override public Block copyWithAppendedNull() { - boolean[] newRowIsNull = copyIsNullAndAppendNull(getRowIsNull(), getOffsetBase(), getPositionCount()); - - int[] newOffsets; - if (getFieldBlockOffsets() == null) { - int desiredLength = getOffsetBase() + positionCount + 2; - newOffsets = new int[desiredLength]; - newOffsets[getOffsetBase()] = getOffsetBase(); - for (int position = getOffsetBase(); position < getOffsetBase() + positionCount; position++) { - // Since there are no nulls in the original array, new offsets are the same as previous ones - newOffsets[position + 1] = newOffsets[position] + 1; - } - - // Null does not change offset - newOffsets[desiredLength - 1] = newOffsets[desiredLength - 2]; + boolean[] newRowIsNull; + if (rowIsNull != null) { + newRowIsNull = Arrays.copyOf(rowIsNull, positionCount + 1); } else { - newOffsets = copyOffsetsAndAppendNull(getFieldBlockOffsets(), getOffsetBase(), getPositionCount()); + newRowIsNull = new boolean[positionCount + 1]; + } + // mark the (new) last element as null + newRowIsNull[positionCount] = true; + + Block[] newBlocks = new Block[fieldBlocks.length]; + for (int i = 0; i < fieldBlocks.length; i++) { + newBlocks[i] = fieldBlocks[i].copyWithAppendedNull(); + } + return new RowBlock(positionCount + 1, newRowIsNull, newBlocks); + } + + /** + * Returns the row fields from the specified block. The block maybe a LazyBlock, RunLengthEncodedBlock, or + * DictionaryBlock, but the underlying block must be a RowBlock. The returned field blocks will be the same + * length as the specified block, which means they are not null suppressed. + */ + public static List getRowFieldsFromBlock(Block block) + { + // if the block is lazy, be careful to not materialize the nested blocks + if (block instanceof LazyBlock lazyBlock) { + block = lazyBlock.getBlock(); } - return createRowBlockInternal(getOffsetBase(), getPositionCount() + 1, newRowIsNull, newOffsets, getRawFieldBlocks()); + if (block instanceof RunLengthEncodedBlock runLengthEncodedBlock) { + RowBlock rowBlock = (RowBlock) runLengthEncodedBlock.getValue(); + return rowBlock.fieldBlocksList.stream() + .map(fieldBlock -> RunLengthEncodedBlock.create(fieldBlock, runLengthEncodedBlock.getPositionCount())) + .toList(); + } + if (block instanceof DictionaryBlock dictionaryBlock) { + RowBlock rowBlock = (RowBlock) dictionaryBlock.getDictionary(); + return rowBlock.fieldBlocksList.stream() + .map(dictionaryBlock::createProjection) + .toList(); + } + if (block instanceof RowBlock rowBlock) { + return List.of(rowBlock.fieldBlocks); + } + throw new IllegalArgumentException("Unexpected block type: " + block.getClass().getSimpleName()); + } + + /** + * Returns the row fields from the specified block with null rows suppressed. The block maybe a LazyBlock, RunLengthEncodedBlock, or + * DictionaryBlock, but the underlying block must be a RowBlock. The returned field blocks will not be the same + * length as the specified block if it contains null rows. + */ + public static List getNullSuppressedRowFieldsFromBlock(Block block) + { + // if the block is lazy, be careful to not materialize the nested blocks + if (block instanceof LazyBlock lazyBlock) { + block = lazyBlock.getBlock(); + } + + if (!block.mayHaveNull()) { + return getRowFieldsFromBlock(block); + } + + if (block instanceof RunLengthEncodedBlock runLengthEncodedBlock) { + RowBlock rowBlock = (RowBlock) runLengthEncodedBlock.getValue(); + if (!rowBlock.isNull(0)) { + throw new IllegalStateException("Expected run length encoded block value to be null"); + } + // all values are null, so return a zero-length block of the correct type + return rowBlock.fieldBlocksList.stream() + .map(fieldBlock -> fieldBlock.getRegion(0, 0)) + .toList(); + } + if (block instanceof DictionaryBlock dictionaryBlock) { + int[] newIds = new int[dictionaryBlock.getPositionCount()]; + int idCount = 0; + for (int position = 0; position < newIds.length; position++) { + if (!dictionaryBlock.isNull(position)) { + newIds[idCount] = dictionaryBlock.getId(position); + idCount++; + } + } + int nonNullPositionCount = idCount; + RowBlock rowBlock = (RowBlock) dictionaryBlock.getDictionary(); + return rowBlock.fieldBlocksList.stream() + .map(field -> DictionaryBlock.create(nonNullPositionCount, field, newIds)) + .toList(); + } + if (block instanceof RowBlock rowBlock) { + int[] nonNullPositions = new int[rowBlock.getPositionCount()]; + int idCount = 0; + for (int position = 0; position < nonNullPositions.length; position++) { + if (!rowBlock.isNull(position)) { + nonNullPositions[idCount] = position; + idCount++; + } + } + int nonNullPositionCount = idCount; + return rowBlock.fieldBlocksList.stream() + .map(field -> DictionaryBlock.create(nonNullPositionCount, field, nonNullPositions)) + .toList(); + } + throw new IllegalArgumentException("Unexpected block type: " + block.getClass().getSimpleName()); } } diff --git a/core/trino-spi/src/main/java/io/trino/spi/block/RowBlockBuilder.java b/core/trino-spi/src/main/java/io/trino/spi/block/RowBlockBuilder.java index 69b2c2ac5822..6d26daa3d50e 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/block/RowBlockBuilder.java +++ b/core/trino-spi/src/main/java/io/trino/spi/block/RowBlockBuilder.java @@ -39,7 +39,6 @@ public class RowBlockBuilder private final BlockBuilderStatus blockBuilderStatus; private int positionCount; - private int[] fieldBlockOffsets; private boolean[] rowIsNull; private final BlockBuilder[] fieldBlockBuilders; private final List fieldBlockBuildersList; @@ -53,17 +52,15 @@ public RowBlockBuilder(List fieldTypes, BlockBuilderStatus blockBuilderSta this( blockBuilderStatus, createFieldBlockBuilders(fieldTypes, blockBuilderStatus, expectedEntries), - new int[expectedEntries + 1], new boolean[expectedEntries]); } - private RowBlockBuilder(@Nullable BlockBuilderStatus blockBuilderStatus, BlockBuilder[] fieldBlockBuilders, int[] fieldBlockOffsets, boolean[] rowIsNull) + private RowBlockBuilder(@Nullable BlockBuilderStatus blockBuilderStatus, BlockBuilder[] fieldBlockBuilders, boolean[] rowIsNull) { super(fieldBlockBuilders.length); this.blockBuilderStatus = blockBuilderStatus; this.positionCount = 0; - this.fieldBlockOffsets = requireNonNull(fieldBlockOffsets, "fieldBlockOffsets is null"); this.rowIsNull = requireNonNull(rowIsNull, "rowIsNull is null"); this.fieldBlockBuilders = requireNonNull(fieldBlockBuilders, "fieldBlockBuilders is null"); this.fieldBlockBuildersList = List.of(fieldBlockBuilders); @@ -85,19 +82,6 @@ protected Block[] getRawFieldBlocks() return fieldBlockBuilders; } - @Override - @Nullable - protected int[] getFieldBlockOffsets() - { - return hasNullRow ? fieldBlockOffsets : null; - } - - @Override - protected int getOffsetBase() - { - return 0; - } - @Nullable @Override protected boolean[] getRowIsNull() @@ -130,7 +114,7 @@ public long getSizeInBytes() @Override public long getRetainedSizeInBytes() { - long size = INSTANCE_SIZE + sizeOf(fieldBlockOffsets) + sizeOf(rowIsNull); + long size = INSTANCE_SIZE + sizeOf(rowIsNull); for (int i = 0; i < numFields; i++) { size += fieldBlockBuilders[i].getRetainedSizeInBytes(); } @@ -146,7 +130,6 @@ public void retainedBytesForEachPart(ObjLongConsumer consumer) for (int i = 0; i < numFields; i++) { consumer.accept(fieldBlockBuilders[i], fieldBlockBuilders[i].getRetainedSizeInBytes()); } - consumer.accept(fieldBlockOffsets, sizeOf(fieldBlockOffsets)); consumer.accept(rowIsNull, sizeOf(rowIsNull)); consumer.accept(this, INSTANCE_SIZE); } @@ -170,32 +153,27 @@ public BlockBuilder appendNull() if (currentEntryOpened) { throw new IllegalStateException("Current entry must be closed before a null can be written"); } + + for (BlockBuilder fieldBlockBuilder : fieldBlockBuilders) { + fieldBlockBuilder.appendNull(); + } + entryAdded(true); return this; } private void entryAdded(boolean isNull) { - if (rowIsNull.length <= positionCount) { - int newSize = BlockUtil.calculateNewArraySize(rowIsNull.length); - rowIsNull = Arrays.copyOf(rowIsNull, newSize); - fieldBlockOffsets = Arrays.copyOf(fieldBlockOffsets, newSize + 1); - } + ensureCapacity(positionCount + 1); - if (isNull) { - fieldBlockOffsets[positionCount + 1] = fieldBlockOffsets[positionCount]; - } - else { - fieldBlockOffsets[positionCount + 1] = fieldBlockOffsets[positionCount] + 1; - } rowIsNull[positionCount] = isNull; hasNullRow |= isNull; hasNonNullRow |= !isNull; positionCount++; - for (int i = 0; i < numFields; i++) { - if (fieldBlockBuilders[i].getPositionCount() != fieldBlockOffsets[positionCount]) { - throw new IllegalStateException(format("field %s has unexpected position count. Expected: %s, actual: %s", i, fieldBlockOffsets[positionCount], fieldBlockBuilders[i].getPositionCount())); + for (int i = 0; i < fieldBlockBuilders.length; i++) { + if (fieldBlockBuilders[i].getPositionCount() != positionCount) { + throw new IllegalStateException(format("field %s has unexpected position count. Expected: %s, actual: %s", i, positionCount, fieldBlockBuilders[i].getPositionCount())); } } @@ -217,7 +195,17 @@ public Block build() for (int i = 0; i < numFields; i++) { fieldBlocks[i] = fieldBlockBuilders[i].build(); } - return createRowBlockInternal(0, positionCount, hasNullRow ? rowIsNull : null, hasNullRow ? fieldBlockOffsets : null, fieldBlocks); + return createRowBlockInternal(positionCount, hasNullRow ? rowIsNull : null, fieldBlocks); + } + + private void ensureCapacity(int capacity) + { + if (rowIsNull.length >= capacity) { + return; + } + + int newSize = BlockUtil.calculateNewArraySize(rowIsNull.length, capacity); + rowIsNull = Arrays.copyOf(rowIsNull, newSize); } @Override @@ -233,7 +221,7 @@ public BlockBuilder newBlockBuilderLike(int expectedEntries, BlockBuilderStatus for (int i = 0; i < numFields; i++) { newBlockBuilders[i] = fieldBlockBuilders[i].newBlockBuilderLike(blockBuilderStatus); } - return new RowBlockBuilder(blockBuilderStatus, newBlockBuilders, new int[expectedEntries + 1], new boolean[expectedEntries]); + return new RowBlockBuilder(blockBuilderStatus, newBlockBuilders, new boolean[expectedEntries]); } @Override @@ -275,10 +263,10 @@ private Block nullRle(int length) { Block[] fieldBlocks = new Block[numFields]; for (int i = 0; i < numFields; i++) { - fieldBlocks[i] = fieldBlockBuilders[i].newBlockBuilderLike(null).build(); + fieldBlocks[i] = fieldBlockBuilders[i].newBlockBuilderLike(null).appendNull().build(); } - RowBlock nullRowBlock = createRowBlockInternal(0, 1, new boolean[] {true}, new int[] {0, 0}, fieldBlocks); + RowBlock nullRowBlock = createRowBlockInternal(1, new boolean[] {true}, fieldBlocks); return RunLengthEncodedBlock.create(nullRowBlock, length); } } diff --git a/core/trino-spi/src/main/java/io/trino/spi/block/RowBlockEncoding.java b/core/trino-spi/src/main/java/io/trino/spi/block/RowBlockEncoding.java index 00fe6302fea4..7e3af8995960 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/block/RowBlockEncoding.java +++ b/core/trino-spi/src/main/java/io/trino/spi/block/RowBlockEncoding.java @@ -17,7 +17,6 @@ import io.airlift.slice.SliceInput; import io.airlift.slice.SliceOutput; -import static io.airlift.slice.Slices.wrappedIntArray; import static io.trino.spi.block.RowBlock.createRowBlockInternal; public class RowBlockEncoding @@ -35,42 +34,19 @@ public String getName() public void writeBlock(BlockEncodingSerde blockEncodingSerde, SliceOutput sliceOutput, Block block) { AbstractRowBlock rowBlock = (AbstractRowBlock) block; - int[] fieldBlockOffsets = rowBlock.getFieldBlockOffsets(); int numFields = rowBlock.numFields; int positionCount = rowBlock.getPositionCount(); - int offsetBase = rowBlock.getOffsetBase(); - - int startFieldBlockOffset = fieldBlockOffsets != null ? fieldBlockOffsets[offsetBase] : offsetBase; - int endFieldBlockOffset = fieldBlockOffsets != null ? fieldBlockOffsets[offsetBase + positionCount] : offsetBase + positionCount; - sliceOutput.appendInt(numFields); sliceOutput.appendInt(positionCount); for (int i = 0; i < numFields; i++) { - blockEncodingSerde.writeBlock(sliceOutput, rowBlock.getRawFieldBlocks()[i].getRegion(startFieldBlockOffset, endFieldBlockOffset - startFieldBlockOffset)); + blockEncodingSerde.writeBlock(sliceOutput, rowBlock.getRawFieldBlocks()[i]); } EncoderUtil.encodeNullsAsBits(sliceOutput, block); - - if ((rowBlock.getRowIsNull() == null) != (fieldBlockOffsets == null)) { - throw new IllegalArgumentException("When rowIsNull is (non) null then fieldBlockOffsets should be (non) null as well"); - } - - if (fieldBlockOffsets != null) { - if (startFieldBlockOffset == 0) { - sliceOutput.writeBytes(wrappedIntArray(fieldBlockOffsets, offsetBase, positionCount + 1)); - } - else { - int[] newFieldBlockOffsets = new int[positionCount + 1]; - for (int position = 0; position < positionCount + 1; position++) { - newFieldBlockOffsets[position] = fieldBlockOffsets[offsetBase + position] - startFieldBlockOffset; - } - sliceOutput.writeBytes(wrappedIntArray(newFieldBlockOffsets)); - } - } } @Override @@ -85,11 +61,6 @@ public Block readBlock(BlockEncodingSerde blockEncodingSerde, SliceInput sliceIn } boolean[] rowIsNull = EncoderUtil.decodeNullBits(sliceInput, positionCount).orElse(null); - int[] fieldBlockOffsets = null; - if (rowIsNull != null) { - fieldBlockOffsets = new int[positionCount + 1]; - sliceInput.readBytes(wrappedIntArray(fieldBlockOffsets)); - } - return createRowBlockInternal(0, positionCount, rowIsNull, fieldBlockOffsets, fieldBlocks); + return createRowBlockInternal(positionCount, rowIsNull, fieldBlocks); } } diff --git a/core/trino-spi/src/main/java/io/trino/spi/block/RunLengthEncodedBlock.java b/core/trino-spi/src/main/java/io/trino/spi/block/RunLengthEncodedBlock.java index 996d35398c87..1b55a0cdd653 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/block/RunLengthEncodedBlock.java +++ b/core/trino-spi/src/main/java/io/trino/spi/block/RunLengthEncodedBlock.java @@ -24,6 +24,9 @@ import java.util.function.ObjLongConsumer; import static io.airlift.slice.SizeOf.instanceSize; +import static io.trino.spi.PageBlockUtil.getUnderlyingValueBlock; +import static io.trino.spi.PageBlockUtil.getUnderlyingValuePosition; +import static io.trino.spi.PageBlockUtil.isValueBlock; import static io.trino.spi.block.BlockUtil.checkArrayRange; import static io.trino.spi.block.BlockUtil.checkReadablePosition; import static io.trino.spi.block.BlockUtil.checkValidPosition; @@ -59,6 +62,21 @@ public static Block create(Block value, int positionCount) if (positionCount == 1) { return value; } + if (isValueBlock(value)) { + return new RunLengthEncodedBlock(value, positionCount); + } + + // if the value is lazy be careful to not materialize it + if (value instanceof LazyBlock lazyBlock) { + return new LazyBlock(positionCount, () -> create(lazyBlock.getBlock(), positionCount)); + } + + // unwrap the value + Block valueBlock = getUnderlyingValueBlock(value); + int valuePosition = getUnderlyingValuePosition(valueBlock, 0); + if (valueBlock.getPositionCount() == 1 && valuePosition == 0) { + return new RunLengthEncodedBlock(valueBlock, positionCount); + } return new RunLengthEncodedBlock(value, positionCount); } diff --git a/core/trino-spi/src/main/java/io/trino/spi/connector/ConnectorBucketNodeMap.java b/core/trino-spi/src/main/java/io/trino/spi/connector/ConnectorBucketNodeMap.java index b8228f732f49..30a9cf0c0e75 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/connector/ConnectorBucketNodeMap.java +++ b/core/trino-spi/src/main/java/io/trino/spi/connector/ConnectorBucketNodeMap.java @@ -17,6 +17,7 @@ import java.util.List; import java.util.Optional; +import java.util.concurrent.ThreadLocalRandom; import static java.lang.String.format; @@ -24,18 +25,19 @@ public final class ConnectorBucketNodeMap { private final int bucketCount; private final Optional> bucketToNode; + private final long cacheKeyHint; public static ConnectorBucketNodeMap createBucketNodeMap(int bucketCount) { - return new ConnectorBucketNodeMap(bucketCount, Optional.empty()); + return new ConnectorBucketNodeMap(bucketCount, Optional.empty(), ThreadLocalRandom.current().nextLong()); } public static ConnectorBucketNodeMap createBucketNodeMap(List bucketToNode) { - return new ConnectorBucketNodeMap(bucketToNode.size(), Optional.of(bucketToNode)); + return new ConnectorBucketNodeMap(bucketToNode.size(), Optional.of(bucketToNode), ThreadLocalRandom.current().nextLong()); } - private ConnectorBucketNodeMap(int bucketCount, Optional> bucketToNode) + private ConnectorBucketNodeMap(int bucketCount, Optional> bucketToNode, long cacheKeyHint) { if (bucketCount <= 0) { throw new IllegalArgumentException("bucketCount must be positive"); @@ -45,6 +47,7 @@ private ConnectorBucketNodeMap(int bucketCount, Optional> bucketToNod } this.bucketCount = bucketCount; this.bucketToNode = bucketToNode.map(List::copyOf); + this.cacheKeyHint = cacheKeyHint; } public int getBucketCount() @@ -61,4 +64,14 @@ public List getFixedMapping() { return bucketToNode.orElseThrow(() -> new IllegalArgumentException("No fixed bucket to node mapping")); } + + public long getCacheKeyHint() + { + return cacheKeyHint; + } + + public ConnectorBucketNodeMap withCacheKeyHint(long cacheKeyHint) + { + return new ConnectorBucketNodeMap(bucketCount, bucketToNode, cacheKeyHint); + } } diff --git a/core/trino-spi/src/main/java/io/trino/spi/connector/MergePage.java b/core/trino-spi/src/main/java/io/trino/spi/connector/MergePage.java index ad40f85f0a9f..4b9745289708 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/connector/MergePage.java +++ b/core/trino-spi/src/main/java/io/trino/spi/connector/MergePage.java @@ -62,15 +62,15 @@ public static MergePage createDeleteAndInsertPages(Page inputPage, int dataColum { // see page description in ConnectorMergeSink int inputChannelCount = inputPage.getChannelCount(); - if (inputChannelCount != dataColumnCount + 2) { - throw new IllegalArgumentException(format("inputPage channelCount (%s) == dataColumns size (%s) + 2", inputChannelCount, dataColumnCount)); + if (inputChannelCount != dataColumnCount + 3) { + throw new IllegalArgumentException(format("inputPage channelCount (%s) == dataColumns size (%s) + 3", inputChannelCount, dataColumnCount)); } int positionCount = inputPage.getPositionCount(); if (positionCount <= 0) { throw new IllegalArgumentException("positionCount should be > 0, but is " + positionCount); } - Block operationBlock = inputPage.getBlock(inputChannelCount - 2); + Block operationBlock = inputPage.getBlock(dataColumnCount); int[] deletePositions = new int[positionCount]; int[] insertPositions = new int[positionCount]; @@ -99,7 +99,7 @@ public static MergePage createDeleteAndInsertPages(Page inputPage, int dataColum for (int i = 0; i < dataColumnCount; i++) { columns[i] = i; } - columns[dataColumnCount] = dataColumnCount + 1; // row ID channel + columns[dataColumnCount] = dataColumnCount + 2; // row ID channel deletePage = Optional.of(inputPage .getColumns(columns) .getPositions(deletePositions, 0, deletePositionCount)); diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/s3select/S3SelectDataType.java b/core/trino-spi/src/main/java/io/trino/spi/connector/RelationType.java similarity index 84% rename from plugin/trino-hive/src/main/java/io/trino/plugin/hive/s3select/S3SelectDataType.java rename to core/trino-spi/src/main/java/io/trino/spi/connector/RelationType.java index 70872574d5db..d038a1de8bc3 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/s3select/S3SelectDataType.java +++ b/core/trino-spi/src/main/java/io/trino/spi/connector/RelationType.java @@ -11,9 +11,11 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package io.trino.plugin.hive.s3select; +package io.trino.spi.connector; -public enum S3SelectDataType { - CSV, - JSON +public enum RelationType +{ + TABLE, + VIEW, + MATERIALIZED_VIEW, } diff --git a/core/trino-spi/src/main/java/io/trino/spi/predicate/TupleDomain.java b/core/trino-spi/src/main/java/io/trino/spi/predicate/TupleDomain.java index 6398ec443829..c7177eda8487 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/predicate/TupleDomain.java +++ b/core/trino-spi/src/main/java/io/trino/spi/predicate/TupleDomain.java @@ -235,6 +235,21 @@ public Optional> getDomains() return domains; } + public Domain getDomain(T column, Type type) + { + if (domains.isEmpty()) { + return Domain.none(type); + } + Domain domain = domains.get().get(column); + if (domain != null && !domain.getType().equals(type)) { + throw new IllegalArgumentException("Provided type %s does not match domain type %s for column %s".formatted(type, domain.getType(), column)); + } + if (domain == null) { + return Domain.all(type); + } + return domain; + } + /** * Returns the strict intersection of the TupleDomains. * The resulting TupleDomain represents the set of tuples that would be valid diff --git a/core/trino-spi/src/main/java/io/trino/spi/type/MapType.java b/core/trino-spi/src/main/java/io/trino/spi/type/MapType.java index 026e65a01f9a..39cd3147e80d 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/type/MapType.java +++ b/core/trino-spi/src/main/java/io/trino/spi/type/MapType.java @@ -87,6 +87,7 @@ public class MapType private final MethodHandle keyBlockNativeNotDistinctFrom; private final MethodHandle keyBlockNotDistinctFrom; + private final MethodHandle keyBlockIdentical; private final MethodHandle keyNativeHashCode; private final MethodHandle keyBlockHashCode; private final MethodHandle keyBlockNativeEqual; @@ -117,6 +118,7 @@ public MapType(Type keyType, Type valueType, TypeOperators typeOperators) keyBlockNativeNotDistinctFrom = filterReturnValue(typeOperators.getDistinctFromOperator(keyType, simpleConvention(FAIL_ON_NULL, BLOCK_POSITION, NEVER_NULL)), NOT) .asType(methodType(boolean.class, Block.class, int.class, keyType.getJavaType().isPrimitive() ? keyType.getJavaType() : Object.class)); keyBlockNotDistinctFrom = filterReturnValue(typeOperators.getDistinctFromOperator(keyType, simpleConvention(FAIL_ON_NULL, BLOCK_POSITION, BLOCK_POSITION)), NOT); + keyBlockIdentical = typeOperators.getDistinctFromOperator(keyType, simpleConvention(FAIL_ON_NULL, BLOCK_POSITION, BLOCK_POSITION)); keyNativeHashCode = typeOperators.getHashCodeOperator(keyType, HASH_CODE_CONVENTION) .asType(methodType(long.class, keyType.getJavaType().isPrimitive() ? keyType.getJavaType() : Object.class)); @@ -341,6 +343,14 @@ public MethodHandle getKeyBlockNotDistinctFrom() return keyBlockNotDistinctFrom; } + /** + * Internal use by this package and io.trino.spi.block only. + */ + public MethodHandle getKeyBlockIdentical() + { + return keyBlockIdentical; + } + private static long hashOperator(MethodHandle keyOperator, MethodHandle valueOperator, Block block) throws Throwable { diff --git a/core/trino-spi/src/main/java/io/trino/spi/type/Type.java b/core/trino-spi/src/main/java/io/trino/spi/type/Type.java index dab302cc06be..e90ebf2a2ef1 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/type/Type.java +++ b/core/trino-spi/src/main/java/io/trino/spi/type/Type.java @@ -98,6 +98,16 @@ default TypeOperatorDeclaration getTypeOperatorDeclaration(TypeOperators typeOpe */ BlockBuilder createBlockBuilder(BlockBuilderStatus blockBuilderStatus, int expectedEntries); + /** + * Creates a block containing as single null values. + */ + default Block createNullBlock() + { + return createBlockBuilder(null, 1, 0) + .appendNull() + .build(); + } + /** * Gets an object representation of the type value in the {@code block} * {@code position}. This is the value returned to the user via the diff --git a/core/trino-spi/src/test/java/io/trino/spi/block/TestBlockUtil.java b/core/trino-spi/src/test/java/io/trino/spi/block/TestBlockUtil.java index 0399387457e1..d3f1167240ae 100644 --- a/core/trino-spi/src/test/java/io/trino/spi/block/TestBlockUtil.java +++ b/core/trino-spi/src/test/java/io/trino/spi/block/TestBlockUtil.java @@ -16,7 +16,8 @@ import org.testng.annotations.Test; import static io.trino.spi.block.BlockUtil.MAX_ARRAY_SIZE; -import static java.lang.String.format; +import static io.trino.spi.block.BlockUtil.calculateNewArraySize; +import static org.assertj.core.api.Assertions.assertThatThrownBy; import static org.testng.Assert.assertEquals; public class TestBlockUtil @@ -24,13 +25,21 @@ public class TestBlockUtil @Test public void testCalculateNewArraySize() { - assertEquals(BlockUtil.calculateNewArraySize(200), 300); - assertEquals(BlockUtil.calculateNewArraySize(Integer.MAX_VALUE), MAX_ARRAY_SIZE); - try { - BlockUtil.calculateNewArraySize(MAX_ARRAY_SIZE); - } - catch (IllegalArgumentException e) { - assertEquals(e.getMessage(), format("Cannot grow array beyond '%s'", MAX_ARRAY_SIZE)); - } + assertEquals(calculateNewArraySize(200), 300); + assertEquals(calculateNewArraySize(200, 10), 300); + assertEquals(calculateNewArraySize(200, 500), 500); + + assertEquals(calculateNewArraySize(MAX_ARRAY_SIZE - 1), MAX_ARRAY_SIZE); + assertEquals(calculateNewArraySize(10, MAX_ARRAY_SIZE), MAX_ARRAY_SIZE); + + assertEquals(calculateNewArraySize(1, 0), 64); + + assertThatThrownBy(() -> calculateNewArraySize(Integer.MAX_VALUE)) + .isInstanceOf(IllegalArgumentException.class); + assertThatThrownBy(() -> calculateNewArraySize(0, Integer.MAX_VALUE)) + .isInstanceOf(IllegalArgumentException.class); + assertThatThrownBy(() -> calculateNewArraySize(MAX_ARRAY_SIZE)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessage("Cannot grow array beyond size %d".formatted(MAX_ARRAY_SIZE)); } } diff --git a/core/trino-spi/src/test/java/io/trino/spi/block/TestLazyBlock.java b/core/trino-spi/src/test/java/io/trino/spi/block/TestLazyBlock.java index 45b3dbed9877..fae6ff7b3437 100644 --- a/core/trino-spi/src/test/java/io/trino/spi/block/TestLazyBlock.java +++ b/core/trino-spi/src/test/java/io/trino/spi/block/TestLazyBlock.java @@ -69,21 +69,21 @@ public void testNestedGetLoadedBlock() List actualNotifications = new ArrayList<>(); Block arrayBlock = new IntArrayBlock(1, Optional.empty(), new int[] {0}); LazyBlock lazyArrayBlock = new LazyBlock(1, () -> arrayBlock); - Block dictionaryBlock = DictionaryBlock.create(2, lazyArrayBlock, new int[] {0, 0}); - LazyBlock lazyBlock = new LazyBlock(2, () -> dictionaryBlock); + Block rowBlock = RowBlock.fromFieldBlocks(2, new Block[] {lazyArrayBlock}); + LazyBlock lazyBlock = new LazyBlock(2, () -> rowBlock); LazyBlock.listenForLoads(lazyBlock, actualNotifications::add); Block loadedBlock = lazyBlock.getBlock(); - assertThat(loadedBlock).isInstanceOf(DictionaryBlock.class); - assertThat(((DictionaryBlock) loadedBlock).getDictionary()).isInstanceOf(LazyBlock.class); + assertThat(loadedBlock).isInstanceOf(RowBlock.class); + assertThat(((RowBlock) loadedBlock).getFieldBlock(0)).isInstanceOf(LazyBlock.class); assertEquals(actualNotifications, ImmutableList.of(loadedBlock)); Block fullyLoadedBlock = lazyBlock.getLoadedBlock(); - assertThat(fullyLoadedBlock).isInstanceOf(DictionaryBlock.class); - assertThat(((DictionaryBlock) fullyLoadedBlock).getDictionary()).isInstanceOf(IntArrayBlock.class); + assertThat(fullyLoadedBlock).isInstanceOf(RowBlock.class); + assertThat(((RowBlock) fullyLoadedBlock).getFieldBlock(0)).isInstanceOf(IntArrayBlock.class); assertEquals(actualNotifications, ImmutableList.of(loadedBlock, arrayBlock)); assertTrue(lazyBlock.isLoaded()); - assertTrue(dictionaryBlock.isLoaded()); + assertTrue(rowBlock.isLoaded()); } private static void assertNotificationsRecursive(int depth, Block lazyBlock, List actualNotifications, List expectedNotifications) diff --git a/core/trino-spi/src/test/java/io/trino/spi/block/TestRowBlock.java b/core/trino-spi/src/test/java/io/trino/spi/block/TestRowBlock.java index 2b7f14637d77..b096eeb41740 100644 --- a/core/trino-spi/src/test/java/io/trino/spi/block/TestRowBlock.java +++ b/core/trino-spi/src/test/java/io/trino/spi/block/TestRowBlock.java @@ -29,7 +29,7 @@ public void testFieldBlockOffsetsIsNullWhenThereIsNoNullRow() Block fieldBlock = new ByteArrayBlock(1, Optional.empty(), new byte[]{10}); AbstractRowBlock rowBlock = (RowBlock) RowBlock.fromFieldBlocks(1, Optional.empty(), new Block[] {fieldBlock}); // Blocks should discard the offset mask during creation if no values are null - assertNull(rowBlock.getFieldBlockOffsets()); + assertNull(rowBlock.getRowIsNull()); } @Test @@ -38,6 +38,6 @@ public void testFieldBlockOffsetsIsNotNullWhenThereIsNullRow() Block fieldBlock = new ByteArrayBlock(1, Optional.empty(), new byte[]{10}); AbstractRowBlock rowBlock = (RowBlock) RowBlock.fromFieldBlocks(1, Optional.of(new boolean[] {true}), new Block[] {fieldBlock}); // Blocks should not discard the offset mask during creation if no values are null - assertNotNull(rowBlock.getFieldBlockOffsets()); + assertNotNull(rowBlock.getRowIsNull()); } } diff --git a/lib/trino-cache/src/main/java/io/trino/cache/SafeCaches.java b/lib/trino-cache/src/main/java/io/trino/cache/SafeCaches.java index a5902fe4d2c2..caf472d7dfb8 100644 --- a/lib/trino-cache/src/main/java/io/trino/cache/SafeCaches.java +++ b/lib/trino-cache/src/main/java/io/trino/cache/SafeCaches.java @@ -70,4 +70,9 @@ private static LoadingCache buildUnsafeCache(CacheBuilder LoadingCache emptyLoadingCache(CacheLoader cacheLoader, boolean recordStats) + { + return new EmptyCache<>(cacheLoader, recordStats); + } } diff --git a/lib/trino-filesystem-azure/src/main/java/io/trino/filesystem/azure/AzureFileSystem.java b/lib/trino-filesystem-azure/src/main/java/io/trino/filesystem/azure/AzureFileSystem.java index b30143bfa914..33118519086e 100644 --- a/lib/trino-filesystem-azure/src/main/java/io/trino/filesystem/azure/AzureFileSystem.java +++ b/lib/trino-filesystem-azure/src/main/java/io/trino/filesystem/azure/AzureFileSystem.java @@ -33,10 +33,12 @@ import com.azure.storage.file.datalake.models.ListPathsOptions; import com.azure.storage.file.datalake.models.PathItem; import com.azure.storage.file.datalake.options.DataLakePathDeleteOptions; +import com.google.common.collect.ImmutableSet; import io.airlift.units.DataSize; import io.trino.filesystem.FileIterator; import io.trino.filesystem.Location; import io.trino.filesystem.TrinoFileSystem; +import io.trino.filesystem.TrinoFileSystemException; import io.trino.filesystem.TrinoInputFile; import io.trino.filesystem.TrinoOutputFile; @@ -44,12 +46,15 @@ import java.util.Objects; import java.util.Optional; import java.util.OptionalLong; +import java.util.Set; import static com.azure.storage.common.implementation.Constants.HeaderConstants.ETAG_WILDCARD; import static com.google.common.base.Preconditions.checkArgument; +import static com.google.common.collect.ImmutableSet.toImmutableSet; import static io.trino.filesystem.azure.AzureUtils.handleAzureException; import static java.lang.Math.toIntExact; import static java.util.Objects.requireNonNull; +import static java.util.UUID.randomUUID; import static java.util.function.Predicate.not; public class AzureFileSystem @@ -299,6 +304,105 @@ public Optional directoryExists(Location location) } } + @Override + public void createDirectory(Location location) + throws IOException + { + AzureLocation azureLocation = new AzureLocation(location); + if (!isHierarchicalNamespaceEnabled(azureLocation)) { + return; + } + try { + DataLakeFileSystemClient fileSystemClient = createFileSystemClient(azureLocation); + DataLakeDirectoryClient directoryClient = createDirectoryIfNotExists(fileSystemClient, azureLocation.path()); + if (!directoryClient.getProperties().isDirectory()) { + throw new TrinoFileSystemException("Location is not a directory: " + azureLocation); + } + } + catch (RuntimeException e) { + throw handleAzureException(e, "creating directory", azureLocation); + } + } + + @Override + public void renameDirectory(Location source, Location target) + throws IOException + { + AzureLocation sourceLocation = new AzureLocation(source); + AzureLocation targetLocation = new AzureLocation(target); + if (!sourceLocation.account().equals(targetLocation.account())) { + throw new TrinoFileSystemException("Cannot rename across storage accounts"); + } + if (!Objects.equals(sourceLocation.container(), targetLocation.container())) { + throw new TrinoFileSystemException("Cannot rename across storage account containers"); + } + if (!isHierarchicalNamespaceEnabled(sourceLocation)) { + throw new TrinoFileSystemException("Azure non-hierarchical does not support directory renames"); + } + if (sourceLocation.path().isEmpty() || targetLocation.path().isEmpty()) { + throw new TrinoFileSystemException("Cannot rename %s to %s".formatted(source, target)); + } + + try { + DataLakeFileSystemClient fileSystemClient = createFileSystemClient(sourceLocation); + DataLakeDirectoryClient directoryClient = createDirectoryClient(fileSystemClient, sourceLocation.path()); + if (!directoryClient.exists()) { + throw new TrinoFileSystemException("Source directory does not exist: " + source); + } + if (!directoryClient.getProperties().isDirectory()) { + throw new TrinoFileSystemException("Source is not a directory: " + source); + } + directoryClient.rename(null, targetLocation.path()); + } + catch (RuntimeException e) { + throw new IOException("Rename directory from %s to %s failed".formatted(source, target), e); + } + } + + @Override + public Set listDirectories(Location location) + throws IOException + { + AzureLocation azureLocation = new AzureLocation(location); + try { + // blob API returns directories as blobs, so it cannot be used when Gen2 is enabled + return isHierarchicalNamespaceEnabled(azureLocation) + ? listGen2Directories(azureLocation) + : listBlobDirectories(azureLocation); + } + catch (RuntimeException e) { + throw handleAzureException(e, "listing files", azureLocation); + } + } + + @Override + public Optional createTemporaryDirectory(Location targetPath, String temporaryPrefix, String relativePrefix) + throws IOException + { + AzureLocation azureLocation = new AzureLocation(targetPath); + if (!isHierarchicalNamespaceEnabled(azureLocation)) { + return Optional.empty(); + } + + // allow for absolute or relative temporary prefix + Location temporary; + if (temporaryPrefix.startsWith("/")) { + String prefix = temporaryPrefix; + while (prefix.startsWith("/")) { + prefix = prefix.substring(1); + } + temporary = azureLocation.baseLocation().appendPath(prefix); + } + else { + temporary = targetPath.appendPath(temporaryPrefix); + } + + temporary = temporary.appendPath(randomUUID().toString()); + + createDirectory(temporary); + return Optional.of(temporary); + } + private boolean isHierarchicalNamespaceEnabled(AzureLocation location) throws IOException { @@ -333,6 +437,41 @@ private BlobContainerClient createBlobContainerClient(AzureLocation location) return builder.buildClient(); } + private Set listGen2Directories(AzureLocation location) + throws IOException + { + DataLakeFileSystemClient fileSystemClient = createFileSystemClient(location); + PagedIterable pathItems; + if (location.path().isEmpty()) { + pathItems = fileSystemClient.listPaths(); + } + else { + DataLakeDirectoryClient directoryClient = createDirectoryClient(fileSystemClient, location.path()); + if (!directoryClient.exists()) { + return ImmutableSet.of(); + } + if (!directoryClient.getProperties().isDirectory()) { + throw new TrinoFileSystemException("Location is not a directory: " + location); + } + pathItems = directoryClient.listPaths(false, false, null, null); + } + Location baseLocation = location.baseLocation(); + return pathItems.stream() + .filter(PathItem::isDirectory) + .map(item -> baseLocation.appendPath(item.getName() + "/")) + .collect(toImmutableSet()); + } + + private Set listBlobDirectories(AzureLocation location) + { + Location baseLocation = location.baseLocation(); + return createBlobContainerClient(location) + .listBlobsByHierarchy(location.directoryPath()).stream() + .filter(BlobItem::isPrefix) + .map(item -> baseLocation.appendPath(item.getName())) + .collect(toImmutableSet()); + } + private DataLakeFileSystemClient createFileSystemClient(AzureLocation location) { requireNonNull(location, "location is null"); @@ -348,4 +487,14 @@ private DataLakeFileSystemClient createFileSystemClient(AzureLocation location) } return fileSystemClient; } + + private static DataLakeDirectoryClient createDirectoryClient(DataLakeFileSystemClient fileSystemClient, String directoryName) + { + return fileSystemClient.getDirectoryClient(directoryName); + } + + private static DataLakeDirectoryClient createDirectoryIfNotExists(DataLakeFileSystemClient fileSystemClient, String name) + { + return fileSystemClient.createDirectoryIfNotExists(name); + } } diff --git a/lib/trino-filesystem-azure/src/main/java/io/trino/filesystem/azure/AzureLocation.java b/lib/trino-filesystem-azure/src/main/java/io/trino/filesystem/azure/AzureLocation.java index 15336a3473cc..6f23dff8ac87 100644 --- a/lib/trino-filesystem-azure/src/main/java/io/trino/filesystem/azure/AzureLocation.java +++ b/lib/trino-filesystem-azure/src/main/java/io/trino/filesystem/azure/AzureLocation.java @@ -23,25 +23,38 @@ class AzureLocation { - private static final String INVALID_LOCATION_MESSAGE = "Invalid Azure location. Expected form is 'abfs://[@].dfs.core.windows.net/': %s"; + private static final String INVALID_ABFS_LOCATION_MESSAGE = "Invalid Azure ABFS location. Expected form is 'abfs[s]://[@].dfs./': %s"; + private static final String INVALID_WASB_LOCATION_MESSAGE = "Invalid Azure WASB location. Expected form is 'wasb[s]://[@].blob./': %s"; // https://learn.microsoft.com/en-us/azure/azure-resource-manager/management/resource-name-rules private static final CharMatcher CONTAINER_VALID_CHARACTERS = CharMatcher.inRange('a', 'z').or(CharMatcher.inRange('0', '9')).or(CharMatcher.is('-')); private static final CharMatcher STORAGE_ACCOUNT_VALID_CHARACTERS = CharMatcher.inRange('a', 'z').or(CharMatcher.inRange('0', '9')); private final Location location; + private final String scheme; private final String account; + private final String endpoint; public AzureLocation(Location location) { this.location = requireNonNull(location, "location is null"); // abfss is also supported but not documented - String scheme = location.scheme().orElseThrow(() -> new IllegalArgumentException(String.format(INVALID_LOCATION_MESSAGE, location))); - checkArgument("abfs".equals(scheme) || "abfss".equals(scheme), INVALID_LOCATION_MESSAGE, location); + scheme = location.scheme().orElseThrow(() -> new IllegalArgumentException(String.format(INVALID_ABFS_LOCATION_MESSAGE, location))); + String invalidLocationMessage; + if ("abfs".equals(scheme) || "abfss".equals(scheme)) { + invalidLocationMessage = INVALID_ABFS_LOCATION_MESSAGE; + } + else if ("wasb".equals(scheme) || "wasbs".equals(scheme)) { + invalidLocationMessage = INVALID_WASB_LOCATION_MESSAGE; + } + else { + // only mention abfs in error message as the other forms are deprecated + throw new IllegalArgumentException(String.format(INVALID_ABFS_LOCATION_MESSAGE, location)); + } // container is interpolated into the URL path, so perform extra checks location.userInfo().ifPresent(container -> { - checkArgument(!container.isEmpty(), INVALID_LOCATION_MESSAGE, location); + checkArgument(!container.isEmpty(), invalidLocationMessage, location); checkArgument( CONTAINER_VALID_CHARACTERS.matchesAllOf(container), "Invalid Azure storage container name. Valid characters are 'a-z', '0-9', and '-': %s", @@ -57,17 +70,27 @@ public AzureLocation(Location location) }); // storage account is the first label of the host - checkArgument(location.host().isPresent(), INVALID_LOCATION_MESSAGE, location); + checkArgument(location.host().isPresent(), invalidLocationMessage, location); String host = location.host().get(); int accountSplit = host.indexOf('.'); checkArgument( accountSplit > 0, - INVALID_LOCATION_MESSAGE, + invalidLocationMessage, this.location); this.account = host.substring(0, accountSplit); - // host must end with ".dfs.core.windows.net" - checkArgument(host.substring(accountSplit).equals(".dfs.core.windows.net"), INVALID_LOCATION_MESSAGE, location); + // abfs[s] host must contain ".dfs.", and wasb[s] host must contain ".blob." before endpoint + if (scheme.equals("abfs") || scheme.equals("abfss")) { + checkArgument(host.substring(accountSplit).startsWith(".dfs."), invalidLocationMessage, location); + // endpoint does not include dfs + this.endpoint = host.substring(accountSplit + ".dfs.".length()); + } + else { + checkArgument(host.substring(accountSplit).startsWith(".blob."), invalidLocationMessage, location); + // endpoint does not include blob + this.endpoint = host.substring(accountSplit + ".blob.".length()); + } + checkArgument(!endpoint.isEmpty(), invalidLocationMessage, location); // storage account is interpolated into URL host name, so perform extra checks checkArgument(STORAGE_ACCOUNT_VALID_CHARACTERS.matchesAllOf(account), @@ -108,9 +131,27 @@ public String path() return location.path(); } + public String directoryPath() + { + String path = location.path(); + if (!path.isEmpty() && !path.endsWith("/")) { + path += "/"; + } + return path; + } + @Override public String toString() { return location.toString(); } + + public Location baseLocation() + { + return Location.of("%s://%s%s.dfs.%s/".formatted( + scheme, + container().map(container -> container + "@").orElse(""), + account(), + endpoint)); + } } diff --git a/lib/trino-filesystem-azure/src/main/java/io/trino/filesystem/azure/AzureOutputFile.java b/lib/trino-filesystem-azure/src/main/java/io/trino/filesystem/azure/AzureOutputFile.java index 261601c197e7..cfcd30317294 100644 --- a/lib/trino-filesystem-azure/src/main/java/io/trino/filesystem/azure/AzureOutputFile.java +++ b/lib/trino-filesystem-azure/src/main/java/io/trino/filesystem/azure/AzureOutputFile.java @@ -13,11 +13,15 @@ */ package io.trino.filesystem.azure; +import com.azure.core.util.BinaryData; import com.azure.storage.blob.BlobClient; +import com.azure.storage.blob.models.BlobErrorCode; +import com.azure.storage.blob.models.BlobStorageException; import io.trino.filesystem.Location; import io.trino.filesystem.TrinoOutputFile; import io.trino.memory.context.AggregatedMemoryContext; +import java.io.FileNotFoundException; import java.io.IOException; import java.io.OutputStream; import java.nio.file.FileAlreadyExistsException; @@ -52,6 +56,21 @@ public boolean exists() return blobClient.exists(); } + @Override + public void createOrOverwrite(byte[] data) + throws IOException + { + try { + blobClient.getBlockBlobClient().upload(BinaryData.fromBytes(data), true); + } + catch (BlobStorageException e) { + if (BlobErrorCode.CONTAINER_NOT_FOUND.equals(e.getErrorCode())) { + throw new FileNotFoundException(location.toString()); + } + throw e; + } + } + @Override public OutputStream create(AggregatedMemoryContext memoryContext) throws IOException diff --git a/lib/trino-filesystem-manager/pom.xml b/lib/trino-filesystem-manager/pom.xml index 7b5347fb44ab..041bbe07c6c9 100644 --- a/lib/trino-filesystem-manager/pom.xml +++ b/lib/trino-filesystem-manager/pom.xml @@ -49,12 +49,12 @@ io.trino - trino-hdfs + trino-spi - io.trino - trino-spi + jakarta.annotation + jakarta.annotation-api diff --git a/lib/trino-filesystem-manager/src/main/java/io/trino/filesystem/manager/FileSystemModule.java b/lib/trino-filesystem-manager/src/main/java/io/trino/filesystem/manager/FileSystemModule.java index d499a628e0ef..0ad32bb4ee1b 100644 --- a/lib/trino-filesystem-manager/src/main/java/io/trino/filesystem/manager/FileSystemModule.java +++ b/lib/trino-filesystem-manager/src/main/java/io/trino/filesystem/manager/FileSystemModule.java @@ -14,37 +14,82 @@ package io.trino.filesystem.manager; import com.google.inject.Binder; -import com.google.inject.Inject; import com.google.inject.Provides; +import com.google.inject.Scopes; import com.google.inject.Singleton; import io.airlift.configuration.AbstractConfigurationAwareModule; +import io.airlift.configuration.ConfigurationFactory; +import io.opentelemetry.api.OpenTelemetry; import io.opentelemetry.api.trace.Tracer; +import io.trino.filesystem.Location; import io.trino.filesystem.TrinoFileSystemFactory; -import io.trino.filesystem.hdfs.HdfsFileSystemFactory; -import io.trino.filesystem.hdfs.HdfsFileSystemModule; +import io.trino.filesystem.cache.CacheFileSystemFactory; +import io.trino.filesystem.cache.CacheKeyProvider; +import io.trino.filesystem.cache.CachingHostAddressProvider; +import io.trino.filesystem.cache.DefaultCacheKeyProvider; +import io.trino.filesystem.cache.DefaultCachingHostAddressProvider; +import io.trino.filesystem.cache.TrinoFileSystemCache; +import io.trino.filesystem.memory.MemoryFileSystemCache; +import io.trino.filesystem.memory.MemoryFileSystemCacheModule; import io.trino.filesystem.s3.S3FileSystemFactory; import io.trino.filesystem.s3.S3FileSystemModule; +import io.trino.filesystem.switching.SwitchingFileSystemFactory; import io.trino.filesystem.tracing.TracingFileSystemFactory; -import io.trino.hdfs.s3.HiveS3Module; +import io.trino.spi.NodeManager; import java.util.Map; import java.util.Optional; +import java.util.function.Function; -import static com.google.inject.Scopes.SINGLETON; import static com.google.inject.multibindings.MapBinder.newMapBinder; +import static com.google.inject.multibindings.OptionalBinder.newOptionalBinder; +import static java.util.Objects.requireNonNull; public class FileSystemModule extends AbstractConfigurationAwareModule { + private final String catalogName; + private final NodeManager nodeManager; + private final OpenTelemetry openTelemetry; + private final boolean coordinatorFileCaching; + private ConfigurationFactory configurationFactory; + + public FileSystemModule(String catalogName, NodeManager nodeManager, OpenTelemetry openTelemetry, boolean coordinatorFileCaching) + { + this.catalogName = requireNonNull(catalogName, "catalogName is null"); + this.nodeManager = requireNonNull(nodeManager, "nodeManager is null"); + this.openTelemetry = requireNonNull(openTelemetry, "openTelemetry is null"); + this.coordinatorFileCaching = coordinatorFileCaching; + } + + @Override + public synchronized void setConfigurationFactory(ConfigurationFactory configurationFactory) + { + this.configurationFactory = requireNonNull(configurationFactory, "configurationFactory is null"); + super.setConfigurationFactory(configurationFactory); + } + @Override protected void setup(Binder binder) { FileSystemConfig config = buildConfigObject(FileSystemConfig.class); - binder.bind(HdfsFileSystemFactoryHolder.class).in(SINGLETON); + newOptionalBinder(binder, HdfsFileSystemLoader.class); if (config.isHadoopEnabled()) { - install(new HdfsFileSystemModule()); + HdfsFileSystemLoader loader = new HdfsFileSystemLoader( + configurationFactory.getProperties(), + false, + false, + !config.isNativeS3Enabled(), + catalogName, + nodeManager, + openTelemetry); + + loader.configure().forEach((name, securitySensitive) -> + configurationFactory.consumeProperty(name)); + + binder.bind(HdfsFileSystemLoader.class).toInstance(loader); } var factories = newMapBinder(binder, String.class, TrinoFileSystemFactory.class); @@ -55,30 +100,47 @@ protected void setup(Binder binder) factories.addBinding("s3a").to(S3FileSystemFactory.class); factories.addBinding("s3n").to(S3FileSystemFactory.class); } - else { - install(new HiveS3Module()); + + newOptionalBinder(binder, CachingHostAddressProvider.class).setDefault().to(DefaultCachingHostAddressProvider.class).in(Scopes.SINGLETON); + newOptionalBinder(binder, CacheKeyProvider.class).setDefault().to(DefaultCacheKeyProvider.class).in(Scopes.SINGLETON); + + newOptionalBinder(binder, TrinoFileSystemCache.class); + newOptionalBinder(binder, MemoryFileSystemCache.class); + + boolean isCoordinator = nodeManager.getCurrentNode().isCoordinator(); + if (coordinatorFileCaching) { + install(new MemoryFileSystemCacheModule(isCoordinator)); } + + newOptionalBinder(binder, CachingHostAddressProvider.class).setDefault().to(DefaultCachingHostAddressProvider.class).in(Scopes.SINGLETON); } @Provides @Singleton - public TrinoFileSystemFactory createFileSystemFactory( - HdfsFileSystemFactoryHolder hdfsFileSystemFactory, + static TrinoFileSystemFactory createFileSystemFactory( + Optional hdfsFileSystemLoader, Map factories, + Optional fileSystemCache, + Optional memoryFileSystemCache, + Optional keyProvider, Tracer tracer) { - TrinoFileSystemFactory delegate = new SwitchingFileSystemFactory(hdfsFileSystemFactory.value(), factories); - return new TracingFileSystemFactory(tracer, delegate); - } + Optional hdfsFactory = hdfsFileSystemLoader.map(HdfsFileSystemLoader::create); - public static class HdfsFileSystemFactoryHolder - { - @Inject(optional = true) - private HdfsFileSystemFactory hdfsFileSystemFactory; + Function loader = location -> location.scheme() + .map(factories::get) + .or(() -> hdfsFactory) + .orElseThrow(() -> new IllegalArgumentException("No factory for location: " + location)); - public Optional value() - { - return Optional.ofNullable(hdfsFileSystemFactory); + TrinoFileSystemFactory delegate = new SwitchingFileSystemFactory(loader); + delegate = new TracingFileSystemFactory(tracer, delegate); + if (fileSystemCache.isPresent()) { + return new CacheFileSystemFactory(tracer, delegate, fileSystemCache.orElseThrow(), keyProvider.orElseThrow()); + } + // use MemoryFileSystemCache only when no other TrinoFileSystemCache is configured + if (memoryFileSystemCache.isPresent()) { + return new CacheFileSystemFactory(tracer, delegate, memoryFileSystemCache.orElseThrow(), keyProvider.orElseThrow()); } + return delegate; } } diff --git a/lib/trino-filesystem-manager/src/main/java/io/trino/filesystem/manager/HdfsClassLoader.java b/lib/trino-filesystem-manager/src/main/java/io/trino/filesystem/manager/HdfsClassLoader.java new file mode 100644 index 000000000000..172e916ebc15 --- /dev/null +++ b/lib/trino-filesystem-manager/src/main/java/io/trino/filesystem/manager/HdfsClassLoader.java @@ -0,0 +1,123 @@ +/* + * 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.filesystem.manager; + +import java.io.IOException; +import java.net.URL; +import java.net.URLClassLoader; +import java.util.Enumeration; +import java.util.List; + +import static com.google.common.base.Preconditions.checkArgument; + +// based on io.trino.server.PluginClassLoader +final class HdfsClassLoader + extends URLClassLoader +{ + public HdfsClassLoader(List urls) + { + // This class loader should not have access to the system (application) class loader + super(urls.toArray(URL[]::new), getPlatformClassLoader()); + } + + @Override + protected Class loadClass(String name, boolean resolve) + throws ClassNotFoundException + { + synchronized (getClassLoadingLock(name)) { + // Check if class is in the loaded classes cache + Class cachedClass = findLoadedClass(name); + if (cachedClass != null) { + return resolveClass(cachedClass, resolve); + } + + // If this is an override class, only check override class loader + if (isOverrideClass(name)) { + return resolveClass(overrideClassLoader().loadClass(name), resolve); + } + + // Look for class locally + return super.loadClass(name, resolve); + } + } + + private Class resolveClass(Class clazz, boolean resolve) + { + if (resolve) { + resolveClass(clazz); + } + return clazz; + } + + @Override + public URL getResource(String name) + { + // If this is an override resource, only check override class loader + if (isOverrideResource(name)) { + return overrideClassLoader().getResource(name); + } + + // Look for resource locally + return super.getResource(name); + } + + @Override + public Enumeration getResources(String name) + throws IOException + { + // If this is an override resource, use override resources + if (isOverrideResource(name)) { + return overrideClassLoader().getResources(name); + } + + // Use local resources + return super.getResources(name); + } + + private ClassLoader overrideClassLoader() + { + return getClass().getClassLoader(); + } + + private static boolean isOverrideResource(String name) + { + return isOverrideClass(name.replace('.', '/')); + } + + private static boolean isOverrideClass(String name) + { + // SPI packages from io.trino.server.PluginManager and dependencies of trino-filesystem + return hasPackage(name, "io.trino.spi.") || + hasPackage(name, "com.fasterxml.jackson.annotation.") || + hasPackage(name, "io.airlift.slice.") || + hasPackage(name, "org.openjdk.jol.") || + hasPackage(name, "io.opentelemetry.api.") || + hasPackage(name, "io.opentelemetry.context.") || + hasPackage(name, "com.google.common.") || + hasExactPackage(name, "io.trino.memory.context.") || + hasExactPackage(name, "io.trino.filesystem."); + } + + private static boolean hasPackage(String name, String packageName) + { + checkArgument(!packageName.isEmpty() && packageName.charAt(packageName.length() - 1) == '.'); + return name.startsWith(packageName); + } + + private static boolean hasExactPackage(String name, String packageName) + { + checkArgument(!packageName.isEmpty() && packageName.charAt(packageName.length() - 1) == '.'); + return name.startsWith(packageName) && (name.lastIndexOf('.') == (packageName.length() - 1)); + } +} diff --git a/lib/trino-filesystem-manager/src/main/java/io/trino/filesystem/manager/HdfsFileSystemLoader.java b/lib/trino-filesystem-manager/src/main/java/io/trino/filesystem/manager/HdfsFileSystemLoader.java new file mode 100644 index 000000000000..6c790c26b624 --- /dev/null +++ b/lib/trino-filesystem-manager/src/main/java/io/trino/filesystem/manager/HdfsFileSystemLoader.java @@ -0,0 +1,176 @@ +/* + * 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.filesystem.manager; + +import io.opentelemetry.api.OpenTelemetry; +import io.trino.filesystem.TrinoFileSystemFactory; +import io.trino.spi.NodeManager; +import io.trino.spi.Plugin; +import io.trino.spi.classloader.ThreadContextClassLoader; +import jakarta.annotation.PreDestroy; + +import java.io.File; +import java.io.IOException; +import java.io.UncheckedIOException; +import java.net.MalformedURLException; +import java.net.URISyntaxException; +import java.net.URL; +import java.nio.file.DirectoryStream; +import java.nio.file.Path; +import java.util.List; +import java.util.Map; + +import static com.google.common.base.Verify.verify; +import static com.google.common.collect.Streams.stream; +import static java.nio.file.Files.newDirectoryStream; + +final class HdfsFileSystemLoader +{ + private final HdfsClassLoader classLoader; + private final Object manager; + + public HdfsFileSystemLoader( + Map config, + boolean azureEnabled, + boolean gcsEnabled, + boolean s3Enabled, + String catalogName, + NodeManager nodeManager, + OpenTelemetry openTelemetry) + { + Class clazz = tryLoadExistingHdfsManager(); + + // check if we are running inside a plugin class loader (full server mode) + if (!getClass().getClassLoader().equals(Plugin.class.getClassLoader())) { + verify(clazz == null, "HDFS should not be on the plugin classpath"); + File sourceFile = getCurrentClassLocation(); + File directory; + if (sourceFile.isDirectory()) { + // running DevelopmentServer in the IDE + verify(sourceFile.getPath().endsWith("/target/classes"), "Source file not in 'target' directory: %s", sourceFile); + directory = new File(sourceFile.getParentFile().getParentFile().getParentFile(), "trino-hdfs/target/hdfs"); + } + else { + // normal server mode where HDFS JARs are in a subdirectory of the plugin + directory = new File(sourceFile.getParentFile(), "hdfs"); + } + verify(directory.isDirectory(), "HDFS directory is missing: %s", directory); + classLoader = createClassLoader(directory); + clazz = loadHdfsManager(classLoader); + } + else { + verify(clazz != null, "HDFS should be on the classpath for tests"); + classLoader = null; + } + + try (var ignore = new ThreadContextClassLoader(classLoader)) { + manager = clazz.getConstructor(Map.class, boolean.class, boolean.class, boolean.class, String.class, NodeManager.class, OpenTelemetry.class) + .newInstance(config, azureEnabled, gcsEnabled, s3Enabled, catalogName, nodeManager, openTelemetry); + } + catch (ReflectiveOperationException e) { + throw new RuntimeException(e); + } + } + + @SuppressWarnings("unchecked") + public Map configure() + { + try (var ignore = new ThreadContextClassLoader(classLoader)) { + return (Map) manager.getClass().getMethod("configure").invoke(manager); + } + catch (ReflectiveOperationException e) { + throw new RuntimeException("Failed to configure HDFS:\n%s\n%s\n%s".formatted("<".repeat(70), e.getCause(), ">".repeat(70)), e); + } + } + + public TrinoFileSystemFactory create() + { + try (var ignore = new ThreadContextClassLoader(classLoader)) { + return (TrinoFileSystemFactory) manager.getClass().getMethod("create").invoke(manager); + } + catch (ReflectiveOperationException e) { + throw new RuntimeException(e); + } + } + + @PreDestroy + public void stop() + throws IOException, ReflectiveOperationException + { + try (classLoader; var ignore = new ThreadContextClassLoader(classLoader)) { + manager.getClass().getMethod("stop").invoke(manager); + } + } + + private File getCurrentClassLocation() + { + try { + return new File(getClass().getProtectionDomain().getCodeSource().getLocation().toURI()); + } + catch (URISyntaxException e) { + throw new RuntimeException(e); + } + } + + private Class tryLoadExistingHdfsManager() + { + try { + return loadHdfsManager(getClass().getClassLoader()); + } + catch (RuntimeException e) { + return null; + } + } + + private static Class loadHdfsManager(ClassLoader classLoader) + { + try { + return classLoader.loadClass("io.trino.filesystem.hdfs.HdfsFileSystemManager"); + } + catch (ClassNotFoundException e) { + throw new RuntimeException(e); + } + } + + private static HdfsClassLoader createClassLoader(File path) + { + List urls = buildClassPath(path); + verify(!urls.isEmpty(), "HDFS directory is empty: %s", path); + return new HdfsClassLoader(urls); + } + + private static List buildClassPath(File path) + { + try (DirectoryStream directoryStream = newDirectoryStream(path.toPath())) { + return stream(directoryStream) + .map(Path::toFile) + .sorted().toList().stream() + .map(HdfsFileSystemLoader::fileToUrl) + .toList(); + } + catch (IOException e) { + throw new UncheckedIOException(e); + } + } + + private static URL fileToUrl(File file) + { + try { + return file.toURI().toURL(); + } + catch (MalformedURLException e) { + throw new UncheckedIOException(e); + } + } +} diff --git a/lib/trino-filesystem-s3/pom.xml b/lib/trino-filesystem-s3/pom.xml index 180fe53c7226..4909b89002de 100644 --- a/lib/trino-filesystem-s3/pom.xml +++ b/lib/trino-filesystem-s3/pom.xml @@ -17,6 +17,16 @@ + + com.fasterxml.jackson.core + jackson-annotations + + + + com.google.errorprone + error_prone_annotations + + com.google.guava guava @@ -27,16 +37,46 @@ guice + + io.airlift + concurrent + + io.airlift configuration + + io.airlift + http-client + + + + io.airlift + log + + + + io.airlift + stats + + io.airlift units + + io.opentelemetry + opentelemetry-api + + + + io.opentelemetry.instrumentation + opentelemetry-aws-sdk-2.2 + + io.trino trino-filesystem @@ -47,6 +87,17 @@ trino-memory-context + + io.trino + trino-plugin-toolkit + + + io.airlift + bootstrap + + + + io.trino trino-spi @@ -62,6 +113,11 @@ jakarta.validation-api + + org.weakref + jmxutils + + software.amazon.awssdk apache-client @@ -88,11 +144,26 @@ http-client-spi + + software.amazon.awssdk + metrics-spi + + software.amazon.awssdk regions + + software.amazon.awssdk + retries + + + + software.amazon.awssdk + retries-spi + + software.amazon.awssdk s3 diff --git a/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/AwsSdkV2ApiCallStats.java b/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/AwsSdkV2ApiCallStats.java new file mode 100644 index 000000000000..e4fab155ccf8 --- /dev/null +++ b/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/AwsSdkV2ApiCallStats.java @@ -0,0 +1,107 @@ +/* + * 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.filesystem.s3; + +import com.google.errorprone.annotations.ThreadSafe; +import io.airlift.stats.CounterStat; +import io.airlift.stats.TimeStat; +import org.weakref.jmx.Managed; +import org.weakref.jmx.Nested; + +import java.time.Duration; + +import static java.util.concurrent.TimeUnit.MILLISECONDS; + +@ThreadSafe +public class AwsSdkV2ApiCallStats +{ + private final TimeStat latency = new TimeStat(MILLISECONDS); + private final CounterStat calls = new CounterStat(); + private final CounterStat failures = new CounterStat(); + private final CounterStat retries = new CounterStat(); + private final CounterStat throttlingExceptions = new CounterStat(); + private final CounterStat serverErrors = new CounterStat(); + + @Managed + @Nested + public TimeStat getLatency() + { + return latency; + } + + @Managed + @Nested + public CounterStat getCalls() + { + return calls; + } + + @Managed + @Nested + public CounterStat getFailures() + { + return failures; + } + + @Managed + @Nested + public CounterStat getRetries() + { + return retries; + } + + @Managed + @Nested + public CounterStat getThrottlingExceptions() + { + return throttlingExceptions; + } + + @Managed + @Nested + public CounterStat getServerErrors() + { + return serverErrors; + } + + public void updateLatency(Duration duration) + { + latency.addNanos(duration.toNanos()); + } + + public void updateCalls() + { + calls.update(1); + } + + public void updateFailures() + { + failures.update(1); + } + + public void updateRetries(int retryCount) + { + retries.update(retryCount); + } + + public void updateThrottlingExceptions() + { + throttlingExceptions.update(1); + } + + public void updateServerErrors() + { + serverErrors.update(1); + } +} diff --git a/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/AwsSdkV2HttpClientStats.java b/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/AwsSdkV2HttpClientStats.java new file mode 100644 index 000000000000..043738e1ec71 --- /dev/null +++ b/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/AwsSdkV2HttpClientStats.java @@ -0,0 +1,80 @@ +/* + * 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.filesystem.s3; + +import com.google.errorprone.annotations.ThreadSafe; +import io.airlift.stats.TimeStat; +import org.weakref.jmx.Managed; +import org.weakref.jmx.Nested; +import software.amazon.awssdk.metrics.SdkMetric; + +import java.time.Duration; +import java.util.concurrent.atomic.AtomicLong; + +import static java.util.concurrent.TimeUnit.MILLISECONDS; +import static software.amazon.awssdk.http.HttpMetric.AVAILABLE_CONCURRENCY; +import static software.amazon.awssdk.http.HttpMetric.LEASED_CONCURRENCY; +import static software.amazon.awssdk.http.HttpMetric.PENDING_CONCURRENCY_ACQUIRES; + +@ThreadSafe +public class AwsSdkV2HttpClientStats +{ + private final TimeStat connectionAcquireLatency = new TimeStat(MILLISECONDS); + private final AtomicLong availableConcurrency = new AtomicLong(); + private final AtomicLong leasedConcurrency = new AtomicLong(); + private final AtomicLong pendingConcurrencyAcquires = new AtomicLong(); + + @Managed + @Nested + public TimeStat getConnectionAcquireLatency() + { + return connectionAcquireLatency; + } + + @Managed + public long getAvailableConcurrency() + { + return availableConcurrency.get(); + } + + @Managed + public long getLeasedConcurrency() + { + return leasedConcurrency.get(); + } + + @Managed + public long getPendingConcurrencyAcquires() + { + return pendingConcurrencyAcquires.get(); + } + + public void updateConcurrencyStats(SdkMetric metric, int value) + { + if (metric.equals(AVAILABLE_CONCURRENCY)) { + availableConcurrency.set(value); + } + else if (metric.equals(PENDING_CONCURRENCY_ACQUIRES)) { + pendingConcurrencyAcquires.set(value); + } + else if (metric.equals(LEASED_CONCURRENCY)) { + leasedConcurrency.set(value); + } + } + + public void updateConcurrencyAcquireDuration(Duration duration) + { + connectionAcquireLatency.addNanos(duration.toNanos()); + } +} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/ForGlueColumnStatisticsRead.java b/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/FileSystemS3.java similarity index 91% rename from plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/ForGlueColumnStatisticsRead.java rename to lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/FileSystemS3.java index 5636c42a9354..b40f5a6793f8 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/ForGlueColumnStatisticsRead.java +++ b/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/FileSystemS3.java @@ -11,7 +11,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package io.trino.plugin.hive.metastore.glue; +package io.trino.filesystem.s3; import com.google.inject.BindingAnnotation; @@ -26,4 +26,4 @@ @Retention(RUNTIME) @Target({FIELD, PARAMETER, METHOD}) @BindingAnnotation -public @interface ForGlueColumnStatisticsRead {} +public @interface FileSystemS3 {} diff --git a/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3Context.java b/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3Context.java index 026acde55240..2a9b3da413b0 100644 --- a/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3Context.java +++ b/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3Context.java @@ -13,25 +13,115 @@ */ package io.trino.filesystem.s3; +import io.trino.filesystem.s3.S3FileSystemConfig.ObjectCannedAcl; import io.trino.filesystem.s3.S3FileSystemConfig.S3SseType; +import io.trino.filesystem.s3.S3FileSystemConfig.StorageClassType; +import io.trino.spi.security.ConnectorIdentity; +import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider; +import software.amazon.awssdk.auth.credentials.AwsSessionCredentials; +import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider; +import software.amazon.awssdk.awscore.AwsRequestOverrideConfiguration; import software.amazon.awssdk.services.s3.model.RequestPayer; +import java.util.Optional; + import static com.google.common.base.Preconditions.checkArgument; +import static io.trino.filesystem.s3.S3FileSystemConfig.S3SseType.CUSTOMER; +import static io.trino.filesystem.s3.S3FileSystemConfig.S3SseType.KMS; +import static io.trino.filesystem.s3.S3FileSystemConstants.EXTRA_CREDENTIALS_ACCESS_KEY_PROPERTY; +import static io.trino.filesystem.s3.S3FileSystemConstants.EXTRA_CREDENTIALS_SECRET_KEY_PROPERTY; +import static io.trino.filesystem.s3.S3FileSystemConstants.EXTRA_CREDENTIALS_SESSION_TOKEN_PROPERTY; import static java.util.Objects.requireNonNull; -record S3Context(int partSize, boolean requesterPays, S3SseType sseType, String sseKmsKeyId) +record S3Context( + int partSize, + boolean requesterPays, + S3SseContext s3SseContext, + Optional credentialsProviderOverride, + StorageClassType storageClass, + ObjectCannedAcl cannedAcl, + boolean exclusiveWriteSupported) { private static final int MIN_PART_SIZE = 5 * 1024 * 1024; // S3 requirement public S3Context { checkArgument(partSize >= MIN_PART_SIZE, "partSize must be at least %s bytes", MIN_PART_SIZE); - requireNonNull(sseType, "sseType is null"); - checkArgument((sseType != S3SseType.KMS) || (sseKmsKeyId != null), "sseKmsKeyId is null for SSE-KMS"); + requireNonNull(s3SseContext, "sseContext is null"); + requireNonNull(credentialsProviderOverride, "credentialsProviderOverride is null"); } public RequestPayer requestPayer() { return requesterPays ? RequestPayer.REQUESTER : null; } + + public S3Context withKmsKeyId(String kmsKeyId) + { + return new S3Context(partSize, requesterPays, S3SseContext.withKmsKeyId(kmsKeyId), credentialsProviderOverride, storageClass, cannedAcl, exclusiveWriteSupported); + } + + public S3Context withCredentials(ConnectorIdentity identity) + { + if (identity.getExtraCredentials().containsKey(EXTRA_CREDENTIALS_ACCESS_KEY_PROPERTY)) { + AwsCredentialsProvider credentialsProvider = StaticCredentialsProvider.create(AwsSessionCredentials.create( + identity.getExtraCredentials().get(EXTRA_CREDENTIALS_ACCESS_KEY_PROPERTY), + identity.getExtraCredentials().get(EXTRA_CREDENTIALS_SECRET_KEY_PROPERTY), + identity.getExtraCredentials().get(EXTRA_CREDENTIALS_SESSION_TOKEN_PROPERTY))); + return withCredentialsProviderOverride(credentialsProvider); + } + return this; + } + + public S3Context withSseCustomerKey(String key) + { + return new S3Context(partSize, requesterPays, S3SseContext.withSseCustomerKey(key), credentialsProviderOverride, storageClass, cannedAcl, exclusiveWriteSupported); + } + + public S3Context withCredentialsProviderOverride(AwsCredentialsProvider credentialsProviderOverride) + { + return new S3Context( + partSize, + requesterPays, + s3SseContext, + Optional.of(credentialsProviderOverride), + storageClass, + cannedAcl, + exclusiveWriteSupported); + } + + public void applyCredentialProviderOverride(AwsRequestOverrideConfiguration.Builder builder) + { + credentialsProviderOverride.ifPresent(builder::credentialsProvider); + } + + record S3SseContext(S3SseType sseType, Optional sseKmsKeyId, Optional sseCustomerKey) + { + S3SseContext + { + requireNonNull(sseType, "sseType is null"); + requireNonNull(sseKmsKeyId, "sseKmsKeyId is null"); + requireNonNull(sseCustomerKey, "sseCustomerKey is null"); + switch (sseType) { + case KMS -> checkArgument(sseKmsKeyId.isPresent(), "sseKmsKeyId is missing for SSE-KMS"); + case CUSTOMER -> checkArgument(sseCustomerKey.isPresent(), "sseCustomerKey is missing for SSE-C"); + case NONE, S3 -> { /* no additional checks */ } + } + } + + public static S3SseContext of(S3SseType sseType, String sseKmsKeyId, String sseCustomerKey) + { + return new S3SseContext(sseType, Optional.ofNullable(sseKmsKeyId), Optional.ofNullable(sseCustomerKey).map(S3SseCustomerKey::onAes256)); + } + + public static S3SseContext withKmsKeyId(String kmsKeyId) + { + return new S3SseContext(KMS, Optional.ofNullable(kmsKeyId), Optional.empty()); + } + + public static S3SseContext withSseCustomerKey(String key) + { + return new S3SseContext(CUSTOMER, Optional.empty(), Optional.ofNullable(key).map(S3SseCustomerKey::onAes256)); + } + } } diff --git a/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3FileSystem.java b/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3FileSystem.java index 48a74456ea4a..45b636d12885 100644 --- a/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3FileSystem.java +++ b/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3FileSystem.java @@ -15,23 +15,35 @@ import com.google.common.collect.HashMultimap; import com.google.common.collect.SetMultimap; +import io.airlift.units.Duration; import io.trino.filesystem.FileIterator; import io.trino.filesystem.Location; import io.trino.filesystem.TrinoFileSystem; +import io.trino.filesystem.TrinoFileSystemException; import io.trino.filesystem.TrinoInputFile; import io.trino.filesystem.TrinoOutputFile; +import io.trino.filesystem.UriLocation; +import io.trino.filesystem.encryption.EncryptionKey; +import software.amazon.awssdk.auth.signer.AwsS3V4Signer; +import software.amazon.awssdk.awscore.AwsRequestOverrideConfiguration; import software.amazon.awssdk.core.exception.SdkException; import software.amazon.awssdk.services.s3.S3Client; +import software.amazon.awssdk.services.s3.model.CommonPrefix; import software.amazon.awssdk.services.s3.model.DeleteObjectRequest; import software.amazon.awssdk.services.s3.model.DeleteObjectsRequest; import software.amazon.awssdk.services.s3.model.DeleteObjectsResponse; +import software.amazon.awssdk.services.s3.model.GetObjectRequest; import software.amazon.awssdk.services.s3.model.ListObjectsV2Request; import software.amazon.awssdk.services.s3.model.ObjectIdentifier; import software.amazon.awssdk.services.s3.model.RequestPayer; import software.amazon.awssdk.services.s3.model.S3Error; -import software.amazon.awssdk.services.s3.paginators.ListObjectsV2Iterable; +import software.amazon.awssdk.services.s3.model.S3Object; +import software.amazon.awssdk.services.s3.presigner.S3Presigner; +import software.amazon.awssdk.services.s3.presigner.model.GetObjectPresignRequest; +import software.amazon.awssdk.services.s3.presigner.model.PresignedGetObjectRequest; import java.io.IOException; +import java.net.URISyntaxException; import java.util.ArrayList; import java.util.Collection; import java.util.HashMap; @@ -39,21 +51,35 @@ import java.util.Map; import java.util.Map.Entry; import java.util.Optional; +import java.util.Set; +import java.util.concurrent.Executor; +import java.util.stream.Stream; +import static com.google.common.base.Verify.verify; +import static com.google.common.collect.ImmutableSet.toImmutableSet; import static com.google.common.collect.Iterables.partition; import static com.google.common.collect.Multimaps.toMultimap; +import static io.trino.filesystem.s3.S3FileSystemConfig.S3SseType.NONE; +import static io.trino.filesystem.s3.S3SseCUtils.encoded; +import static io.trino.filesystem.s3.S3SseCUtils.md5Checksum; +import static io.trino.filesystem.s3.S3SseRequestConfigurator.setEncryptionSettings; import static java.util.Objects.requireNonNull; +import static java.util.stream.Collectors.toMap; final class S3FileSystem implements TrinoFileSystem { + private final Executor uploadExecutor; private final S3Client client; + private final S3Presigner preSigner; private final S3Context context; private final RequestPayer requestPayer; - public S3FileSystem(S3Client client, S3Context context) + public S3FileSystem(Executor uploadExecutor, S3Client client, S3Presigner preSigner, S3Context context) { + this.uploadExecutor = requireNonNull(uploadExecutor, "uploadExecutor is null"); this.client = requireNonNull(client, "client is null"); + this.preSigner = requireNonNull(preSigner, "preSigner is null"); this.context = requireNonNull(context, "context is null"); this.requestPayer = context.requestPayer(); } @@ -61,19 +87,19 @@ public S3FileSystem(S3Client client, S3Context context) @Override public TrinoInputFile newInputFile(Location location) { - return new S3InputFile(client, context, new S3Location(location), null); + return new S3InputFile(client, context, new S3Location(location), null, null, Optional.empty()); } @Override public TrinoInputFile newInputFile(Location location, long length) { - return new S3InputFile(client, context, new S3Location(location), length); + return new S3InputFile(client, context, new S3Location(location), length, null, Optional.empty()); } @Override public TrinoOutputFile newOutputFile(Location location) { - return new S3OutputFile(client, context, new S3Location(location)); + return new S3OutputFile(uploadExecutor, client, context, new S3Location(location), Optional.empty()); } @Override @@ -83,6 +109,7 @@ public void deleteFile(Location location) location.verifyValidFileLocation(); S3Location s3Location = new S3Location(location); DeleteObjectRequest request = DeleteObjectRequest.builder() + .overrideConfiguration(context::applyCredentialProviderOverride) .requestPayer(requestPayer) .key(s3Location.key()) .bucket(s3Location.bucket()) @@ -92,7 +119,7 @@ public void deleteFile(Location location) client.deleteObject(request); } catch (SdkException e) { - throw new IOException("Failed to delete file: " + location, e); + throw new TrinoFileSystemException("Failed to delete file: " + location, e); } } @@ -100,13 +127,13 @@ public void deleteFile(Location location) public void deleteDirectory(Location location) throws IOException { - FileIterator iterator = listFiles(location); + FileIterator iterator = listObjects(location, true); while (iterator.hasNext()) { List files = new ArrayList<>(); while ((files.size() < 1000) && iterator.hasNext()) { files.add(iterator.next().location()); } - deleteFiles(files); + deleteObjects(files); } } @@ -115,7 +142,12 @@ public void deleteFiles(Collection locations) throws IOException { locations.forEach(Location::verifyValidFileLocation); + deleteObjects(locations); + } + private void deleteObjects(Collection locations) + throws IOException + { SetMultimap bucketToKeys = locations.stream() .map(S3Location::new) .collect(toMultimap(S3Location::bucket, S3Location::key, HashMultimap::create)); @@ -132,8 +164,10 @@ public void deleteFiles(Collection locations) .toList(); DeleteObjectsRequest request = DeleteObjectsRequest.builder() + .overrideConfiguration(context::applyCredentialProviderOverride) .requestPayer(requestPayer) .bucket(bucket) + .overrideConfiguration(disableStrongIntegrityChecksums()) .delete(builder -> builder.objects(objects).quiet(true)) .build(); @@ -144,7 +178,7 @@ public void deleteFiles(Collection locations) } } catch (SdkException e) { - throw new IOException("Error while batch deleting files", e); + throw new TrinoFileSystemException("Error while batch deleting files", e); } } } @@ -164,6 +198,12 @@ public void renameFile(Location source, Location target) @Override public FileIterator listFiles(Location location) throws IOException + { + return listObjects(location, false); + } + + private FileIterator listObjects(Location location, boolean includeDirectoryObjects) + throws IOException { S3Location s3Location = new S3Location(location); @@ -173,16 +213,23 @@ public FileIterator listFiles(Location location) } ListObjectsV2Request request = ListObjectsV2Request.builder() + .overrideConfiguration(context::applyCredentialProviderOverride) + // Restore status will not be added to the response without requested + //.optionalObjectAttributes(OptionalObjectAttributes.RESTORE_STATUS) + .requestPayer(requestPayer) .bucket(s3Location.bucket()) .prefix(key) .build(); try { - ListObjectsV2Iterable iterable = client.listObjectsV2Paginator(request); - return new S3FileIterator(s3Location, iterable.contents().iterator()); + Stream s3ObjectStream = client.listObjectsV2Paginator(request).contents().stream(); + if (!includeDirectoryObjects) { + s3ObjectStream = s3ObjectStream.filter(object -> !object.key().endsWith("/")); + } + return new S3FileIterator(s3Location, s3ObjectStream.iterator()); } catch (SdkException e) { - throw new IOException("Failed to list location: " + location, e); + throw new TrinoFileSystemException("Failed to list location: " + location, e); } } @@ -197,9 +244,119 @@ public Optional directoryExists(Location location) return Optional.empty(); } + @Override + public void createDirectory(Location location) + { + validateS3Location(location); + // S3 does not have directories + } + + @Override + public void renameDirectory(Location source, Location target) + throws IOException + { + throw new IOException("S3 does not support directory renames"); + } + + @Override + public Set listDirectories(Location location) + throws IOException + { + S3Location s3Location = new S3Location(location); + Location baseLocation = s3Location.baseLocation(); + + String key = s3Location.key(); + if (!key.isEmpty() && !key.endsWith("/")) { + key += "/"; + } + + ListObjectsV2Request request = ListObjectsV2Request.builder() + .overrideConfiguration(context::applyCredentialProviderOverride) + .requestPayer(requestPayer) + .bucket(s3Location.bucket()) + .prefix(key) + .delimiter("/") + .build(); + + try { + return client.listObjectsV2Paginator(request) + .commonPrefixes().stream() + .map(CommonPrefix::prefix) + .map(baseLocation::appendPath) + .collect(toImmutableSet()); + } + catch (SdkException e) { + throw new TrinoFileSystemException("Failed to list location: " + location, e); + } + } + + @Override + public Optional createTemporaryDirectory(Location targetPath, String temporaryPrefix, String relativePrefix) + throws IOException + { + validateS3Location(targetPath); + // S3 does not have directories + return Optional.empty(); + } + + public Optional encryptedPreSignedUri(Location location, Duration ttl, Optional key) + throws IOException + { + location.verifyValidFileLocation(); + S3Location s3Location = new S3Location(location); + + verify(key.isEmpty() || context.s3SseContext().sseType() == NONE, "Encryption key cannot be used with SSE configuration"); + + GetObjectRequest request = GetObjectRequest.builder() + .overrideConfiguration(context::applyCredentialProviderOverride) + .requestPayer(requestPayer) + .key(s3Location.key()) + .bucket(s3Location.bucket()) + .applyMutation(builder -> + key.ifPresentOrElse( + encryption -> + builder.sseCustomerKeyMD5(md5Checksum(encryption)) + .sseCustomerAlgorithm(encryption.algorithm()) + .sseCustomerKey(encoded(encryption)), + () -> setEncryptionSettings(builder, context.s3SseContext()))) + .build(); + + GetObjectPresignRequest preSignRequest = GetObjectPresignRequest.builder() + .signatureDuration(ttl.toJavaTime()) + .getObjectRequest(request) + .build(); + try { + PresignedGetObjectRequest preSigned = preSigner.presignGetObject(preSignRequest); + return Optional.of(new UriLocation(preSigned.url().toURI(), filterHeaders(preSigned.httpRequest().headers()))); + } + catch (SdkException e) { + throw new IOException("Failed to generate pre-signed URI", e); + } + catch (URISyntaxException e) { + throw new TrinoFileSystemException("Failed to convert pre-signed URI to URI", e); + } + } + + private static Map> filterHeaders(Map> headers) + { + return headers.entrySet().stream() + .filter(entry -> !entry.getKey().equalsIgnoreCase("host")) + .collect(toMap(Entry::getKey, Entry::getValue)); + } + @SuppressWarnings("ResultOfObjectAllocationIgnored") private static void validateS3Location(Location location) { new S3Location(location); } + + // TODO (https://github.com/trinodb/trino/issues/24955): + // remove me once all of the S3-compatible storage support strong integrity checks + @SuppressWarnings("deprecation") + static AwsRequestOverrideConfiguration disableStrongIntegrityChecksums() + { + return AwsRequestOverrideConfiguration.builder() + .signer(AwsS3V4Signer.create()) + .build(); + } } diff --git a/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3FileSystemConfig.java b/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3FileSystemConfig.java index 616e79e7c54b..2940546bd0aa 100644 --- a/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3FileSystemConfig.java +++ b/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3FileSystemConfig.java @@ -13,23 +13,94 @@ */ package io.trino.filesystem.s3; +import com.google.common.base.Splitter; +import com.google.common.collect.ImmutableSet; import com.google.common.net.HostAndPort; import io.airlift.configuration.Config; import io.airlift.configuration.ConfigDescription; import io.airlift.configuration.ConfigSecuritySensitive; import io.airlift.units.DataSize; +import io.airlift.units.Duration; import io.airlift.units.MaxDataSize; import io.airlift.units.MinDataSize; +import jakarta.validation.constraints.AssertTrue; import jakarta.validation.constraints.Min; import jakarta.validation.constraints.NotNull; +import jakarta.validation.constraints.Size; +import software.amazon.awssdk.retries.api.RetryStrategy; +import software.amazon.awssdk.services.s3.model.ObjectCannedACL; +import software.amazon.awssdk.services.s3.model.StorageClass; +import java.util.Optional; +import java.util.Set; + +import static com.google.common.base.Strings.nullToEmpty; import static io.airlift.units.DataSize.Unit.MEGABYTE; +import static software.amazon.awssdk.awscore.retry.AwsRetryStrategy.adaptiveRetryStrategy; +import static software.amazon.awssdk.awscore.retry.AwsRetryStrategy.legacyRetryStrategy; +import static software.amazon.awssdk.awscore.retry.AwsRetryStrategy.standardRetryStrategy; public class S3FileSystemConfig { public enum S3SseType { - NONE, S3, KMS + NONE, S3, KMS, CUSTOMER + } + + public enum StorageClassType + { + STANDARD, + STANDARD_IA, + INTELLIGENT_TIERING; + + public static StorageClass toStorageClass(StorageClassType storageClass) + { + return switch (storageClass) { + case STANDARD -> StorageClass.STANDARD; + case STANDARD_IA -> StorageClass.STANDARD_IA; + case INTELLIGENT_TIERING -> StorageClass.INTELLIGENT_TIERING; + }; + } + } + + public enum ObjectCannedAcl + { + NONE, + PRIVATE, + PUBLIC_READ, + PUBLIC_READ_WRITE, + AUTHENTICATED_READ, + BUCKET_OWNER_READ, + BUCKET_OWNER_FULL_CONTROL; + + public static ObjectCannedACL getCannedAcl(S3FileSystemConfig.ObjectCannedAcl cannedAcl) + { + return switch (cannedAcl) { + case NONE -> null; + case PRIVATE -> ObjectCannedACL.PRIVATE; + case PUBLIC_READ -> ObjectCannedACL.PUBLIC_READ; + case PUBLIC_READ_WRITE -> ObjectCannedACL.PUBLIC_READ_WRITE; + case AUTHENTICATED_READ -> ObjectCannedACL.AUTHENTICATED_READ; + case BUCKET_OWNER_READ -> ObjectCannedACL.BUCKET_OWNER_READ; + case BUCKET_OWNER_FULL_CONTROL -> ObjectCannedACL.BUCKET_OWNER_FULL_CONTROL; + }; + } + } + + public enum RetryMode + { + STANDARD, + LEGACY, + ADAPTIVE; + + public static RetryStrategy getRetryStrategy(RetryMode retryMode) + { + return switch (retryMode) { + case STANDARD -> standardRetryStrategy(); + case LEGACY -> legacyRetryStrategy(); + case ADAPTIVE -> adaptiveRetryStrategy(); + }; + } } private String awsAccessKey; @@ -37,6 +108,7 @@ public enum S3SseType private String endpoint; private String region; private boolean pathStyleAccess; + private StorageClassType storageClass = StorageClassType.STANDARD; private String iamRole; private String roleSessionName = "trino-filesystem"; private String externalId; @@ -44,11 +116,27 @@ public enum S3SseType private String stsRegion; private S3SseType sseType = S3SseType.NONE; private String sseKmsKeyId; + private String sseCustomerKey; + private boolean useWebIdentityTokenCredentialsProvider; private DataSize streamingPartSize = DataSize.of(16, MEGABYTE); private boolean requesterPays; - private Integer maxConnections; + private Integer maxConnections = 500; + private Duration connectionTtl; + private Duration connectionMaxIdleTime; + private Duration socketConnectTimeout; + private Duration socketReadTimeout; + private boolean tcpKeepAlive; private HostAndPort httpProxy; private boolean httpProxySecure; + private String httpProxyUsername; + private String httpProxyPassword; + private boolean preemptiveBasicProxyAuth; + private Set nonProxyHosts = ImmutableSet.of(); + private ObjectCannedAcl objectCannedAcl = ObjectCannedAcl.NONE; + private RetryMode retryMode = RetryMode.LEGACY; + private int maxErrorRetries = 10; + private boolean supportsExclusiveCreate = true; + private String applicationId = "Trino"; public String getAwsAccessKey() { @@ -112,6 +200,19 @@ public S3FileSystemConfig setPathStyleAccess(boolean pathStyleAccess) return this; } + public StorageClassType getStorageClass() + { + return storageClass; + } + + @Config("s3.storage-class") + @ConfigDescription("The S3 storage class to use when writing the data") + public S3FileSystemConfig setStorageClass(StorageClassType storageClass) + { + this.storageClass = storageClass; + return this; + } + public String getIamRole() { return iamRole; @@ -176,6 +277,46 @@ public S3FileSystemConfig setStsRegion(String stsRegion) return this; } + @NotNull + public ObjectCannedAcl getCannedAcl() + { + return objectCannedAcl; + } + + @Config("s3.canned-acl") + @ConfigDescription("Canned ACL (predefined grants) to manage access to objects") + public S3FileSystemConfig setCannedAcl(ObjectCannedAcl objectCannedAcl) + { + this.objectCannedAcl = objectCannedAcl; + return this; + } + + public RetryMode getRetryMode() + { + return retryMode; + } + + @Config("s3.retry-mode") + @ConfigDescription("Specifies how the AWS SDK attempts retries, default is LEGACY") + public S3FileSystemConfig setRetryMode(RetryMode retryMode) + { + this.retryMode = retryMode; + return this; + } + + @Min(1) // minimum set to 1 as the SDK validates this has to be > 0 + public int getMaxErrorRetries() + { + return maxErrorRetries; + } + + @Config("s3.max-error-retries") + public S3FileSystemConfig setMaxErrorRetries(int maxErrorRetries) + { + this.maxErrorRetries = maxErrorRetries; + return this; + } + @NotNull public S3SseType getSseType() { @@ -202,6 +343,41 @@ public S3FileSystemConfig setSseKmsKeyId(String sseKmsKeyId) return this; } + public boolean isUseWebIdentityTokenCredentialsProvider() + { + return useWebIdentityTokenCredentialsProvider; + } + + @Config("s3.use-web-identity-token-credentials-provider") + public S3FileSystemConfig setUseWebIdentityTokenCredentialsProvider(boolean useWebIdentityTokenCredentialsProvider) + { + this.useWebIdentityTokenCredentialsProvider = useWebIdentityTokenCredentialsProvider; + return this; + } + + public String getSseCustomerKey() + { + return sseCustomerKey; + } + + @Config("s3.sse.customer-key") + @ConfigDescription("Customer Key to use for S3 server-side encryption with Customer key (SSE-C)") + @ConfigSecuritySensitive + public S3FileSystemConfig setSseCustomerKey(String sseCustomerKey) + { + this.sseCustomerKey = sseCustomerKey; + return this; + } + + @AssertTrue(message = "s3.sse.customer-key has to be set for server-side encryption with customer-provided key") + public boolean isSseWithCustomerKeyConfigValid() + { + if (sseType == S3SseType.CUSTOMER) { + return sseCustomerKey != null; + } + return true; + } + @NotNull @MinDataSize("5MB") @MaxDataSize("256MB") @@ -243,6 +419,71 @@ public S3FileSystemConfig setMaxConnections(Integer maxConnections) return this; } + public Optional getConnectionTtl() + { + return Optional.ofNullable(connectionTtl); + } + + @Config("s3.connection-ttl") + @ConfigDescription("Maximum time allowed for connections to be reused before being replaced in the connection pool") + public S3FileSystemConfig setConnectionTtl(Duration connectionTtl) + { + this.connectionTtl = connectionTtl; + return this; + } + + public Optional getConnectionMaxIdleTime() + { + return Optional.ofNullable(connectionMaxIdleTime); + } + + @Config("s3.connection-max-idle-time") + @ConfigDescription("Maximum time allowed for connections to remain idle in the connection pool before being closed") + public S3FileSystemConfig setConnectionMaxIdleTime(Duration connectionMaxIdleTime) + { + this.connectionMaxIdleTime = connectionMaxIdleTime; + return this; + } + + public Optional getSocketConnectTimeout() + { + return Optional.ofNullable(socketConnectTimeout); + } + + @Config("s3.socket-connect-timeout") + @ConfigDescription("Maximum time allowed for socket connect to complete before timing out") + public S3FileSystemConfig setSocketConnectTimeout(Duration socketConnectTimeout) + { + this.socketConnectTimeout = socketConnectTimeout; + return this; + } + + public Optional getSocketReadTimeout() + { + return Optional.ofNullable(socketReadTimeout); + } + + @Config("s3.socket-read-timeout") + @ConfigDescription("Maximum time allowed for socket reads before timing out") + public S3FileSystemConfig setSocketReadTimeout(Duration socketReadTimeout) + { + this.socketReadTimeout = socketReadTimeout; + return this; + } + + public boolean getTcpKeepAlive() + { + return tcpKeepAlive; + } + + @Config("s3.tcp-keep-alive") + @ConfigDescription("Enable TCP keep alive on created connections") + public S3FileSystemConfig setTcpKeepAlive(boolean tcpKeepAlive) + { + this.tcpKeepAlive = tcpKeepAlive; + return this; + } + public HostAndPort getHttpProxy() { return httpProxy; @@ -266,4 +507,81 @@ public S3FileSystemConfig setHttpProxySecure(boolean httpProxySecure) this.httpProxySecure = httpProxySecure; return this; } + + public String getHttpProxyUsername() + { + return httpProxyUsername; + } + + @Config("s3.http-proxy.username") + public S3FileSystemConfig setHttpProxyUsername(String httpProxyUsername) + { + this.httpProxyUsername = httpProxyUsername; + return this; + } + + public String getHttpProxyPassword() + { + return httpProxyPassword; + } + + @Config("s3.http-proxy.password") + @ConfigSecuritySensitive + public S3FileSystemConfig setHttpProxyPassword(String httpProxyPassword) + { + this.httpProxyPassword = httpProxyPassword; + return this; + } + + public boolean getHttpProxyPreemptiveBasicProxyAuth() + { + return preemptiveBasicProxyAuth; + } + + @Config("s3.http-proxy.preemptive-basic-auth") + public S3FileSystemConfig setHttpProxyPreemptiveBasicProxyAuth(boolean preemptiveBasicProxyAuth) + { + this.preemptiveBasicProxyAuth = preemptiveBasicProxyAuth; + return this; + } + + public Set getNonProxyHosts() + { + return nonProxyHosts; + } + + @Config("s3.http-proxy.non-proxy-hosts") + public S3FileSystemConfig setNonProxyHosts(String nonProxyHosts) + { + this.nonProxyHosts = ImmutableSet.copyOf(Splitter.on(',').omitEmptyStrings().trimResults().split(nullToEmpty(nonProxyHosts))); + return this; + } + + public boolean isSupportsExclusiveCreate() + { + return supportsExclusiveCreate; + } + + @Config("s3.exclusive-create") + @ConfigDescription("Whether S3-compatible storage supports exclusive create (true for Minio and AWS S3)") + public S3FileSystemConfig setSupportsExclusiveCreate(boolean supportsExclusiveCreate) + { + this.supportsExclusiveCreate = supportsExclusiveCreate; + return this; + } + + @Size(max = 50) + @NotNull + public String getApplicationId() + { + return applicationId; + } + + @Config("s3.application-id") + @ConfigDescription("Suffix that will be added to HTTP User-Agent header to identify the application") + public S3FileSystemConfig setApplicationId(String applicationId) + { + this.applicationId = applicationId; + return this; + } } diff --git a/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3FileSystemFactory.java b/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3FileSystemFactory.java index 065ede41c99f..c165e41addf1 100644 --- a/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3FileSystemFactory.java +++ b/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3FileSystemFactory.java @@ -14,95 +14,46 @@ package io.trino.filesystem.s3; import com.google.inject.Inject; +import io.opentelemetry.api.OpenTelemetry; import io.trino.filesystem.TrinoFileSystem; import io.trino.filesystem.TrinoFileSystemFactory; import io.trino.spi.security.ConnectorIdentity; import jakarta.annotation.PreDestroy; -import software.amazon.awssdk.auth.credentials.AwsBasicCredentials; -import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider; -import software.amazon.awssdk.http.apache.ApacheHttpClient; -import software.amazon.awssdk.http.apache.ProxyConfiguration; -import software.amazon.awssdk.regions.Region; import software.amazon.awssdk.services.s3.S3Client; -import software.amazon.awssdk.services.s3.S3ClientBuilder; -import software.amazon.awssdk.services.sts.StsClient; -import software.amazon.awssdk.services.sts.StsClientBuilder; -import software.amazon.awssdk.services.sts.auth.StsAssumeRoleCredentialsProvider; +import software.amazon.awssdk.services.s3.presigner.S3Presigner; -import java.net.URI; -import java.util.Optional; - -import static java.lang.Math.toIntExact; +import java.util.concurrent.Executor; public final class S3FileSystemFactory implements TrinoFileSystemFactory { + private final S3FileSystemLoader loader; private final S3Client client; private final S3Context context; + private final Executor uploadExecutor; + private final S3Presigner preSigner; @Inject - public S3FileSystemFactory(S3FileSystemConfig config) + public S3FileSystemFactory(OpenTelemetry openTelemetry, S3FileSystemConfig config, S3FileSystemStats stats) { - S3ClientBuilder s3 = S3Client.builder(); - - if ((config.getAwsAccessKey() != null) && (config.getAwsSecretKey() != null)) { - s3.credentialsProvider(StaticCredentialsProvider.create( - AwsBasicCredentials.create(config.getAwsAccessKey(), config.getAwsSecretKey()))); - } - - Optional.ofNullable(config.getRegion()).map(Region::of).ifPresent(s3::region); - Optional.ofNullable(config.getEndpoint()).map(URI::create).ifPresent(s3::endpointOverride); - s3.forcePathStyle(config.isPathStyleAccess()); - - if (config.getIamRole() != null) { - StsClientBuilder sts = StsClient.builder(); - Optional.ofNullable(config.getStsEndpoint()).map(URI::create).ifPresent(sts::endpointOverride); - Optional.ofNullable(config.getStsRegion()) - .or(() -> Optional.ofNullable(config.getRegion())) - .map(Region::of).ifPresent(sts::region); - - s3.credentialsProvider(StsAssumeRoleCredentialsProvider.builder() - .refreshRequest(request -> request - .roleArn(config.getIamRole()) - .roleSessionName(config.getRoleSessionName()) - .externalId(config.getExternalId())) - .stsClient(sts.build()) - .asyncCredentialUpdateEnabled(true) - .build()); - } - - ApacheHttpClient.Builder httpClient = ApacheHttpClient.builder() - .maxConnections(config.getMaxConnections()); - - if (config.getHttpProxy() != null) { - URI endpoint = URI.create("%s://%s".formatted( - config.isHttpProxySecure() ? "https" : "http", - config.getHttpProxy())); - httpClient.proxyConfiguration(ProxyConfiguration.builder() - .endpoint(endpoint) - .build()); - } - - s3.httpClientBuilder(httpClient); - - this.client = s3.build(); - - context = new S3Context( - toIntExact(config.getStreamingPartSize().toBytes()), - config.isRequesterPays(), - config.getSseType(), - config.getSseKmsKeyId()); + this.loader = new S3FileSystemLoader(openTelemetry, config, stats); + this.client = loader.createClient(); + this.preSigner = loader.createPreSigner(); + this.context = loader.context(); + this.uploadExecutor = loader.uploadExecutor(); } @PreDestroy public void destroy() { - client.close(); + try (client) { + loader.destroy(); + } } @Override public TrinoFileSystem create(ConnectorIdentity identity) { - return new S3FileSystem(client, context); + return new S3FileSystem(uploadExecutor, client, preSigner, context.withCredentials(identity)); } } diff --git a/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3FileSystemLoader.java b/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3FileSystemLoader.java new file mode 100644 index 000000000000..3c3188049435 --- /dev/null +++ b/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3FileSystemLoader.java @@ -0,0 +1,320 @@ +/* + * 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.filesystem.s3; + +import com.google.inject.Inject; +import io.opentelemetry.api.OpenTelemetry; +import io.opentelemetry.instrumentation.awssdk.v2_2.AwsSdkTelemetry; +import io.trino.filesystem.Location; +import io.trino.filesystem.TrinoFileSystemFactory; +import io.trino.filesystem.s3.S3Context.S3SseContext; +import jakarta.annotation.PreDestroy; +import software.amazon.awssdk.auth.credentials.AwsBasicCredentials; +import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider; +import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider; +import software.amazon.awssdk.auth.credentials.WebIdentityTokenFileCredentialsProvider; +import software.amazon.awssdk.core.checksums.RequestChecksumCalculation; +import software.amazon.awssdk.core.client.config.ClientOverrideConfiguration; +import software.amazon.awssdk.http.SdkHttpClient; +import software.amazon.awssdk.http.apache.ApacheHttpClient; +import software.amazon.awssdk.http.apache.ProxyConfiguration; +import software.amazon.awssdk.metrics.MetricPublisher; +import software.amazon.awssdk.regions.Region; +import software.amazon.awssdk.services.s3.S3Client; +import software.amazon.awssdk.services.s3.S3ClientBuilder; +import software.amazon.awssdk.services.s3.S3Configuration; +import software.amazon.awssdk.services.s3.presigner.S3Presigner; +import software.amazon.awssdk.services.sts.StsClient; +import software.amazon.awssdk.services.sts.StsClientBuilder; +import software.amazon.awssdk.services.sts.auth.StsAssumeRoleCredentialsProvider; + +import java.net.URI; +import java.util.Map; +import java.util.Optional; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.Executor; +import java.util.concurrent.ExecutorService; +import java.util.function.Function; + +import static com.google.common.base.Preconditions.checkState; +import static io.airlift.concurrent.Threads.daemonThreadsNamed; +import static io.trino.filesystem.s3.S3FileSystemConfig.RetryMode.getRetryStrategy; +import static java.lang.Math.toIntExact; +import static java.util.Objects.requireNonNull; +import static java.util.concurrent.Executors.newCachedThreadPool; +import static software.amazon.awssdk.core.checksums.ResponseChecksumValidation.WHEN_REQUIRED; + +final class S3FileSystemLoader + implements Function +{ + private final Optional mappingProvider; + private final SdkHttpClient httpClient; + private final S3ClientFactory clientFactory; + private final S3Presigner preSigner; + private final S3Context context; + private final ExecutorService uploadExecutor = newCachedThreadPool(daemonThreadsNamed("s3-upload-%s")); + private final Map, S3Client> clients = new ConcurrentHashMap<>(); + + @Inject + public S3FileSystemLoader(S3SecurityMappingProvider mappingProvider, OpenTelemetry openTelemetry, S3FileSystemConfig config, S3FileSystemStats stats) + { + this(Optional.of(mappingProvider), openTelemetry, config, stats); + } + + S3FileSystemLoader(OpenTelemetry openTelemetry, S3FileSystemConfig config, S3FileSystemStats stats) + { + this(Optional.empty(), openTelemetry, config, stats); + } + + private S3FileSystemLoader(Optional mappingProvider, OpenTelemetry openTelemetry, S3FileSystemConfig config, S3FileSystemStats stats) + { + this.mappingProvider = requireNonNull(mappingProvider, "mappingProvider is null"); + this.httpClient = createHttpClient(config); + + requireNonNull(stats, "stats is null"); + + MetricPublisher metricPublisher = stats.newMetricPublisher(); + this.clientFactory = s3ClientFactory(httpClient, openTelemetry, config, metricPublisher); + + this.preSigner = s3PreSigner(httpClient, openTelemetry, config, metricPublisher); + + this.context = new S3Context( + toIntExact(config.getStreamingPartSize().toBytes()), + config.isRequesterPays(), + S3SseContext.of( + config.getSseType(), + config.getSseKmsKeyId(), + config.getSseCustomerKey()), + Optional.empty(), + config.getStorageClass(), + config.getCannedAcl(), + config.isSupportsExclusiveCreate()); + } + + @Override + public TrinoFileSystemFactory apply(Location location) + { + return identity -> { + Optional mapping = mappingProvider.orElseThrow().getMapping(identity, location); + + S3Client client = clients.computeIfAbsent(mapping, ignore -> clientFactory.create(mapping)); + S3Context context = this.context.withCredentials(identity); + + if (mapping.isPresent() && mapping.get().kmsKeyId().isPresent()) { + checkState(mapping.get().sseCustomerKey().isEmpty(), "Both SSE-C and KMS-managed keys cannot be used at the same time"); + context = context.withKmsKeyId(mapping.get().kmsKeyId().get()); + } + + if (mapping.isPresent() && mapping.get().sseCustomerKey().isPresent()) { + context = context.withSseCustomerKey(mapping.get().sseCustomerKey().get()); + } + + return new S3FileSystem(uploadExecutor, client, preSigner, context); + }; + } + + @PreDestroy + public void destroy() + { + try (httpClient) { + uploadExecutor.shutdownNow(); + } + } + + S3Client createClient() + { + return clientFactory.create(Optional.empty()); + } + + S3Presigner createPreSigner() + { + return preSigner; + } + + S3Context context() + { + return context; + } + + Executor uploadExecutor() + { + return uploadExecutor; + } + + private static S3ClientFactory s3ClientFactory(SdkHttpClient httpClient, OpenTelemetry openTelemetry, S3FileSystemConfig config, MetricPublisher metricPublisher) + { + ClientOverrideConfiguration overrideConfiguration = createOverrideConfiguration(openTelemetry, config, metricPublisher); + + Optional staticCredentialsProvider = createStaticCredentialsProvider(config); + Optional staticRegion = Optional.ofNullable(config.getRegion()); + Optional staticEndpoint = Optional.ofNullable(config.getEndpoint()); + boolean pathStyleAccess = config.isPathStyleAccess(); + boolean useWebIdentityTokenCredentialsProvider = config.isUseWebIdentityTokenCredentialsProvider(); + Optional staticIamRole = Optional.ofNullable(config.getIamRole()); + String staticRoleSessionName = config.getRoleSessionName(); + String externalId = config.getExternalId(); + + return mapping -> { + Optional credentialsProvider = mapping + .flatMap(S3SecurityMappingResult::credentialsProvider) + .or(() -> staticCredentialsProvider); + + Optional region = mapping.flatMap(S3SecurityMappingResult::region).or(() -> staticRegion); + Optional endpoint = mapping.flatMap(S3SecurityMappingResult::endpoint).or(() -> staticEndpoint); + + Optional iamRole = mapping.flatMap(S3SecurityMappingResult::iamRole).or(() -> staticIamRole); + String roleSessionName = mapping.flatMap(S3SecurityMappingResult::roleSessionName).orElse(staticRoleSessionName); + + S3ClientBuilder s3 = S3Client.builder(); + s3.overrideConfiguration(overrideConfiguration); + s3.httpClient(httpClient); + s3.responseChecksumValidation(WHEN_REQUIRED); + s3.requestChecksumCalculation(RequestChecksumCalculation.WHEN_REQUIRED); + + region.map(Region::of).ifPresent(s3::region); + endpoint.map(URI::create).ifPresent(s3::endpointOverride); + s3.forcePathStyle(pathStyleAccess); + + if (useWebIdentityTokenCredentialsProvider) { + s3.credentialsProvider(WebIdentityTokenFileCredentialsProvider.builder() + .asyncCredentialUpdateEnabled(true) + .build()); + } + else if (iamRole.isPresent()) { + s3.credentialsProvider(StsAssumeRoleCredentialsProvider.builder() + .refreshRequest(request -> request + .roleArn(iamRole.get()) + .roleSessionName(roleSessionName) + .externalId(externalId)) + .stsClient(createStsClient(config, credentialsProvider)) + .asyncCredentialUpdateEnabled(true) + .build()); + } + else { + credentialsProvider.ifPresent(s3::credentialsProvider); + } + + return s3.build(); + }; + } + + private static S3Presigner s3PreSigner(SdkHttpClient httpClient, OpenTelemetry openTelemetry, S3FileSystemConfig config, MetricPublisher metricPublisher) + { + Optional staticCredentialsProvider = createStaticCredentialsProvider(config); + Optional staticRegion = Optional.ofNullable(config.getRegion()); + Optional staticEndpoint = Optional.ofNullable(config.getEndpoint()); + boolean pathStyleAccess = config.isPathStyleAccess(); + boolean useWebIdentityTokenCredentialsProvider = config.isUseWebIdentityTokenCredentialsProvider(); + Optional staticIamRole = Optional.ofNullable(config.getIamRole()); + String staticRoleSessionName = config.getRoleSessionName(); + String externalId = config.getExternalId(); + + S3Presigner.Builder s3 = S3Presigner.builder(); + s3.s3Client(s3ClientFactory(httpClient, openTelemetry, config, metricPublisher) + .create(Optional.empty())); + + staticRegion.map(Region::of).ifPresent(s3::region); + staticEndpoint.map(URI::create).ifPresent(s3::endpointOverride); + s3.serviceConfiguration(S3Configuration.builder() + .pathStyleAccessEnabled(pathStyleAccess) + .build()); + + if (useWebIdentityTokenCredentialsProvider) { + s3.credentialsProvider(WebIdentityTokenFileCredentialsProvider.builder() + .asyncCredentialUpdateEnabled(true) + .build()); + } + else if (staticIamRole.isPresent()) { + s3.credentialsProvider(StsAssumeRoleCredentialsProvider.builder() + .refreshRequest(request -> request + .roleArn(staticIamRole.get()) + .roleSessionName(staticRoleSessionName) + .externalId(externalId)) + .stsClient(createStsClient(config, staticCredentialsProvider)) + .asyncCredentialUpdateEnabled(true) + .build()); + } + else { + staticCredentialsProvider.ifPresent(s3::credentialsProvider); + } + + return s3.build(); + } + + private static Optional createStaticCredentialsProvider(S3FileSystemConfig config) + { + if ((config.getAwsAccessKey() != null) || (config.getAwsSecretKey() != null)) { + return Optional.of(StaticCredentialsProvider.create( + AwsBasicCredentials.create(config.getAwsAccessKey(), config.getAwsSecretKey()))); + } + return Optional.empty(); + } + + private static StsClient createStsClient(S3FileSystemConfig config, Optional credentialsProvider) + { + StsClientBuilder sts = StsClient.builder(); + Optional.ofNullable(config.getStsEndpoint()).map(URI::create).ifPresent(sts::endpointOverride); + Optional.ofNullable(config.getStsRegion()) + .or(() -> Optional.ofNullable(config.getRegion())) + .map(Region::of).ifPresent(sts::region); + credentialsProvider.ifPresent(sts::credentialsProvider); + return sts.build(); + } + + private static ClientOverrideConfiguration createOverrideConfiguration(OpenTelemetry openTelemetry, S3FileSystemConfig config, MetricPublisher metricPublisher) + { + return ClientOverrideConfiguration.builder() + .addExecutionInterceptor(AwsSdkTelemetry.builder(openTelemetry) + .setCaptureExperimentalSpanAttributes(true) + //.setRecordIndividualHttpError(true) + .build().newExecutionInterceptor()) + .retryStrategy(getRetryStrategy(config.getRetryMode()).toBuilder() + .maxAttempts(config.getMaxErrorRetries()) + .build()) + .appId(config.getApplicationId()) + .addMetricPublisher(metricPublisher) + .build(); + } + + private static SdkHttpClient createHttpClient(S3FileSystemConfig config) + { + ApacheHttpClient.Builder client = ApacheHttpClient.builder() + .maxConnections(config.getMaxConnections()) + .tcpKeepAlive(config.getTcpKeepAlive()); + + config.getConnectionTtl().ifPresent(ttl -> client.connectionTimeToLive(ttl.toJavaTime())); + config.getConnectionMaxIdleTime().ifPresent(time -> client.connectionMaxIdleTime(time.toJavaTime())); + config.getSocketConnectTimeout().ifPresent(timeout -> client.connectionTimeout(timeout.toJavaTime())); + config.getSocketReadTimeout().ifPresent(timeout -> client.socketTimeout(timeout.toJavaTime())); + + if (config.getHttpProxy() != null) { + client.proxyConfiguration(ProxyConfiguration.builder() + .endpoint(URI.create("%s://%s".formatted( + config.isHttpProxySecure() ? "https" : "http", + config.getHttpProxy()))) + .username(config.getHttpProxyUsername()) + .password(config.getHttpProxyPassword()) + .nonProxyHosts(config.getNonProxyHosts()) + .preemptiveBasicAuthenticationEnabled(config.getHttpProxyPreemptiveBasicProxyAuth()) + .build()); + } + + return client.build(); + } + + interface S3ClientFactory + { + S3Client create(Optional mapping); + } +} diff --git a/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3FileSystemModule.java b/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3FileSystemModule.java index ac96377cfffb..a7db2f289634 100644 --- a/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3FileSystemModule.java +++ b/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3FileSystemModule.java @@ -13,19 +13,90 @@ */ package io.trino.filesystem.s3; +import com.google.common.base.VerifyException; import com.google.inject.Binder; -import com.google.inject.Module; +import com.google.inject.BindingAnnotation; +import com.google.inject.Key; +import com.google.inject.Provides; +import com.google.inject.Singleton; +import io.airlift.configuration.AbstractConfigurationAwareModule; +import io.airlift.units.Duration; +import io.trino.filesystem.TrinoFileSystemFactory; +import io.trino.filesystem.switching.SwitchingFileSystemFactory; + +import java.lang.annotation.Retention; +import java.lang.annotation.Target; +import java.util.function.Supplier; import static com.google.inject.Scopes.SINGLETON; import static io.airlift.configuration.ConfigBinder.configBinder; +import static io.airlift.http.client.HttpClientBinder.httpClientBinder; +import static java.lang.annotation.ElementType.FIELD; +import static java.lang.annotation.ElementType.METHOD; +import static java.lang.annotation.ElementType.PARAMETER; +import static java.lang.annotation.RetentionPolicy.RUNTIME; +import static java.util.concurrent.TimeUnit.SECONDS; +import static org.weakref.jmx.guice.ExportBinder.newExporter; public class S3FileSystemModule - implements Module + extends AbstractConfigurationAwareModule { @Override - public void configure(Binder binder) + protected void setup(Binder binder) { configBinder(binder).bindConfig(S3FileSystemConfig.class); - binder.bind(S3FileSystemFactory.class).in(SINGLETON); + + if (buildConfigObject(S3SecurityMappingEnabledConfig.class).isEnabled()) { + install(new S3SecurityMappingModule()); + } + else { + binder.bind(TrinoFileSystemFactory.class).annotatedWith(FileSystemS3.class) + .to(S3FileSystemFactory.class).in(SINGLETON); + } + + binder.bind(S3FileSystemStats.class).in(SINGLETON); + newExporter(binder).export(S3FileSystemStats.class).withGeneratedName(); + } + + public static class S3SecurityMappingModule + extends AbstractConfigurationAwareModule + { + @Override + protected void setup(Binder binder) + { + S3SecurityMappingConfig config = buildConfigObject(S3SecurityMappingConfig.class); + + binder.bind(S3SecurityMappingProvider.class).in(SINGLETON); + binder.bind(S3FileSystemLoader.class).in(SINGLETON); + + var mappingsBinder = binder.bind(new Key>() {}); + if (config.getConfigFile().isPresent()) { + mappingsBinder.to(S3SecurityMappingsFileSource.class).in(SINGLETON); + } + else if (config.getConfigUri().isPresent()) { + mappingsBinder.to(S3SecurityMappingsUriSource.class).in(SINGLETON); + httpClientBinder(binder).bindHttpClient("s3-security-mapping", ForS3SecurityMapping.class) + .withConfigDefaults(httpConfig -> httpConfig + .setRequestTimeout(new Duration(10, SECONDS)) + .setSelectorCount(1) + .setMinThreads(1)); + } + else { + throw new VerifyException("No security mapping source configured"); + } + } + + @Provides + @Singleton + @FileSystemS3 + static TrinoFileSystemFactory createFileSystemFactory(S3FileSystemLoader loader) + { + return new SwitchingFileSystemFactory(loader); + } } + + @Retention(RUNTIME) + @Target({FIELD, PARAMETER, METHOD}) + @BindingAnnotation + public @interface ForS3SecurityMapping {} } diff --git a/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3FileSystemStats.java b/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3FileSystemStats.java new file mode 100644 index 000000000000..00cca1077bf0 --- /dev/null +++ b/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3FileSystemStats.java @@ -0,0 +1,284 @@ +/* + * 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.filesystem.s3; + +import io.airlift.log.Logger; +import org.weakref.jmx.Managed; +import org.weakref.jmx.Nested; +import software.amazon.awssdk.metrics.MetricCollection; +import software.amazon.awssdk.metrics.MetricPublisher; +import software.amazon.awssdk.metrics.SdkMetric; + +import java.time.Duration; +import java.util.Objects; +import java.util.Optional; +import java.util.Set; + +import static java.util.Objects.requireNonNull; +import static software.amazon.awssdk.core.internal.metrics.SdkErrorType.SERVER_ERROR; +import static software.amazon.awssdk.core.internal.metrics.SdkErrorType.THROTTLING; +import static software.amazon.awssdk.core.metrics.CoreMetric.API_CALL_DURATION; +import static software.amazon.awssdk.core.metrics.CoreMetric.API_CALL_SUCCESSFUL; +import static software.amazon.awssdk.core.metrics.CoreMetric.ERROR_TYPE; +import static software.amazon.awssdk.core.metrics.CoreMetric.OPERATION_NAME; +import static software.amazon.awssdk.core.metrics.CoreMetric.RETRY_COUNT; +import static software.amazon.awssdk.core.metrics.CoreMetric.SERVICE_ID; +import static software.amazon.awssdk.http.HttpMetric.AVAILABLE_CONCURRENCY; +import static software.amazon.awssdk.http.HttpMetric.CONCURRENCY_ACQUIRE_DURATION; +import static software.amazon.awssdk.http.HttpMetric.LEASED_CONCURRENCY; +import static software.amazon.awssdk.http.HttpMetric.PENDING_CONCURRENCY_ACQUIRES; + +public class S3FileSystemStats +{ + private final AwsSdkV2ApiCallStats total = new AwsSdkV2ApiCallStats(); + + private final AwsSdkV2ApiCallStats headObject = new AwsSdkV2ApiCallStats(); + private final AwsSdkV2ApiCallStats getObject = new AwsSdkV2ApiCallStats(); + private final AwsSdkV2ApiCallStats listObjectsV2 = new AwsSdkV2ApiCallStats(); + private final AwsSdkV2ApiCallStats putObject = new AwsSdkV2ApiCallStats(); + private final AwsSdkV2ApiCallStats deleteObject = new AwsSdkV2ApiCallStats(); + private final AwsSdkV2ApiCallStats deleteObjects = new AwsSdkV2ApiCallStats(); + private final AwsSdkV2ApiCallStats createMultipartUpload = new AwsSdkV2ApiCallStats(); + private final AwsSdkV2ApiCallStats completeMultipartUpload = new AwsSdkV2ApiCallStats(); + private final AwsSdkV2ApiCallStats abortMultipartUpload = new AwsSdkV2ApiCallStats(); + private final AwsSdkV2ApiCallStats uploadPart = new AwsSdkV2ApiCallStats(); + private final AwsSdkV2HttpClientStats httpClientStats = new AwsSdkV2HttpClientStats(); + + private static final AwsSdkV2ApiCallStats dummy = new DummyAwsSdkV2ApiCallStats(); + + @Managed + @Nested + public AwsSdkV2ApiCallStats getTotal() + { + return total; + } + + @Managed + @Nested + public AwsSdkV2ApiCallStats getHeadObject() + { + return headObject; + } + + @Managed + @Nested + public AwsSdkV2ApiCallStats getGetObject() + { + return getObject; + } + + @Managed + @Nested + public AwsSdkV2ApiCallStats getListObjectsV2() + { + return listObjectsV2; + } + + @Managed + @Nested + public AwsSdkV2ApiCallStats getPutObject() + { + return putObject; + } + + @Managed + @Nested + public AwsSdkV2ApiCallStats getDeleteObject() + { + return deleteObject; + } + + @Managed + @Nested + public AwsSdkV2ApiCallStats getDeleteObjects() + { + return deleteObjects; + } + + @Managed + @Nested + public AwsSdkV2ApiCallStats getCreateMultipartUpload() + { + return createMultipartUpload; + } + + @Managed + @Nested + public AwsSdkV2ApiCallStats getCompleteMultipartUpload() + { + return completeMultipartUpload; + } + + @Managed + @Nested + public AwsSdkV2ApiCallStats getAbortMultipartUpload() + { + return abortMultipartUpload; + } + + @Managed + @Nested + public AwsSdkV2ApiCallStats uploadPart() + { + return uploadPart; + } + + @Managed + @Nested + public AwsSdkV2HttpClientStats getHttpClientStats() + { + return httpClientStats; + } + + public MetricPublisher newMetricPublisher() + { + return new JmxMetricPublisher(this); + } + + public static final class JmxMetricPublisher + implements MetricPublisher + { + private static final Set> ALLOWED_METRICS = Set.of( + API_CALL_SUCCESSFUL, RETRY_COUNT, API_CALL_DURATION, ERROR_TYPE, + AVAILABLE_CONCURRENCY, LEASED_CONCURRENCY, PENDING_CONCURRENCY_ACQUIRES, + CONCURRENCY_ACQUIRE_DURATION); + + private static final Logger log = Logger.get(JmxMetricPublisher.class); + + private final S3FileSystemStats stats; + + public JmxMetricPublisher(S3FileSystemStats stats) + { + this.stats = requireNonNull(stats, "stats is null"); + } + + @Override + public void publish(MetricCollection metricCollection) + { + try { + Optional serviceId = metricCollection.metricValues(SERVICE_ID).stream().filter(Objects::nonNull).findFirst(); + Optional operationName = metricCollection.metricValues(OPERATION_NAME).stream().filter(Objects::nonNull).findFirst(); + if (serviceId.isEmpty() || operationName.isEmpty()) { + log.warn("ServiceId or OperationName is empty for AWS MetricCollection: %s", metricCollection); + return; + } + + if (!serviceId.get().equals("S3")) { + return; + } + + AwsSdkV2ApiCallStats apiCallStats = getApiCallStats(operationName.get()); + publishMetrics(metricCollection, apiCallStats, stats.httpClientStats); + } + catch (Exception e) { + log.warn(e, "Publishing AWS metrics failed"); + } + } + + private void publishMetrics(MetricCollection metricCollection, AwsSdkV2ApiCallStats apiCallStats, AwsSdkV2HttpClientStats httpClientStats) + { + metricCollection.stream() + .filter(metricRecord -> metricRecord.value() != null && ALLOWED_METRICS.contains(metricRecord.metric())) + .forEach(metricRecord -> { + if (metricRecord.metric().equals(API_CALL_SUCCESSFUL)) { + Boolean value = (Boolean) metricRecord.value(); + + stats.total.updateCalls(); + apiCallStats.updateCalls(); + + if (value.equals(Boolean.FALSE)) { + stats.total.updateFailures(); + apiCallStats.updateFailures(); + } + } + else if (metricRecord.metric().equals(RETRY_COUNT)) { + int value = (int) metricRecord.value(); + + stats.total.updateRetries(value); + apiCallStats.updateRetries(value); + } + else if (metricRecord.metric().equals(API_CALL_DURATION)) { + Duration value = (Duration) metricRecord.value(); + + stats.total.updateLatency(value); + apiCallStats.updateLatency(value); + } + else if (metricRecord.metric().equals(ERROR_TYPE)) { + String value = (String) metricRecord.value(); + + if (value.equals(THROTTLING.toString())) { + stats.total.updateThrottlingExceptions(); + apiCallStats.updateThrottlingExceptions(); + } + else if (value.equals(SERVER_ERROR.toString())) { + stats.total.updateServerErrors(); + apiCallStats.updateServerErrors(); + } + } + else if (metricRecord.metric().equals(LEASED_CONCURRENCY) || metricRecord.metric().equals(AVAILABLE_CONCURRENCY) || metricRecord.metric().equals(PENDING_CONCURRENCY_ACQUIRES)) { + int value = (int) metricRecord.value(); + httpClientStats.updateConcurrencyStats(metricRecord.metric(), value); + } + else if (metricRecord.metric().equals(CONCURRENCY_ACQUIRE_DURATION)) { + Duration duration = (Duration) metricRecord.value(); + httpClientStats.updateConcurrencyAcquireDuration(duration); + } + }); + + metricCollection.children().forEach(child -> publishMetrics(child, apiCallStats, httpClientStats)); + } + + @Override + public void close() {} + + private AwsSdkV2ApiCallStats getApiCallStats(String operationName) + { + return switch (operationName) { + case "HeadObject" -> stats.headObject; + case "GetObject" -> stats.getObject; + case "ListObjectsV2" -> stats.listObjectsV2; + case "PutObject" -> stats.putObject; + case "DeleteObject" -> stats.deleteObject; + case "DeleteObjects" -> stats.deleteObjects; + case "CreateMultipartUpload" -> stats.createMultipartUpload; + case "CompleteMultipartUpload" -> stats.completeMultipartUpload; + case "AbortMultipartUpload" -> stats.abortMultipartUpload; + case "UploadPart" -> stats.uploadPart; + default -> S3FileSystemStats.dummy; + }; + } + } + + private static class DummyAwsSdkV2ApiCallStats + extends AwsSdkV2ApiCallStats + { + @Override + public void updateLatency(Duration duration) {} + + @Override + public void updateCalls() {} + + @Override + public void updateFailures() {} + + @Override + public void updateRetries(int retryCount) {} + + @Override + public void updateThrottlingExceptions() {} + + @Override + public void updateServerErrors() {} + } +} diff --git a/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3InputFile.java b/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3InputFile.java index 9df5d5169854..f9a2f77f41e0 100644 --- a/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3InputFile.java +++ b/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3InputFile.java @@ -14,9 +14,11 @@ package io.trino.filesystem.s3; import io.trino.filesystem.Location; +import io.trino.filesystem.TrinoFileSystemException; import io.trino.filesystem.TrinoInput; import io.trino.filesystem.TrinoInputFile; import io.trino.filesystem.TrinoInputStream; +import io.trino.filesystem.encryption.EncryptionKey; import software.amazon.awssdk.core.exception.SdkException; import software.amazon.awssdk.services.s3.S3Client; import software.amazon.awssdk.services.s3.model.GetObjectRequest; @@ -28,7 +30,13 @@ import java.io.FileNotFoundException; import java.io.IOException; import java.time.Instant; +import java.util.Optional; +import static com.google.common.base.Verify.verify; +import static io.trino.filesystem.s3.S3FileSystemConfig.S3SseType.NONE; +import static io.trino.filesystem.s3.S3SseCUtils.encoded; +import static io.trino.filesystem.s3.S3SseCUtils.md5Checksum; +import static io.trino.filesystem.s3.S3SseRequestConfigurator.setEncryptionSettings; import static java.util.Objects.requireNonNull; final class S3InputFile @@ -36,17 +44,24 @@ final class S3InputFile { private final S3Client client; private final S3Location location; + private final S3Context context; private final RequestPayer requestPayer; + private final Optional key; private Long length; private Instant lastModified; - public S3InputFile(S3Client client, S3Context context, S3Location location, Long length) + public S3InputFile(S3Client client, S3Context context, S3Location location, Long length, Instant lastModified, Optional key) { this.client = requireNonNull(client, "client is null"); this.location = requireNonNull(location, "location is null"); + this.context = requireNonNull(context, "context is null"); this.requestPayer = context.requestPayer(); this.length = length; + this.lastModified = lastModified; + this.key = requireNonNull(key, "key is null"); location.location().verifyValidFileLocation(); + + verify(key.isEmpty() || context.s3SseContext().sseType() == NONE, "Encryption key cannot be used with SSE configuration"); } @Override @@ -97,9 +112,17 @@ public Location location() private GetObjectRequest newGetObjectRequest() { return GetObjectRequest.builder() + .overrideConfiguration(context::applyCredentialProviderOverride) .requestPayer(requestPayer) .bucket(location.bucket()) .key(location.key()) + .applyMutation(builder -> + key.ifPresentOrElse( + encryption -> + builder.sseCustomerKey(encoded(encryption)) + .sseCustomerAlgorithm(encryption.algorithm()) + .sseCustomerKeyMD5(md5Checksum(encryption)), + () -> setEncryptionSettings(builder, context.s3SseContext()))) .build(); } @@ -107,9 +130,17 @@ private boolean headObject() throws IOException { HeadObjectRequest request = HeadObjectRequest.builder() + .overrideConfiguration(context::applyCredentialProviderOverride) .requestPayer(requestPayer) .bucket(location.bucket()) .key(location.key()) + .applyMutation(builder -> + key.ifPresentOrElse( + encryption -> + builder.sseCustomerKey(encoded(encryption)) + .sseCustomerAlgorithm(encryption.algorithm()) + .sseCustomerKeyMD5(md5Checksum(encryption)), + () -> setEncryptionSettings(builder, context.s3SseContext()))) .build(); try { @@ -126,7 +157,7 @@ private boolean headObject() return false; } catch (SdkException e) { - throw new IOException("S3 HEAD request failed for file: " + location, e); + throw new TrinoFileSystemException("S3 HEAD request failed for file: " + location, e); } } } diff --git a/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3Location.java b/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3Location.java index 059453b16b03..d4e2932c731a 100644 --- a/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3Location.java +++ b/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3Location.java @@ -52,4 +52,9 @@ public String toString() { return location.toString(); } + + public Location baseLocation() + { + return Location.of("%s://%s/".formatted(scheme(), bucket())); + } } diff --git a/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3OutputFile.java b/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3OutputFile.java index 5a6bf934c1b8..8303d97bc393 100644 --- a/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3OutputFile.java +++ b/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3OutputFile.java @@ -15,25 +15,34 @@ import io.trino.filesystem.Location; import io.trino.filesystem.TrinoOutputFile; +import io.trino.filesystem.encryption.EncryptionKey; import io.trino.memory.context.AggregatedMemoryContext; import software.amazon.awssdk.services.s3.S3Client; +import java.io.IOException; import java.io.OutputStream; +import java.util.Optional; +import java.util.concurrent.Executor; +import static io.trino.memory.context.AggregatedMemoryContext.newSimpleAggregatedMemoryContext; import static java.util.Objects.requireNonNull; final class S3OutputFile implements TrinoOutputFile { + private final Executor uploadExecutor; private final S3Client client; private final S3Context context; private final S3Location location; + private final Optional key; - public S3OutputFile(S3Client client, S3Context context, S3Location location) + public S3OutputFile(Executor uploadExecutor, S3Client client, S3Context context, S3Location location, Optional key) { + this.uploadExecutor = requireNonNull(uploadExecutor, "uploadExecutor is null"); this.client = requireNonNull(client, "client is null"); this.context = requireNonNull(context, "context is null"); this.location = requireNonNull(location, "location is null"); + this.key = requireNonNull(key, "key is null"); location.location().verifyValidFileLocation(); } @@ -44,10 +53,37 @@ public OutputStream create(AggregatedMemoryContext memoryContext) return createOrOverwrite(memoryContext); } + @Override + public void createOrOverwrite(byte[] data) + throws IOException + { + try (OutputStream out = create(newSimpleAggregatedMemoryContext(), false)) { + out.write(data); + } + } + + @Override + public void createExclusive(byte[] data) + throws IOException + { + if (!context.exclusiveWriteSupported()) { + throw new UnsupportedOperationException("createExclusive not supported by " + getClass()); + } + + try (OutputStream out = create(newSimpleAggregatedMemoryContext(), true)) { + out.write(data); + } + } + @Override public OutputStream createOrOverwrite(AggregatedMemoryContext memoryContext) { - return new S3OutputStream(memoryContext, client, context, location); + return new S3OutputStream(memoryContext, uploadExecutor, client, context, location, true, key); + } + + public OutputStream create(AggregatedMemoryContext memoryContext, boolean exclusive) + { + return new S3OutputStream(memoryContext, uploadExecutor, client, context, location, exclusive, key); } @Override diff --git a/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3OutputStream.java b/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3OutputStream.java index 8259e2928900..e2210ee6cc50 100644 --- a/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3OutputStream.java +++ b/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3OutputStream.java @@ -13,7 +13,7 @@ */ package io.trino.filesystem.s3; -import io.trino.filesystem.s3.S3FileSystemConfig.S3SseType; +import io.trino.filesystem.encryption.EncryptionKey; import io.trino.memory.context.AggregatedMemoryContext; import io.trino.memory.context.LocalMemoryContext; import software.amazon.awssdk.core.exception.SdkException; @@ -23,8 +23,11 @@ import software.amazon.awssdk.services.s3.model.CompleteMultipartUploadRequest; import software.amazon.awssdk.services.s3.model.CompletedPart; import software.amazon.awssdk.services.s3.model.CreateMultipartUploadRequest; +import software.amazon.awssdk.services.s3.model.ObjectCannedACL; import software.amazon.awssdk.services.s3.model.PutObjectRequest; import software.amazon.awssdk.services.s3.model.RequestPayer; +import software.amazon.awssdk.services.s3.model.S3Exception; +import software.amazon.awssdk.services.s3.model.StorageClass; import software.amazon.awssdk.services.s3.model.UploadPartRequest; import software.amazon.awssdk.services.s3.model.UploadPartResponse; @@ -32,33 +35,44 @@ import java.io.InterruptedIOException; import java.io.OutputStream; import java.nio.ByteBuffer; +import java.nio.file.FileAlreadyExistsException; import java.util.ArrayList; import java.util.Arrays; import java.util.List; import java.util.Optional; import java.util.concurrent.ExecutionException; +import java.util.concurrent.Executor; import java.util.concurrent.Future; -import static com.google.common.primitives.Ints.constrainToRange; +import static com.google.common.base.Verify.verify; +import static io.trino.filesystem.s3.S3FileSystemConfig.ObjectCannedAcl.getCannedAcl; +import static io.trino.filesystem.s3.S3FileSystemConfig.S3SseType.NONE; +import static io.trino.filesystem.s3.S3FileSystemConfig.StorageClassType.toStorageClass; +import static io.trino.filesystem.s3.S3SseCUtils.encoded; +import static io.trino.filesystem.s3.S3SseCUtils.md5Checksum; +import static io.trino.filesystem.s3.S3SseRequestConfigurator.setEncryptionSettings; import static java.lang.Math.max; import static java.lang.Math.min; import static java.lang.System.arraycopy; +import static java.net.HttpURLConnection.HTTP_PRECON_FAILED; import static java.util.Objects.requireNonNull; import static java.util.concurrent.CompletableFuture.supplyAsync; -import static software.amazon.awssdk.services.s3.model.ServerSideEncryption.AES256; -import static software.amazon.awssdk.services.s3.model.ServerSideEncryption.AWS_KMS; final class S3OutputStream extends OutputStream { private final List parts = new ArrayList<>(); private final LocalMemoryContext memoryContext; + private final Executor uploadExecutor; private final S3Client client; private final S3Location location; + private final S3Context context; private final int partSize; private final RequestPayer requestPayer; - private final S3SseType sseType; - private final String sseKmsKeyId; + private final StorageClass storageClass; + private final ObjectCannedACL cannedAcl; + private final boolean exclusiveCreate; + private final Optional key; private int currentPartNumber; private byte[] buffer = new byte[0]; @@ -75,15 +89,21 @@ final class S3OutputStream // Visibility is ensured by calling get() on inProgressUploadFuture. private Optional uploadId = Optional.empty(); - public S3OutputStream(AggregatedMemoryContext memoryContext, S3Client client, S3Context context, S3Location location) + public S3OutputStream(AggregatedMemoryContext memoryContext, Executor uploadExecutor, S3Client client, S3Context context, S3Location location, boolean exclusiveCreate, Optional key) { this.memoryContext = memoryContext.newLocalMemoryContext(S3OutputStream.class.getSimpleName()); + this.uploadExecutor = requireNonNull(uploadExecutor, "uploadExecutor is null"); this.client = requireNonNull(client, "client is null"); this.location = requireNonNull(location, "location is null"); + this.exclusiveCreate = exclusiveCreate; + this.context = requireNonNull(context, "context is null"); this.partSize = context.partSize(); this.requestPayer = context.requestPayer(); - this.sseType = context.sseType(); - this.sseKmsKeyId = context.sseKmsKeyId(); + this.storageClass = toStorageClass(context.storageClass()); + this.cannedAcl = getCannedAcl(context.cannedAcl()); + this.key = requireNonNull(key, "key is null"); + + verify(key.isEmpty() || context.s3SseContext().sseType() == NONE, "Encryption key cannot be used with SSE configuration"); } @SuppressWarnings("NumericCastThatLosesPrecision") @@ -179,7 +199,7 @@ private void ensureCapacity(int extra) int target = max(buffer.length, initialBufferSize); if (target < capacity) { target += target / 2; // increase 50% - target = constrainToRange(target, capacity, partSize); + target = clamp(target, capacity, partSize); } buffer = Arrays.copyOf(buffer, target); memoryContext.setBytes(buffer.length); @@ -192,16 +212,23 @@ private void flushBuffer(boolean finished) // skip multipart upload if there would only be one part if (finished && !multipartUploadStarted) { PutObjectRequest request = PutObjectRequest.builder() + .overrideConfiguration(context::applyCredentialProviderOverride) + .acl(cannedAcl) .requestPayer(requestPayer) .bucket(location.bucket()) .key(location.key()) + .storageClass(storageClass) .contentLength((long) bufferSize) .applyMutation(builder -> { - switch (sseType) { - case NONE -> { /* ignored */ } - case S3 -> builder.serverSideEncryption(AES256); - case KMS -> builder.serverSideEncryption(AWS_KMS).ssekmsKeyId(sseKmsKeyId); - } + //if (exclusiveCreate) { + // builder.ifNoneMatch("*"); + //} + key.ifPresent(encryption -> { + builder.sseCustomerKey(encoded(encryption)); + builder.sseCustomerAlgorithm(encryption.algorithm()); + builder.sseCustomerKeyMD5(md5Checksum(encryption)); + }); + setEncryptionSettings(builder, context.s3SseContext()); }) .build(); @@ -211,9 +238,17 @@ private void flushBuffer(boolean finished) client.putObject(request, RequestBody.fromByteBuffer(bytes)); return; } + catch (S3Exception e) { + failed = true; + // when `location` already exists, the operation will fail with `412 Precondition Failed` + if (e.statusCode() == HTTP_PRECON_FAILED) { + throw new FileAlreadyExistsException(location.toString()); + } + throw new IOException("Put failed for bucket [%s] key [%s]: %s".formatted(location.bucket(), location.key(), e), e); + } catch (SdkException e) { failed = true; - throw new IOException(e); + throw new IOException("Put failed for bucket [%s] key [%s]: %s".formatted(location.bucket(), location.key(), e), e); } } @@ -241,7 +276,7 @@ private void flushBuffer(boolean finished) throw e; } multipartUploadStarted = true; - inProgressUploadFuture = supplyAsync(() -> uploadPage(data, length)); + inProgressUploadFuture = supplyAsync(() -> uploadPage(data, length), uploadExecutor); } } @@ -268,16 +303,19 @@ private CompletedPart uploadPage(byte[] data, int length) { if (uploadId.isEmpty()) { CreateMultipartUploadRequest request = CreateMultipartUploadRequest.builder() + .overrideConfiguration(context::applyCredentialProviderOverride) + .acl(cannedAcl) .requestPayer(requestPayer) .bucket(location.bucket()) .key(location.key()) - .applyMutation(builder -> { - switch (sseType) { - case NONE -> { /* ignored */ } - case S3 -> builder.serverSideEncryption(AES256); - case KMS -> builder.serverSideEncryption(AWS_KMS).ssekmsKeyId(sseKmsKeyId); - } - }) + .storageClass(storageClass) + .applyMutation(builder -> + key.ifPresentOrElse( + encryption -> + builder.sseCustomerKey(encoded(encryption)) + .sseCustomerAlgorithm(encryption.algorithm()) + .sseCustomerKeyMD5(md5Checksum(encryption)), + () -> setEncryptionSettings(builder, context.s3SseContext()))) .build(); uploadId = Optional.of(client.createMultipartUpload(request).uploadId()); @@ -285,12 +323,20 @@ private CompletedPart uploadPage(byte[] data, int length) currentPartNumber++; UploadPartRequest request = UploadPartRequest.builder() + .overrideConfiguration(context::applyCredentialProviderOverride) .requestPayer(requestPayer) .bucket(location.bucket()) .key(location.key()) .contentLength((long) length) .uploadId(uploadId.get()) .partNumber(currentPartNumber) + .applyMutation(builder -> + key.ifPresentOrElse( + encryption -> + builder.sseCustomerKey(encoded(encryption)) + .sseCustomerAlgorithm(encryption.algorithm()) + .sseCustomerKeyMD5(md5Checksum(encryption)), + () -> setEncryptionSettings(builder, context.s3SseContext()))) .build(); ByteBuffer bytes = ByteBuffer.wrap(data, 0, length); @@ -309,11 +355,23 @@ private CompletedPart uploadPage(byte[] data, int length) private void finishUpload(String uploadId) { CompleteMultipartUploadRequest request = CompleteMultipartUploadRequest.builder() + .overrideConfiguration(context::applyCredentialProviderOverride) .requestPayer(requestPayer) .bucket(location.bucket()) .key(location.key()) .uploadId(uploadId) .multipartUpload(x -> x.parts(parts)) + .applyMutation(builder -> { + key.ifPresentOrElse( + encryption -> + builder.sseCustomerKey(encoded(encryption)) + .sseCustomerAlgorithm(encryption.algorithm()) + .sseCustomerKeyMD5(md5Checksum(encryption)), + () -> setEncryptionSettings(builder, context.s3SseContext())); + //if (exclusiveCreate) { + // builder.ifNoneMatch("*"); + //} + }) .build(); client.completeMultipartUpload(request); @@ -322,6 +380,7 @@ private void finishUpload(String uploadId) private void abortUpload() { uploadId.map(id -> AbortMultipartUploadRequest.builder() + .overrideConfiguration(context::applyCredentialProviderOverride) .requestPayer(requestPayer) .bucket(location.bucket()) .key(location.key()) @@ -342,4 +401,12 @@ private void abortUploadSuppressed(Throwable throwable) } } } + + public static int clamp(long value, int min, int max) + { + if (min > max) { + throw new IllegalArgumentException(min + " > " + max); + } + return (int) Math.min(max, Math.max(value, min)); + } } diff --git a/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3SecurityMapping.java b/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3SecurityMapping.java new file mode 100644 index 000000000000..ff49aad13594 --- /dev/null +++ b/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3SecurityMapping.java @@ -0,0 +1,191 @@ +/* + * 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.filesystem.s3; + +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableSet; +import io.trino.filesystem.Location; +import io.trino.spi.security.ConnectorIdentity; +import software.amazon.awssdk.auth.credentials.AwsBasicCredentials; +import software.amazon.awssdk.auth.credentials.AwsCredentials; + +import java.util.Collection; +import java.util.List; +import java.util.Optional; +import java.util.Set; +import java.util.function.Predicate; +import java.util.regex.Pattern; + +import static com.google.common.base.Preconditions.checkArgument; +import static java.util.Objects.requireNonNull; + +public final class S3SecurityMapping +{ + private final Predicate user; + private final Predicate> group; + private final Predicate prefix; + private final Optional iamRole; + private final Optional roleSessionName; + private final Set allowedIamRoles; + private final Optional kmsKeyId; + private final Set allowedKmsKeyIds; + private final Optional sseCustomerKey; + private final Set allowedSseCustomerKeys; + private final Optional credentials; + private final boolean useClusterDefault; + private final Optional endpoint; + private final Optional region; + + @JsonCreator + public S3SecurityMapping( + @JsonProperty("user") Optional user, + @JsonProperty("group") Optional group, + @JsonProperty("prefix") Optional prefix, + @JsonProperty("iamRole") Optional iamRole, + @JsonProperty("roleSessionName") Optional roleSessionName, + @JsonProperty("allowedIamRoles") Optional> allowedIamRoles, + @JsonProperty("kmsKeyId") Optional kmsKeyId, + @JsonProperty("allowedKmsKeyIds") Optional> allowedKmsKeyIds, + @JsonProperty("sseCustomerKey") Optional sseCustomerKey, + @JsonProperty("allowedSseCustomerKeys") Optional> allowedSseCustomerKeys, + @JsonProperty("accessKey") Optional accessKey, + @JsonProperty("secretKey") Optional secretKey, + @JsonProperty("useClusterDefault") Optional useClusterDefault, + @JsonProperty("endpoint") Optional endpoint, + @JsonProperty("region") Optional region) + { + this.user = user + .map(S3SecurityMapping::toPredicate) + .orElse(ignore -> true); + this.group = group + .map(S3SecurityMapping::toPredicate) + .map(S3SecurityMapping::anyMatch) + .orElse(ignore -> true); + this.prefix = prefix + .map(Location::of) + .map(S3Location::new) + .map(S3SecurityMapping::prefixPredicate) + .orElse(ignore -> true); + + this.iamRole = requireNonNull(iamRole, "iamRole is null"); + this.roleSessionName = requireNonNull(roleSessionName, "roleSessionName is null"); + checkArgument(roleSessionName.isEmpty() || iamRole.isPresent(), "iamRole must be provided when roleSessionName is provided"); + + this.allowedIamRoles = ImmutableSet.copyOf(allowedIamRoles.orElse(ImmutableList.of())); + + this.kmsKeyId = requireNonNull(kmsKeyId, "kmsKeyId is null"); + + this.allowedKmsKeyIds = ImmutableSet.copyOf(allowedKmsKeyIds.orElse(ImmutableList.of())); + + this.sseCustomerKey = requireNonNull(sseCustomerKey, "sseCustomerKey is null"); + + this.allowedSseCustomerKeys = allowedSseCustomerKeys.map(ImmutableSet::copyOf).orElse(ImmutableSet.of()); + + requireNonNull(accessKey, "accessKey is null"); + requireNonNull(secretKey, "secretKey is null"); + checkArgument(accessKey.isPresent() == secretKey.isPresent(), "accessKey and secretKey must be provided together"); + this.credentials = accessKey.map(access -> AwsBasicCredentials.create(access, secretKey.get())); + + this.useClusterDefault = useClusterDefault.orElse(false); + boolean roleOrCredentialsArePresent = !this.allowedIamRoles.isEmpty() || iamRole.isPresent() || credentials.isPresent(); + checkArgument(this.useClusterDefault != roleOrCredentialsArePresent, "must either allow useClusterDefault role or provide role and/or credentials"); + + checkArgument(!this.useClusterDefault || this.kmsKeyId.isEmpty(), "KMS key ID cannot be provided together with useClusterDefault"); + checkArgument(!this.useClusterDefault || this.sseCustomerKey.isEmpty(), "SSE Customer key cannot be provided together with useClusterDefault"); + checkArgument(this.kmsKeyId.isEmpty() || this.sseCustomerKey.isEmpty(), "SSE Customer key cannot be provided together with KMS key ID"); + + this.endpoint = requireNonNull(endpoint, "endpoint is null"); + this.region = requireNonNull(region, "region is null"); + } + + boolean matches(ConnectorIdentity identity, S3Location location) + { + return user.test(identity.getUser()) && + group.test(identity.getGroups()) && + prefix.test(location); + } + + public Optional iamRole() + { + return iamRole; + } + + public Optional roleSessionName() + { + return roleSessionName; + } + + public Set allowedIamRoles() + { + return allowedIamRoles; + } + + public Optional kmsKeyId() + { + return kmsKeyId; + } + + public Set allowedKmsKeyIds() + { + return allowedKmsKeyIds; + } + + public Optional sseCustomerKey() + { + return sseCustomerKey; + } + + public Set allowedSseCustomerKeys() + { + return allowedSseCustomerKeys; + } + + public Optional credentials() + { + return credentials; + } + + public boolean useClusterDefault() + { + return useClusterDefault; + } + + public Optional endpoint() + { + return endpoint; + } + + public Optional region() + { + return region; + } + + private static Predicate prefixPredicate(S3Location prefix) + { + return value -> prefix.bucket().equals(value.bucket()) && + value.key().startsWith(prefix.key()); + } + + private static Predicate toPredicate(Pattern pattern) + { + return value -> pattern.matcher(value).matches(); + } + + private static Predicate> anyMatch(Predicate predicate) + { + return values -> values.stream().anyMatch(predicate); + } +} diff --git a/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3SecurityMappingConfig.java b/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3SecurityMappingConfig.java new file mode 100644 index 000000000000..8bf2b8cbe499 --- /dev/null +++ b/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3SecurityMappingConfig.java @@ -0,0 +1,148 @@ +/* + * 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.filesystem.s3; + +import io.airlift.configuration.Config; +import io.airlift.configuration.ConfigDescription; +import io.airlift.configuration.validation.FileExists; +import io.airlift.units.Duration; +import jakarta.validation.constraints.AssertTrue; +import jakarta.validation.constraints.NotNull; + +import java.io.File; +import java.net.URI; +import java.util.Optional; + +public class S3SecurityMappingConfig +{ + private File configFile; + private URI configUri; + private String jsonPointer = ""; + private String roleCredentialName; + private String kmsKeyIdCredentialName; + private String sseCustomerKeyCredentialName; + private Duration refreshPeriod; + private String colonReplacement; + + public Optional<@FileExists File> getConfigFile() + { + return Optional.ofNullable(configFile); + } + + @Config("s3.security-mapping.config-file") + @ConfigDescription("Path to the JSON security mappings file") + public S3SecurityMappingConfig setConfigFile(File configFile) + { + this.configFile = configFile; + return this; + } + + public Optional getConfigUri() + { + return Optional.ofNullable(configUri); + } + + @Config("s3.security-mapping.config-uri") + @ConfigDescription("HTTP URI of the JSON security mappings") + public S3SecurityMappingConfig setConfigUri(URI configUri) + { + this.configUri = configUri; + return this; + } + + @NotNull + public String getJsonPointer() + { + return jsonPointer; + } + + @Config("s3.security-mapping.json-pointer") + @ConfigDescription("JSON pointer (RFC 6901) to mappings inside JSON config") + public S3SecurityMappingConfig setJsonPointer(String jsonPointer) + { + this.jsonPointer = jsonPointer; + return this; + } + + public Optional getRoleCredentialName() + { + return Optional.ofNullable(roleCredentialName); + } + + @Config("s3.security-mapping.iam-role-credential-name") + @ConfigDescription("Name of the extra credential used to provide IAM role") + public S3SecurityMappingConfig setRoleCredentialName(String roleCredentialName) + { + this.roleCredentialName = roleCredentialName; + return this; + } + + public Optional getKmsKeyIdCredentialName() + { + return Optional.ofNullable(kmsKeyIdCredentialName); + } + + @Config("s3.security-mapping.kms-key-id-credential-name") + @ConfigDescription("Name of the extra credential used to provide KMS Key ID") + public S3SecurityMappingConfig setKmsKeyIdCredentialName(String kmsKeyIdCredentialName) + { + this.kmsKeyIdCredentialName = kmsKeyIdCredentialName; + return this; + } + + public Optional getSseCustomerKeyCredentialName() + { + return Optional.ofNullable(sseCustomerKeyCredentialName); + } + + @Config("s3.security-mapping.sse-customer-key-credential-name") + @ConfigDescription("Name of the extra credential used to provide SSE Customer key") + public S3SecurityMappingConfig setSseCustomerKeyCredentialName(String sseCustomerKeyCredentialName) + { + this.sseCustomerKeyCredentialName = sseCustomerKeyCredentialName; + return this; + } + + public Optional getRefreshPeriod() + { + return Optional.ofNullable(refreshPeriod); + } + + @Config("s3.security-mapping.refresh-period") + @ConfigDescription("How often to refresh the security mapping configuration") + public S3SecurityMappingConfig setRefreshPeriod(Duration refreshPeriod) + { + this.refreshPeriod = refreshPeriod; + return this; + } + + public Optional getColonReplacement() + { + return Optional.ofNullable(colonReplacement); + } + + @Config("s3.security-mapping.colon-replacement") + @ConfigDescription("Value used in place of colon for IAM role name in extra credentials") + public S3SecurityMappingConfig setColonReplacement(String colonReplacement) + { + this.colonReplacement = colonReplacement; + return this; + } + + @AssertTrue(message = "Exactly one of s3.security-mapping.config-file or s3.security-mapping.config-uri must be set") + public boolean validateMappingsConfig() + { + return (configFile == null) != (configUri == null); + } +} diff --git a/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3SecurityMappingEnabledConfig.java b/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3SecurityMappingEnabledConfig.java new file mode 100644 index 000000000000..dd2b2867d078 --- /dev/null +++ b/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3SecurityMappingEnabledConfig.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.filesystem.s3; + +import io.airlift.configuration.Config; + +public class S3SecurityMappingEnabledConfig +{ + private boolean enabled; + + public boolean isEnabled() + { + return enabled; + } + + @Config("s3.security-mapping.enabled") + public S3SecurityMappingEnabledConfig setEnabled(boolean enabled) + { + this.enabled = enabled; + return this; + } +} diff --git a/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3SecurityMappingProvider.java b/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3SecurityMappingProvider.java new file mode 100644 index 000000000000..008fb835f6bf --- /dev/null +++ b/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3SecurityMappingProvider.java @@ -0,0 +1,171 @@ +/* + * 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.filesystem.s3; + +import com.google.inject.Inject; +import io.airlift.units.Duration; +import io.trino.filesystem.Location; +import io.trino.spi.security.AccessDeniedException; +import io.trino.spi.security.ConnectorIdentity; + +import java.util.Optional; +import java.util.function.Supplier; + +import static com.google.common.base.Suppliers.memoize; +import static com.google.common.base.Suppliers.memoizeWithExpiration; +import static com.google.common.base.Verify.verify; +import static java.util.Objects.requireNonNull; +import static java.util.concurrent.TimeUnit.MILLISECONDS; + +final class S3SecurityMappingProvider +{ + private final Supplier mappingsProvider; + private final Optional roleCredentialName; + private final Optional kmsKeyIdCredentialName; + private final Optional sseCustomerKeyCredentialName; + private final Optional colonReplacement; + + @Inject + public S3SecurityMappingProvider(S3SecurityMappingConfig config, Supplier mappingsProvider) + { + this(mappingsProvider(mappingsProvider, config.getRefreshPeriod()), + config.getRoleCredentialName(), + config.getKmsKeyIdCredentialName(), + config.getSseCustomerKeyCredentialName(), + config.getColonReplacement()); + } + + public S3SecurityMappingProvider( + Supplier mappingsProvider, + Optional roleCredentialName, + Optional kmsKeyIdCredentialName, + Optional sseCustomerKeyCredentialName, + Optional colonReplacement) + { + this.mappingsProvider = requireNonNull(mappingsProvider, "mappingsProvider is null"); + this.roleCredentialName = requireNonNull(roleCredentialName, "roleCredentialName is null"); + this.kmsKeyIdCredentialName = requireNonNull(kmsKeyIdCredentialName, "kmsKeyIdCredentialName is null"); + this.sseCustomerKeyCredentialName = requireNonNull(sseCustomerKeyCredentialName, "customerKeyCredentialName is null"); + this.colonReplacement = requireNonNull(colonReplacement, "colonReplacement is null"); + } + + public Optional getMapping(ConnectorIdentity identity, Location location) + { + S3SecurityMapping mapping = mappingsProvider.get().getMapping(identity, new S3Location(location)) + .orElseThrow(() -> new AccessDeniedException("No matching S3 security mapping")); + + if (mapping.useClusterDefault()) { + return Optional.empty(); + } + + return Optional.of(new S3SecurityMappingResult( + mapping.credentials(), + selectRole(mapping, identity), + mapping.roleSessionName().map(name -> name.replace("${USER}", identity.getUser())), + selectKmsKeyId(mapping, identity), + getSseCustomerKey(mapping, identity), + mapping.endpoint(), + mapping.region())); + } + + private Optional selectRole(S3SecurityMapping mapping, ConnectorIdentity identity) + { + Optional optionalSelected = getRoleFromExtraCredential(identity); + + if (optionalSelected.isEmpty()) { + if (!mapping.allowedIamRoles().isEmpty() && mapping.iamRole().isEmpty()) { + throw new AccessDeniedException("No S3 role selected and mapping has no default role"); + } + verify(mapping.iamRole().isPresent() || mapping.credentials().isPresent(), "mapping must have role or credential"); + return mapping.iamRole(); + } + + String selected = optionalSelected.get(); + + // selected role must match default or be allowed + if (!selected.equals(mapping.iamRole().orElse(null)) && + !mapping.allowedIamRoles().contains(selected)) { + throw new AccessDeniedException("Selected S3 role is not allowed: " + selected); + } + + return optionalSelected; + } + + private Optional getRoleFromExtraCredential(ConnectorIdentity identity) + { + return roleCredentialName + .map(name -> identity.getExtraCredentials().get(name)) + .map(role -> colonReplacement + .map(replacement -> role.replace(replacement, ":")) + .orElse(role)); + } + + private Optional selectKmsKeyId(S3SecurityMapping mapping, ConnectorIdentity identity) + { + Optional userSelected = getKmsKeyIdFromExtraCredential(identity); + + if (userSelected.isEmpty()) { + return mapping.kmsKeyId(); + } + + String selected = userSelected.get(); + + // selected key ID must match default or be allowed + if (!selected.equals(mapping.kmsKeyId().orElse(null)) && + !mapping.allowedKmsKeyIds().contains(selected) && + !mapping.allowedKmsKeyIds().contains("*")) { + throw new AccessDeniedException("Selected KMS Key ID is not allowed"); + } + + return userSelected; + } + + private Optional getKmsKeyIdFromExtraCredential(ConnectorIdentity identity) + { + return kmsKeyIdCredentialName.map(name -> identity.getExtraCredentials().get(name)); + } + + private Optional getSseCustomerKey(S3SecurityMapping mapping, ConnectorIdentity identity) + { + Optional providedKey = getSseCustomerKeyFromExtraCredential(identity); + + if (providedKey.isEmpty()) { + return mapping.sseCustomerKey(); + } + if (mapping.sseCustomerKey().isPresent() && mapping.allowedSseCustomerKeys().isEmpty()) { + throw new AccessDeniedException("allowedSseCustomerKeys must be set if sseCustomerKey is provided"); + } + + String selected = providedKey.get(); + + if (selected.equals(mapping.sseCustomerKey().orElse(null)) || + mapping.allowedSseCustomerKeys().contains(selected) || + mapping.allowedSseCustomerKeys().contains("*")) { + return providedKey; + } + throw new AccessDeniedException("Provided SSE Customer Key is not allowed"); + } + + private Optional getSseCustomerKeyFromExtraCredential(ConnectorIdentity identity) + { + return sseCustomerKeyCredentialName.map(name -> identity.getExtraCredentials().get(name)); + } + + private static Supplier mappingsProvider(Supplier supplier, Optional refreshPeriod) + { + return refreshPeriod + .map(refresh -> memoizeWithExpiration(supplier::get, refresh.toMillis(), MILLISECONDS)) + .orElseGet(() -> memoize(supplier::get)); + } +} diff --git a/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3SecurityMappingResult.java b/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3SecurityMappingResult.java new file mode 100644 index 000000000000..2d2ef7dc2e92 --- /dev/null +++ b/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3SecurityMappingResult.java @@ -0,0 +1,48 @@ +/* + * 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.filesystem.s3; + +import software.amazon.awssdk.auth.credentials.AwsCredentials; +import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider; +import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider; + +import java.util.Optional; + +import static java.util.Objects.requireNonNull; + +record S3SecurityMappingResult( + Optional credentials, + Optional iamRole, + Optional roleSessionName, + Optional kmsKeyId, + Optional sseCustomerKey, + Optional endpoint, + Optional region) +{ + public S3SecurityMappingResult + { + requireNonNull(credentials, "credentials is null"); + requireNonNull(iamRole, "iamRole is null"); + requireNonNull(roleSessionName, "roleSessionName is null"); + requireNonNull(kmsKeyId, "kmsKeyId is null"); + requireNonNull(sseCustomerKey, "sseCustomerKey is null"); + requireNonNull(endpoint, "endpoint is null"); + requireNonNull(region, "region is null"); + } + + public Optional credentialsProvider() + { + return credentials.map(StaticCredentialsProvider::create); + } +} diff --git a/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3SecurityMappings.java b/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3SecurityMappings.java new file mode 100644 index 000000000000..22ecb7945e03 --- /dev/null +++ b/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3SecurityMappings.java @@ -0,0 +1,42 @@ +/* + * 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.filesystem.s3; + +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; +import com.google.common.collect.ImmutableList; +import io.trino.spi.security.ConnectorIdentity; + +import java.util.List; +import java.util.Optional; + +import static java.util.Objects.requireNonNull; + +public final class S3SecurityMappings +{ + private final List mappings; + + @JsonCreator + public S3SecurityMappings(@JsonProperty("mappings") List mappings) + { + this.mappings = ImmutableList.copyOf(requireNonNull(mappings, "mappings is null")); + } + + Optional getMapping(ConnectorIdentity identity, S3Location location) + { + return mappings.stream() + .filter(mapping -> mapping.matches(identity, location)) + .findFirst(); + } +} diff --git a/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3SecurityMappingsFileSource.java b/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3SecurityMappingsFileSource.java new file mode 100644 index 000000000000..aa96fdc26912 --- /dev/null +++ b/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3SecurityMappingsFileSource.java @@ -0,0 +1,41 @@ +/* + * 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.filesystem.s3; + +import com.google.inject.Inject; + +import java.io.File; +import java.util.function.Supplier; + +import static io.trino.plugin.base.util.JsonUtils.parseJson; + +class S3SecurityMappingsFileSource + implements Supplier +{ + private final File configFile; + private final String jsonPointer; + + @Inject + public S3SecurityMappingsFileSource(S3SecurityMappingConfig config) + { + this.configFile = config.getConfigFile().orElseThrow(); + this.jsonPointer = config.getJsonPointer(); + } + + @Override + public S3SecurityMappings get() + { + return parseJson(configFile.toPath(), jsonPointer, S3SecurityMappings.class); + } +} diff --git a/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3SecurityMappingsUriSource.java b/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3SecurityMappingsUriSource.java new file mode 100644 index 000000000000..e77a17348b77 --- /dev/null +++ b/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3SecurityMappingsUriSource.java @@ -0,0 +1,64 @@ +/* + * 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.filesystem.s3; + +import com.google.common.annotations.VisibleForTesting; +import com.google.inject.Inject; +import io.airlift.http.client.HttpClient; +import io.airlift.http.client.HttpStatus; +import io.airlift.http.client.Request; +import io.airlift.http.client.StringResponseHandler.StringResponse; +import io.trino.filesystem.s3.S3FileSystemModule.ForS3SecurityMapping; + +import java.net.URI; +import java.util.function.Supplier; + +import static io.airlift.http.client.Request.Builder.prepareGet; +import static io.airlift.http.client.StringResponseHandler.createStringResponseHandler; +import static io.trino.plugin.base.util.JsonUtils.parseJson; +import static java.util.Objects.requireNonNull; + +class S3SecurityMappingsUriSource + implements Supplier +{ + private final URI configUri; + private final HttpClient httpClient; + private final String jsonPointer; + + @Inject + public S3SecurityMappingsUriSource(S3SecurityMappingConfig config, @ForS3SecurityMapping HttpClient httpClient) + { + this.configUri = config.getConfigUri().orElseThrow(); + this.httpClient = requireNonNull(httpClient, "httpClient is null"); + this.jsonPointer = config.getJsonPointer(); + } + + @Override + public S3SecurityMappings get() + { + return parseJson(getRawJsonString(), jsonPointer, S3SecurityMappings.class); + } + + @VisibleForTesting + String getRawJsonString() + { + Request request = prepareGet().setUri(configUri).build(); + StringResponse response = httpClient.execute(request, createStringResponseHandler()); + int status = response.getStatusCode(); + if (status != HttpStatus.OK.code()) { + throw new RuntimeException("Request to '%s' returned unexpected status code: %s".formatted(configUri, status)); + } + return response.getBody(); + } +} diff --git a/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3SseCUtils.java b/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3SseCUtils.java new file mode 100644 index 000000000000..47d1975132b8 --- /dev/null +++ b/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3SseCUtils.java @@ -0,0 +1,42 @@ +/* + * 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.filesystem.s3; + +import io.trino.filesystem.encryption.EncryptionKey; + +import java.security.MessageDigest; +import java.security.NoSuchAlgorithmException; +import java.util.Base64; + +final class S3SseCUtils +{ + private S3SseCUtils() {} + + public static String encoded(EncryptionKey key) + { + return Base64.getEncoder().encodeToString(key.key()); + } + + public static String md5Checksum(EncryptionKey key) + { + try { + MessageDigest digest = MessageDigest.getInstance("MD5"); + digest.update(key.key()); + return Base64.getEncoder().encodeToString(digest.digest()); + } + catch (NoSuchAlgorithmException e) { + throw new RuntimeException(e); + } + } +} diff --git a/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3SseCustomerKey.java b/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3SseCustomerKey.java new file mode 100644 index 000000000000..9c9a4be3bed2 --- /dev/null +++ b/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3SseCustomerKey.java @@ -0,0 +1,35 @@ +/* + * 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.filesystem.s3; + +import static java.util.Objects.requireNonNull; +import static software.amazon.awssdk.utils.BinaryUtils.fromBase64; +import static software.amazon.awssdk.utils.Md5Utils.md5AsBase64; + +public record S3SseCustomerKey(String key, String md5, String algorithm) +{ + private static final String SSE_C_ALGORITHM = "AES256"; + + public S3SseCustomerKey + { + requireNonNull(key, "key is null"); + requireNonNull(md5, "md5 is null"); + requireNonNull(algorithm, "algorithm is null"); + } + + public static S3SseCustomerKey onAes256(String key) + { + return new S3SseCustomerKey(key, md5AsBase64(fromBase64(key)), SSE_C_ALGORITHM); + } +} diff --git a/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3SseRequestConfigurator.java b/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3SseRequestConfigurator.java new file mode 100644 index 000000000000..8559117a187c --- /dev/null +++ b/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3SseRequestConfigurator.java @@ -0,0 +1,101 @@ +/* + * 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.filesystem.s3; + +import io.trino.filesystem.s3.S3Context.S3SseContext; +import software.amazon.awssdk.services.s3.model.CompleteMultipartUploadRequest; +import software.amazon.awssdk.services.s3.model.CreateMultipartUploadRequest; +import software.amazon.awssdk.services.s3.model.GetObjectRequest; +import software.amazon.awssdk.services.s3.model.HeadObjectRequest; +import software.amazon.awssdk.services.s3.model.PutObjectRequest; +import software.amazon.awssdk.services.s3.model.UploadPartRequest; + +import static io.trino.filesystem.s3.S3FileSystemConfig.S3SseType.CUSTOMER; +import static software.amazon.awssdk.services.s3.model.ServerSideEncryption.AES256; +import static software.amazon.awssdk.services.s3.model.ServerSideEncryption.AWS_KMS; + +public final class S3SseRequestConfigurator +{ + private S3SseRequestConfigurator() {} + + public static void setEncryptionSettings(PutObjectRequest.Builder builder, S3SseContext context) + { + switch (context.sseType()) { + case NONE -> { /* ignored */ } + case S3 -> builder.serverSideEncryption(AES256); + case KMS -> context.sseKmsKeyId().ifPresent(builder.serverSideEncryption(AWS_KMS)::ssekmsKeyId); + case CUSTOMER -> { + context.sseCustomerKey().ifPresent(s3SseCustomerKey -> + builder.sseCustomerAlgorithm(s3SseCustomerKey.algorithm()) + .sseCustomerKey(s3SseCustomerKey.key()) + .sseCustomerKeyMD5(s3SseCustomerKey.md5())); + } + } + } + + public static void setEncryptionSettings(CreateMultipartUploadRequest.Builder builder, S3SseContext context) + { + switch (context.sseType()) { + case NONE -> { /* ignored */ } + case S3 -> builder.serverSideEncryption(AES256); + case KMS -> context.sseKmsKeyId().ifPresent(builder.serverSideEncryption(AWS_KMS)::ssekmsKeyId); + case CUSTOMER -> { + context.sseCustomerKey().ifPresent(s3SseCustomerKey -> + builder.sseCustomerAlgorithm(s3SseCustomerKey.algorithm()) + .sseCustomerKey(s3SseCustomerKey.key()) + .sseCustomerKeyMD5(s3SseCustomerKey.md5())); + } + } + } + + public static void setEncryptionSettings(CompleteMultipartUploadRequest.Builder builder, S3SseContext context) + { + if (context.sseType() == CUSTOMER) { + context.sseCustomerKey().ifPresent(s3SseCustomerKey -> + builder.sseCustomerAlgorithm(s3SseCustomerKey.algorithm()) + .sseCustomerKey(s3SseCustomerKey.key()) + .sseCustomerKeyMD5(s3SseCustomerKey.md5())); + } + } + + public static void setEncryptionSettings(GetObjectRequest.Builder builder, S3SseContext context) + { + if (context.sseType().equals(CUSTOMER)) { + context.sseCustomerKey().ifPresent(s3SseCustomerKey -> + builder.sseCustomerAlgorithm(s3SseCustomerKey.algorithm()) + .sseCustomerKey(s3SseCustomerKey.key()) + .sseCustomerKeyMD5(s3SseCustomerKey.md5())); + } + } + + public static void setEncryptionSettings(HeadObjectRequest.Builder builder, S3SseContext context) + { + if (context.sseType().equals(CUSTOMER)) { + context.sseCustomerKey().ifPresent(s3SseCustomerKey -> + builder.sseCustomerAlgorithm(s3SseCustomerKey.algorithm()) + .sseCustomerKey(s3SseCustomerKey.key()) + .sseCustomerKeyMD5(s3SseCustomerKey.md5())); + } + } + + public static void setEncryptionSettings(UploadPartRequest.Builder builder, S3SseContext context) + { + if (context.sseType() == CUSTOMER) { + context.sseCustomerKey().ifPresent(s3SseCustomerKey -> + builder.sseCustomerAlgorithm(s3SseCustomerKey.algorithm()) + .sseCustomerKey(s3SseCustomerKey.key()) + .sseCustomerKeyMD5(s3SseCustomerKey.md5())); + } + } +} diff --git a/lib/trino-filesystem-s3/src/test/java/io/trino/filesystem/s3/AbstractTestS3FileSystem.java b/lib/trino-filesystem-s3/src/test/java/io/trino/filesystem/s3/AbstractTestS3FileSystem.java index 39076e1aeca4..523ba84c9a62 100644 --- a/lib/trino-filesystem-s3/src/test/java/io/trino/filesystem/s3/AbstractTestS3FileSystem.java +++ b/lib/trino-filesystem-s3/src/test/java/io/trino/filesystem/s3/AbstractTestS3FileSystem.java @@ -36,6 +36,7 @@ import java.util.List; import static com.google.common.collect.Iterables.getOnlyElement; +import static io.trino.filesystem.s3.S3FileSystem.disableStrongIntegrityChecksums; import static java.nio.charset.StandardCharsets.UTF_8; import static org.assertj.core.api.Assertions.assertThat; @@ -74,7 +75,8 @@ public void testFileWithTrailingWhitespaceAgainstNativeClient() String key = "foo/bar with whitespace "; byte[] contents = "abc foo bar".getBytes(UTF_8); s3Client.putObject( - request -> request.bucket(bucket()).key(key), + request -> request.bucket(bucket()).key(key) + .overrideConfiguration(disableStrongIntegrityChecksums()), RequestBody.fromBytes(contents.clone())); try { // Verify listing diff --git a/lib/trino-filesystem-s3/src/test/java/io/trino/filesystem/s3/TestS3FileSystemAwsS3.java b/lib/trino-filesystem-s3/src/test/java/io/trino/filesystem/s3/TestS3FileSystemAwsS3.java index e98876f3ee6e..572a218dc42a 100644 --- a/lib/trino-filesystem-s3/src/test/java/io/trino/filesystem/s3/TestS3FileSystemAwsS3.java +++ b/lib/trino-filesystem-s3/src/test/java/io/trino/filesystem/s3/TestS3FileSystemAwsS3.java @@ -14,6 +14,7 @@ package io.trino.filesystem.s3; import io.airlift.units.DataSize; +import io.opentelemetry.api.OpenTelemetry; import software.amazon.awssdk.auth.credentials.AwsBasicCredentials; import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider; import software.amazon.awssdk.regions.Region; @@ -56,11 +57,12 @@ protected S3Client createS3Client() @Override protected S3FileSystemFactory createS3FileSystemFactory() { - return new S3FileSystemFactory(new S3FileSystemConfig() + return new S3FileSystemFactory(OpenTelemetry.noop(), new S3FileSystemConfig() .setAwsAccessKey(accessKey) .setAwsSecretKey(secretKey) .setRegion(region) - .setStreamingPartSize(DataSize.valueOf("5.5MB"))); + .setSupportsExclusiveCreate(true) + .setStreamingPartSize(DataSize.valueOf("5.5MB")), new S3FileSystemStats()); } private static String environmentVariable(String name) diff --git a/lib/trino-filesystem-s3/src/test/java/io/trino/filesystem/s3/TestS3FileSystemConfig.java b/lib/trino-filesystem-s3/src/test/java/io/trino/filesystem/s3/TestS3FileSystemConfig.java index dd07e1f6edf5..0c26f8e6d10a 100644 --- a/lib/trino-filesystem-s3/src/test/java/io/trino/filesystem/s3/TestS3FileSystemConfig.java +++ b/lib/trino-filesystem-s3/src/test/java/io/trino/filesystem/s3/TestS3FileSystemConfig.java @@ -16,6 +16,7 @@ import com.google.common.collect.ImmutableMap; import com.google.common.net.HostAndPort; import io.airlift.units.DataSize; +import io.airlift.units.Duration; import io.trino.filesystem.s3.S3FileSystemConfig.S3SseType; import org.junit.jupiter.api.Test; @@ -25,6 +26,10 @@ import static io.airlift.configuration.testing.ConfigAssertions.assertRecordedDefaults; import static io.airlift.configuration.testing.ConfigAssertions.recordDefaults; import static io.airlift.units.DataSize.Unit.MEGABYTE; +import static io.trino.filesystem.s3.S3FileSystemConfig.RetryMode.LEGACY; +import static io.trino.filesystem.s3.S3FileSystemConfig.RetryMode.STANDARD; +import static io.trino.filesystem.s3.S3FileSystemConfig.StorageClassType.STANDARD_IA; +import static java.util.concurrent.TimeUnit.MINUTES; public class TestS3FileSystemConfig { @@ -42,13 +47,30 @@ public void testDefaults() .setExternalId(null) .setStsEndpoint(null) .setStsRegion(null) + .setStorageClass(S3FileSystemConfig.StorageClassType.STANDARD) + .setCannedAcl(S3FileSystemConfig.ObjectCannedAcl.NONE) .setSseType(S3SseType.NONE) + .setRetryMode(LEGACY) + .setMaxErrorRetries(10) .setSseKmsKeyId(null) + .setUseWebIdentityTokenCredentialsProvider(false) + .setSseCustomerKey(null) .setStreamingPartSize(DataSize.of(16, MEGABYTE)) .setRequesterPays(false) - .setMaxConnections(null) + .setMaxConnections(500) + .setConnectionTtl(null) + .setConnectionMaxIdleTime(null) + .setSocketConnectTimeout(null) + .setSocketReadTimeout(null) + .setTcpKeepAlive(false) .setHttpProxy(null) - .setHttpProxySecure(false)); + .setHttpProxySecure(false) + .setNonProxyHosts(null) + .setHttpProxyUsername(null) + .setHttpProxyPassword(null) + .setHttpProxyPreemptiveBasicProxyAuth(false) + .setSupportsExclusiveCreate(true) + .setApplicationId("Trino")); } @Test @@ -65,13 +87,30 @@ public void testExplicitPropertyMappings() .put("s3.external-id", "myid") .put("s3.sts.endpoint", "sts.example.com") .put("s3.sts.region", "us-west-2") + .put("s3.storage-class", "STANDARD_IA") + .put("s3.canned-acl", "BUCKET_OWNER_FULL_CONTROL") + .put("s3.retry-mode", "STANDARD") + .put("s3.max-error-retries", "12") .put("s3.sse.type", "KMS") .put("s3.sse.kms-key-id", "mykey") + .put("s3.sse.customer-key", "customerKey") + .put("s3.use-web-identity-token-credentials-provider", "true") .put("s3.streaming.part-size", "42MB") .put("s3.requester-pays", "true") .put("s3.max-connections", "42") + .put("s3.connection-ttl", "1m") + .put("s3.connection-max-idle-time", "2m") + .put("s3.socket-connect-timeout", "3m") + .put("s3.socket-read-timeout", "4m") + .put("s3.tcp-keep-alive", "true") .put("s3.http-proxy", "localhost:8888") .put("s3.http-proxy.secure", "true") + .put("s3.http-proxy.non-proxy-hosts", "test1,test2,test3") + .put("s3.http-proxy.username", "test") + .put("s3.http-proxy.password", "test") + .put("s3.http-proxy.preemptive-basic-auth", "true") + .put("s3.exclusive-create", "false") + .put("s3.application-id", "application id") .buildOrThrow(); S3FileSystemConfig expected = new S3FileSystemConfig() @@ -85,13 +124,30 @@ public void testExplicitPropertyMappings() .setExternalId("myid") .setStsEndpoint("sts.example.com") .setStsRegion("us-west-2") + .setStorageClass(STANDARD_IA) + .setCannedAcl(S3FileSystemConfig.ObjectCannedAcl.BUCKET_OWNER_FULL_CONTROL) .setStreamingPartSize(DataSize.of(42, MEGABYTE)) + .setRetryMode(STANDARD) + .setMaxErrorRetries(12) .setSseType(S3SseType.KMS) .setSseKmsKeyId("mykey") + .setUseWebIdentityTokenCredentialsProvider(true) + .setSseCustomerKey("customerKey") .setRequesterPays(true) .setMaxConnections(42) + .setConnectionTtl(new Duration(1, MINUTES)) + .setConnectionMaxIdleTime(new Duration(2, MINUTES)) + .setSocketConnectTimeout(new Duration(3, MINUTES)) + .setSocketReadTimeout(new Duration(4, MINUTES)) + .setTcpKeepAlive(true) .setHttpProxy(HostAndPort.fromParts("localhost", 8888)) - .setHttpProxySecure(true); + .setHttpProxySecure(true) + .setNonProxyHosts("test1, test2, test3") + .setHttpProxyUsername("test") + .setHttpProxyPassword("test") + .setHttpProxyPreemptiveBasicProxyAuth(true) + .setSupportsExclusiveCreate(false) + .setApplicationId("application id"); assertFullMapping(properties, expected); } diff --git a/lib/trino-filesystem-s3/src/test/java/io/trino/filesystem/s3/TestS3FileSystemLocalStack.java b/lib/trino-filesystem-s3/src/test/java/io/trino/filesystem/s3/TestS3FileSystemLocalStack.java index 213a16f991e8..02443e129b64 100644 --- a/lib/trino-filesystem-s3/src/test/java/io/trino/filesystem/s3/TestS3FileSystemLocalStack.java +++ b/lib/trino-filesystem-s3/src/test/java/io/trino/filesystem/s3/TestS3FileSystemLocalStack.java @@ -14,6 +14,7 @@ package io.trino.filesystem.s3; import io.airlift.units.DataSize; +import io.opentelemetry.api.OpenTelemetry; import org.testcontainers.containers.localstack.LocalStackContainer; import org.testcontainers.containers.localstack.LocalStackContainer.Service; import org.testcontainers.junit.jupiter.Container; @@ -62,11 +63,11 @@ protected S3Client createS3Client() @Override protected S3FileSystemFactory createS3FileSystemFactory() { - return new S3FileSystemFactory(new S3FileSystemConfig() + return new S3FileSystemFactory(OpenTelemetry.noop(), new S3FileSystemConfig() .setAwsAccessKey(LOCALSTACK.getAccessKey()) .setAwsSecretKey(LOCALSTACK.getSecretKey()) .setEndpoint(LOCALSTACK.getEndpointOverride(Service.S3).toString()) .setRegion(LOCALSTACK.getRegion()) - .setStreamingPartSize(DataSize.valueOf("5.5MB"))); + .setStreamingPartSize(DataSize.valueOf("5.5MB")), new S3FileSystemStats()); } } diff --git a/lib/trino-filesystem-s3/src/test/java/io/trino/filesystem/s3/TestS3FileSystemMinIo.java b/lib/trino-filesystem-s3/src/test/java/io/trino/filesystem/s3/TestS3FileSystemMinIo.java index 99b9251ba828..c8cf8429792e 100644 --- a/lib/trino-filesystem-s3/src/test/java/io/trino/filesystem/s3/TestS3FileSystemMinIo.java +++ b/lib/trino-filesystem-s3/src/test/java/io/trino/filesystem/s3/TestS3FileSystemMinIo.java @@ -14,11 +14,13 @@ package io.trino.filesystem.s3; import io.airlift.units.DataSize; +import io.opentelemetry.api.OpenTelemetry; import io.trino.testing.containers.Minio; import org.junit.jupiter.api.AfterAll; import org.junit.jupiter.api.Test; import software.amazon.awssdk.auth.credentials.AwsBasicCredentials; import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider; +import software.amazon.awssdk.core.checksums.RequestChecksumCalculation; import software.amazon.awssdk.regions.Region; import software.amazon.awssdk.services.s3.S3Client; @@ -26,6 +28,7 @@ import java.net.URI; import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static software.amazon.awssdk.core.checksums.ResponseChecksumValidation.WHEN_REQUIRED; public class TestS3FileSystemMinIo extends AbstractTestS3FileSystem @@ -64,6 +67,8 @@ protected S3Client createS3Client() .endpointOverride(URI.create(minio.getMinioAddress())) .region(Region.of(Minio.MINIO_REGION)) .forcePathStyle(true) + .responseChecksumValidation(WHEN_REQUIRED) + .requestChecksumCalculation(RequestChecksumCalculation.WHEN_REQUIRED) .credentialsProvider(StaticCredentialsProvider.create( AwsBasicCredentials.create(Minio.MINIO_ACCESS_KEY, Minio.MINIO_SECRET_KEY))) .build(); @@ -72,13 +77,14 @@ protected S3Client createS3Client() @Override protected S3FileSystemFactory createS3FileSystemFactory() { - return new S3FileSystemFactory(new S3FileSystemConfig() + return new S3FileSystemFactory(OpenTelemetry.noop(), new S3FileSystemConfig() .setEndpoint(minio.getMinioAddress()) .setRegion(Minio.MINIO_REGION) .setPathStyleAccess(true) .setAwsAccessKey(Minio.MINIO_ACCESS_KEY) .setAwsSecretKey(Minio.MINIO_SECRET_KEY) - .setStreamingPartSize(DataSize.valueOf("5.5MB"))); + .setSupportsExclusiveCreate(true) + .setStreamingPartSize(DataSize.valueOf("5.5MB")), new S3FileSystemStats()); } @Test diff --git a/lib/trino-filesystem-s3/src/test/java/io/trino/filesystem/s3/TestS3FileSystemS3Mock.java b/lib/trino-filesystem-s3/src/test/java/io/trino/filesystem/s3/TestS3FileSystemS3Mock.java index 40eab3373c21..373916b802e6 100644 --- a/lib/trino-filesystem-s3/src/test/java/io/trino/filesystem/s3/TestS3FileSystemS3Mock.java +++ b/lib/trino-filesystem-s3/src/test/java/io/trino/filesystem/s3/TestS3FileSystemS3Mock.java @@ -15,6 +15,7 @@ import com.adobe.testing.s3mock.testcontainers.S3MockContainer; import io.airlift.units.DataSize; +import io.opentelemetry.api.OpenTelemetry; import org.junit.jupiter.api.Test; import org.testcontainers.junit.jupiter.Container; import org.testcontainers.junit.jupiter.Testcontainers; @@ -58,13 +59,14 @@ protected S3Client createS3Client() @Override protected S3FileSystemFactory createS3FileSystemFactory() { - return new S3FileSystemFactory(new S3FileSystemConfig() + return new S3FileSystemFactory(OpenTelemetry.noop(), new S3FileSystemConfig() .setAwsAccessKey("accesskey") .setAwsSecretKey("secretkey") .setEndpoint(S3_MOCK.getHttpEndpoint()) .setRegion(Region.US_EAST_1.id()) .setPathStyleAccess(true) - .setStreamingPartSize(DataSize.valueOf("5.5MB"))); + .setStreamingPartSize(DataSize.valueOf("5.5MB")) + .setSupportsExclusiveCreate(false), new S3FileSystemStats()); } // TODO: remove when fixed in S3Mock: https://github.com/adobe/S3Mock/pull/1131 diff --git a/lib/trino-filesystem/pom.xml b/lib/trino-filesystem/pom.xml index d942ecc71304..29aab23c3fe6 100644 --- a/lib/trino-filesystem/pom.xml +++ b/lib/trino-filesystem/pom.xml @@ -22,11 +22,26 @@ guava + + com.google.inject + guice + + + + io.airlift + configuration + + io.airlift slice + + io.airlift + units + + io.opentelemetry opentelemetry-api @@ -42,6 +57,11 @@ opentelemetry-semconv + + io.trino + trino-cache + + io.trino trino-memory-context @@ -52,6 +72,16 @@ trino-spi + + jakarta.validation + jakarta.validation-api + + + + org.weakref + jmxutils + + org.jetbrains annotations diff --git a/lib/trino-filesystem/src/main/java/io/trino/filesystem/FileEntry.java b/lib/trino-filesystem/src/main/java/io/trino/filesystem/FileEntry.java index bfbae5ecb6e6..5ee65f2f0cd2 100644 --- a/lib/trino-filesystem/src/main/java/io/trino/filesystem/FileEntry.java +++ b/lib/trino-filesystem/src/main/java/io/trino/filesystem/FileEntry.java @@ -14,10 +14,12 @@ package io.trino.filesystem; import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableSet; import java.time.Instant; import java.util.List; import java.util.Optional; +import java.util.Set; import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.collect.ImmutableList.toImmutableList; @@ -26,7 +28,7 @@ import static java.util.Comparator.comparing; import static java.util.Objects.requireNonNull; -public record FileEntry(Location location, long length, Instant lastModified, Optional> blocks) +public record FileEntry(Location location, long length, Instant lastModified, Optional> blocks, Set tags) { public FileEntry { @@ -34,6 +36,12 @@ public record FileEntry(Location location, long length, Instant lastModified, Op requireNonNull(location, "location is null"); requireNonNull(blocks, "blocks is null"); blocks = blocks.map(locations -> validatedBlocks(locations, length)); + tags = ImmutableSet.copyOf(requireNonNull(tags, "tags is null")); + } + + public FileEntry(Location location, long length, Instant lastModified, Optional> blocks) + { + this(location, length, lastModified, blocks, ImmutableSet.of()); } public record Block(List hosts, long offset, long length) diff --git a/lib/trino-filesystem/src/main/java/io/trino/filesystem/Location.java b/lib/trino-filesystem/src/main/java/io/trino/filesystem/Location.java index 08cbe7ef03df..62cf2502bfd2 100644 --- a/lib/trino-filesystem/src/main/java/io/trino/filesystem/Location.java +++ b/lib/trino-filesystem/src/main/java/io/trino/filesystem/Location.java @@ -23,6 +23,7 @@ import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.base.Preconditions.checkState; import static com.google.common.collect.Iterables.getLast; +import static io.trino.filesystem.Locations.isS3Tables; import static java.lang.Integer.parseInt; import static java.util.Objects.requireNonNull; import static java.util.function.Predicate.not; @@ -95,7 +96,10 @@ public static Location of(String location) } } - checkArgument((userInfo.isEmpty() && host.isEmpty() && port.isEmpty()) || authoritySplit.size() == 2, "Path missing in file system location: %s", location); + if (!isS3Tables(location)) { + // S3 Tables create tables under the bucket like 's3://e97725d9-dbfb-4334-784sox7edps35ncq16arh546frqa1use2b--table-s3' + checkArgument((userInfo.isEmpty() && host.isEmpty() && port.isEmpty()) || authoritySplit.size() == 2, "Path missing in file system location: %s", location); + } String path = (authoritySplit.size() == 2) ? authoritySplit.get(1) : ""; return new Location(location, Optional.of(scheme), userInfo, host, port, path); diff --git a/lib/trino-filesystem/src/main/java/io/trino/filesystem/Locations.java b/lib/trino-filesystem/src/main/java/io/trino/filesystem/Locations.java index 13694965dc76..d29b80df9acf 100644 --- a/lib/trino-filesystem/src/main/java/io/trino/filesystem/Locations.java +++ b/lib/trino-filesystem/src/main/java/io/trino/filesystem/Locations.java @@ -13,10 +13,14 @@ */ package io.trino.filesystem; +import java.util.regex.Pattern; + import static com.google.common.base.Preconditions.checkArgument; public final class Locations { + private static final Pattern S3_TABLES = Pattern.compile("s3://(?!.*/).*--table-s3"); + private Locations() {} /** @@ -47,4 +51,9 @@ public static boolean areDirectoryLocationsEquivalent(Location leftLocation, Loc return leftLocation.equals(rightLocation) || leftLocation.removeOneTrailingSlash().equals(rightLocation.removeOneTrailingSlash()); } + + public static boolean isS3Tables(String location) + { + return S3_TABLES.matcher(location).matches(); + } } diff --git a/lib/trino-filesystem/src/main/java/io/trino/filesystem/TrinoFileSystem.java b/lib/trino-filesystem/src/main/java/io/trino/filesystem/TrinoFileSystem.java index e50d8fec5d47..334a92e0a529 100644 --- a/lib/trino-filesystem/src/main/java/io/trino/filesystem/TrinoFileSystem.java +++ b/lib/trino-filesystem/src/main/java/io/trino/filesystem/TrinoFileSystem.java @@ -13,9 +13,14 @@ */ package io.trino.filesystem; +import com.google.common.base.Throwables; + +import java.io.FileNotFoundException; import java.io.IOException; +import java.time.Instant; import java.util.Collection; import java.util.Optional; +import java.util.Set; /** * TrinoFileSystem is the main abstraction for Trino to interact with data in cloud-like storage @@ -67,6 +72,20 @@ public interface TrinoFileSystem */ TrinoInputFile newInputFile(Location location, long length); + /** + * Creates a TrinoInputFile with a predeclared length and lastModifiedTime which can be used to read the file data. + * The length will be returned from {@link TrinoInputFile#length()} and the actual file length + * will never be checked. The lastModified will be returned from {@link TrinoInputFile#lastModified()} and the + * actual file last modified time will never be checked. The file location path cannot be empty, and must not end + * with a slash or whitespace. + * + * @throws IllegalArgumentException if location is not valid for this file system + */ + default TrinoInputFile newInputFile(Location location, long length, Instant lastModified) + { + return newInputFile(location, length); + } + /** * Creates a TrinoOutputFile which can be used to create or overwrite the file. The file * location path cannot be empty, and must not end with a slash or whitespace. @@ -169,4 +188,71 @@ FileIterator listFiles(Location location) */ Optional directoryExists(Location location) throws IOException; + + /** + * Creates the specified directory and any parent directories that do not exist. + * For hierarchical file systems, if the location already exists but is not a + * directory, or if the directory cannot be created, an exception is raised. + * This method does nothing for non-hierarchical file systems or if the directory + * already exists. + * + * @throws IllegalArgumentException if location is not valid for this file system + */ + void createDirectory(Location location) + throws IOException; + + /** + * Renames source to target. An exception is raised if the target already exists, + * or on non-hierarchical file systems. + * + * @throws IllegalArgumentException if location is not valid for this file system + */ + void renameDirectory(Location source, Location target) + throws IOException; + + /** + * Lists all directories that are direct descendants of the specified directory. + * If the path is empty, all directories at the root of the file system are returned. + * Otherwise, the path must end with a slash. + * If the location does not exist, an empty set is returned. + *

+ * For hierarchical file systems, if the path is not a directory, an exception is raised. + * For hierarchical file systems, if the path does not reference an existing directory, + * an empty iterator is returned. For blob file systems, all directories containing + * blobs that start with the location are listed. + * + * @throws IllegalArgumentException if location is not valid for this file system + */ + Set listDirectories(Location location) + throws IOException; + + /** + * Creates a temporary directory for the target path. The directory will be created + * using the (possibly absolute) prefix such that the directory can be renamed to + * the target path. The relative prefix will be used if the target path does not + * support the temporary prefix (which is typically absolute). + *

+ * The temporary directory is not created for non-hierarchical file systems or for + * target paths that do not support renaming, and an empty optional is returned. + * + * @throws IllegalArgumentException If the target path is not valid for this file system. + */ + Optional createTemporaryDirectory(Location targetPath, String temporaryPrefix, String relativePrefix) + throws IOException; + + /** + * Checks whether given exception is unrecoverable, so that further retries won't help + *

+ * By default, all third party (AWS, Azure, GCP) SDKs will retry appropriate exceptions + * (either client side IO errors, or 500/503), so there is no need to retry those additionally. + *

+ * If any custom retry behavior is needed, it is advised to change SDK's retry handlers, + * rather than introducing outer retry loop, which combined with SDKs default retries, + * could lead to prolonged, unnecessary retries + */ + static boolean isUnrecoverableException(Throwable throwable) + { + return Throwables.getCausalChain(throwable).stream() + .anyMatch(t -> t instanceof TrinoFileSystemException || t instanceof FileNotFoundException || t instanceof UnsupportedOperationException); + } } diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/s3select/TestUnrecoverableS3OperationException.java b/lib/trino-filesystem/src/main/java/io/trino/filesystem/TrinoFileSystemException.java similarity index 52% rename from plugin/trino-hive/src/test/java/io/trino/plugin/hive/s3select/TestUnrecoverableS3OperationException.java rename to lib/trino-filesystem/src/main/java/io/trino/filesystem/TrinoFileSystemException.java index 8bc3abadbea5..30d09b429d98 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/s3select/TestUnrecoverableS3OperationException.java +++ b/lib/trino-filesystem/src/main/java/io/trino/filesystem/TrinoFileSystemException.java @@ -11,21 +11,25 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package io.trino.plugin.hive.s3select; - -import org.testng.annotations.Test; +package io.trino.filesystem; import java.io.IOException; -import static io.trino.plugin.hive.s3select.S3SelectLineRecordReader.UnrecoverableS3OperationException; -import static org.assertj.core.api.Assertions.assertThat; - -public class TestUnrecoverableS3OperationException +/** + * Unrecoverable file system exception. + * This exception is thrown for fatal errors, or after retries have already been performed, + * so additional retries must not be performed when this is caught. + */ +public class TrinoFileSystemException + extends IOException { - @Test - public void testMessage() + public TrinoFileSystemException(String message, Throwable cause) + { + super(message, cause); + } + + public TrinoFileSystemException(String message) { - assertThat(new UnrecoverableS3OperationException("test-bucket", "test-key", new IOException("test io exception"))) - .hasMessage("java.io.IOException: test io exception (Bucket: test-bucket, Key: test-key)"); + super(message); } } diff --git a/lib/trino-filesystem/src/main/java/io/trino/filesystem/TrinoOutputFile.java b/lib/trino-filesystem/src/main/java/io/trino/filesystem/TrinoOutputFile.java index 2737f4b22a22..3144a1846166 100644 --- a/lib/trino-filesystem/src/main/java/io/trino/filesystem/TrinoOutputFile.java +++ b/lib/trino-filesystem/src/main/java/io/trino/filesystem/TrinoOutputFile.java @@ -18,6 +18,7 @@ import java.io.IOException; import java.io.OutputStream; +import java.nio.file.FileAlreadyExistsException; import static io.trino.memory.context.AggregatedMemoryContext.newSimpleAggregatedMemoryContext; @@ -29,6 +30,29 @@ default OutputStream create() return create(newSimpleAggregatedMemoryContext()); } + /** + * Create file with the specified content, atomically if possible. + * The file will be replaced if it already exists. + * If an error occurs while writing and the implementation does not + * support atomic writes, then a partial file may be written, + * or the original file may be deleted or left unchanged. + */ + void createOrOverwrite(byte[] data) + throws IOException; + + /** + * Create file exclusively and atomically with the specified content. + * If an error occurs while writing, the file will not be created. + * + * @throws FileAlreadyExistsException if the file already exists + * @throws UnsupportedOperationException if the file system does not support this operation + */ + default void createExclusive(byte[] data) + throws IOException + { + throw new UnsupportedOperationException("createExclusive not supported by " + getClass()); + } + default OutputStream createOrOverwrite() throws IOException { diff --git a/lib/trino-filesystem/src/main/java/io/trino/filesystem/UriLocation.java b/lib/trino-filesystem/src/main/java/io/trino/filesystem/UriLocation.java new file mode 100644 index 000000000000..fd55f7186305 --- /dev/null +++ b/lib/trino-filesystem/src/main/java/io/trino/filesystem/UriLocation.java @@ -0,0 +1,44 @@ +/* + * 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.filesystem; + +import com.google.common.collect.ImmutableMap; + +import java.net.URI; +import java.util.List; +import java.util.Map; + +import static java.util.Objects.requireNonNull; + +public class UriLocation +{ + private final URI uri; + private final Map> headers; + + public UriLocation(URI uri, Map> headers) + { + this.uri = requireNonNull(uri, "uri is null"); + this.headers = ImmutableMap.copyOf(requireNonNull(headers, "headers is null")); + } + + public URI uri() + { + return uri; + } + + public Map> headers() + { + return headers; + } +} diff --git a/lib/trino-filesystem/src/main/java/io/trino/filesystem/cache/CacheFileSystem.java b/lib/trino-filesystem/src/main/java/io/trino/filesystem/cache/CacheFileSystem.java new file mode 100644 index 000000000000..cdf338825d48 --- /dev/null +++ b/lib/trino-filesystem/src/main/java/io/trino/filesystem/cache/CacheFileSystem.java @@ -0,0 +1,151 @@ +/* + * 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.filesystem.cache; + +import io.trino.filesystem.FileIterator; +import io.trino.filesystem.Location; +import io.trino.filesystem.TrinoFileSystem; +import io.trino.filesystem.TrinoInputFile; +import io.trino.filesystem.TrinoOutputFile; + +import java.io.IOException; +import java.io.UncheckedIOException; +import java.time.Instant; +import java.util.Collection; +import java.util.Optional; +import java.util.OptionalLong; +import java.util.Set; + +import static java.util.Objects.requireNonNull; + +public final class CacheFileSystem + implements TrinoFileSystem +{ + private final TrinoFileSystem delegate; + private final TrinoFileSystemCache cache; + private final CacheKeyProvider keyProvider; + + public CacheFileSystem(TrinoFileSystem delegate, TrinoFileSystemCache cache, CacheKeyProvider keyProvider) + { + this.delegate = requireNonNull(delegate, "delegate is null"); + this.cache = requireNonNull(cache, "cache is null"); + this.keyProvider = requireNonNull(keyProvider, "keyProvider is null"); + } + + @Override + public TrinoInputFile newInputFile(Location location) + { + return new CacheInputFile(delegate.newInputFile(location), cache, keyProvider, OptionalLong.empty(), Optional.empty()); + } + + @Override + public TrinoInputFile newInputFile(Location location, long length) + { + return new CacheInputFile(delegate.newInputFile(location, length), cache, keyProvider, OptionalLong.of(length), Optional.empty()); + } + + @Override + public TrinoInputFile newInputFile(Location location, long length, Instant lastModified) + { + return new CacheInputFile(delegate.newInputFile(location, length, lastModified), cache, keyProvider, OptionalLong.of(length), Optional.of(lastModified)); + } + + @Override + public TrinoOutputFile newOutputFile(Location location) + { + TrinoOutputFile output = delegate.newOutputFile(location); + try { + cache.expire(location); + } + catch (IOException e) { + throw new UncheckedIOException(e); + } + return output; + } + + @Override + public void deleteFile(Location location) + throws IOException + { + delegate.deleteFile(location); + cache.expire(location); + } + + @Override + public void deleteDirectory(Location location) + throws IOException + { + delegate.deleteDirectory(location); + cache.expire(location); + } + + @Override + public void renameFile(Location source, Location target) + throws IOException + { + delegate.renameFile(source, target); + cache.expire(source); + cache.expire(target); + } + + @Override + public FileIterator listFiles(Location location) + throws IOException + { + return delegate.listFiles(location); + } + + @Override + public Optional directoryExists(Location location) + throws IOException + { + return delegate.directoryExists(location); + } + + @Override + public void createDirectory(Location location) + throws IOException + { + delegate.createDirectory(location); + } + + @Override + public void renameDirectory(Location source, Location target) + throws IOException + { + delegate.renameDirectory(source, target); + } + + @Override + public Set listDirectories(Location location) + throws IOException + { + return delegate.listDirectories(location); + } + + @Override + public Optional createTemporaryDirectory(Location targetPath, String temporaryPrefix, String relativePrefix) + throws IOException + { + return delegate.createTemporaryDirectory(targetPath, temporaryPrefix, relativePrefix); + } + + @Override + public void deleteFiles(Collection locations) + throws IOException + { + delegate.deleteFiles(locations); + cache.expire(locations); + } +} diff --git a/lib/trino-filesystem/src/main/java/io/trino/filesystem/cache/CacheFileSystemFactory.java b/lib/trino-filesystem/src/main/java/io/trino/filesystem/cache/CacheFileSystemFactory.java new file mode 100644 index 000000000000..3c02e0a3c4f3 --- /dev/null +++ b/lib/trino-filesystem/src/main/java/io/trino/filesystem/cache/CacheFileSystemFactory.java @@ -0,0 +1,45 @@ +/* + * 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.filesystem.cache; + +import io.opentelemetry.api.trace.Tracer; +import io.trino.filesystem.TrinoFileSystem; +import io.trino.filesystem.TrinoFileSystemFactory; +import io.trino.filesystem.tracing.TracingFileSystemCache; +import io.trino.spi.security.ConnectorIdentity; + +import static java.util.Objects.requireNonNull; + +public final class CacheFileSystemFactory + implements TrinoFileSystemFactory +{ + private final Tracer tracer; + private final TrinoFileSystemFactory delegate; + private final TrinoFileSystemCache cache; + private final CacheKeyProvider keyProvider; + + public CacheFileSystemFactory(Tracer tracer, TrinoFileSystemFactory delegate, TrinoFileSystemCache cache, CacheKeyProvider keyProvider) + { + this.tracer = requireNonNull(tracer, "tracer is null"); + this.delegate = requireNonNull(delegate, "delegate is null"); + this.cache = requireNonNull(cache, "cache is null"); + this.keyProvider = requireNonNull(keyProvider, "keyProvider is null"); + } + + @Override + public TrinoFileSystem create(ConnectorIdentity identity) + { + return new CacheFileSystem(delegate.create(identity), new TracingFileSystemCache(tracer, cache), keyProvider); + } +} diff --git a/lib/trino-filesystem/src/main/java/io/trino/filesystem/cache/CacheInputFile.java b/lib/trino-filesystem/src/main/java/io/trino/filesystem/cache/CacheInputFile.java new file mode 100644 index 000000000000..3c7f6e1fe3d2 --- /dev/null +++ b/lib/trino-filesystem/src/main/java/io/trino/filesystem/cache/CacheInputFile.java @@ -0,0 +1,106 @@ +/* + * 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.filesystem.cache; + +import io.trino.filesystem.Location; +import io.trino.filesystem.TrinoInput; +import io.trino.filesystem.TrinoInputFile; +import io.trino.filesystem.TrinoInputStream; + +import java.io.IOException; +import java.time.Instant; +import java.util.Optional; +import java.util.OptionalLong; + +import static java.util.Objects.requireNonNull; + +public final class CacheInputFile + implements TrinoInputFile +{ + private final TrinoInputFile delegate; + private final TrinoFileSystemCache cache; + private final CacheKeyProvider keyProvider; + private OptionalLong length; + private Optional lastModified; + + public CacheInputFile(TrinoInputFile delegate, TrinoFileSystemCache cache, CacheKeyProvider keyProvider, OptionalLong length, Optional lastModified) + { + this.delegate = requireNonNull(delegate, "delegate is null"); + this.cache = requireNonNull(cache, "cache is null"); + this.keyProvider = requireNonNull(keyProvider, "keyProvider is null"); + this.length = requireNonNull(length, "length is null"); + this.lastModified = requireNonNull(lastModified, "lastModified is null"); + } + + @Override + public TrinoInput newInput() + throws IOException + { + Optional key = keyProvider.getCacheKey(delegate); + if (key.isPresent()) { + return cache.cacheInput(delegate, key.orElseThrow()); + } + return delegate.newInput(); + } + + @Override + public TrinoInputStream newStream() + throws IOException + { + Optional key = keyProvider.getCacheKey(delegate); + if (key.isPresent()) { + return cache.cacheStream(delegate, key.orElseThrow()); + } + return delegate.newStream(); + } + + @Override + public long length() + throws IOException + { + if (length.isEmpty()) { + Optional key = keyProvider.getCacheKey(delegate); + if (key.isPresent()) { + length = OptionalLong.of(cache.cacheLength(delegate, key.orElseThrow())); + } + else { + length = OptionalLong.of(delegate.length()); + } + } + return length.getAsLong(); + } + + @Override + public Instant lastModified() + throws IOException + { + if (lastModified.isEmpty()) { + lastModified = Optional.of(delegate.lastModified()); + } + return lastModified.orElseThrow(); + } + + @Override + public boolean exists() + throws IOException + { + return delegate.exists(); + } + + @Override + public Location location() + { + return delegate.location(); + } +} diff --git a/lib/trino-filesystem/src/main/java/io/trino/filesystem/cache/CacheKeyProvider.java b/lib/trino-filesystem/src/main/java/io/trino/filesystem/cache/CacheKeyProvider.java new file mode 100644 index 000000000000..f557e938a3ea --- /dev/null +++ b/lib/trino-filesystem/src/main/java/io/trino/filesystem/cache/CacheKeyProvider.java @@ -0,0 +1,28 @@ +/* + * 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.filesystem.cache; + +import io.trino.filesystem.TrinoInputFile; + +import java.io.IOException; +import java.util.Optional; + +public interface CacheKeyProvider +{ + /** + * Get the cache key of a TrinoInputFile. Returns Optional.empty() if the file is not cacheable. + */ + Optional getCacheKey(TrinoInputFile inputFile) + throws IOException; +} diff --git a/lib/trino-filesystem/src/main/java/io/trino/filesystem/cache/CachingHostAddressProvider.java b/lib/trino-filesystem/src/main/java/io/trino/filesystem/cache/CachingHostAddressProvider.java new file mode 100644 index 000000000000..c35c4049999a --- /dev/null +++ b/lib/trino-filesystem/src/main/java/io/trino/filesystem/cache/CachingHostAddressProvider.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.filesystem.cache; + +import io.trino.spi.HostAddress; + +import java.util.List; + +public interface CachingHostAddressProvider +{ + /** + * Returns a lists of hosts which are preferred to cache the split with the given path. + */ + List getHosts(String splitPath, List defaultAddresses); +} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/DisabledGlueColumnStatisticsProviderFactory.java b/lib/trino-filesystem/src/main/java/io/trino/filesystem/cache/DefaultCacheKeyProvider.java similarity index 58% rename from plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/DisabledGlueColumnStatisticsProviderFactory.java rename to lib/trino-filesystem/src/main/java/io/trino/filesystem/cache/DefaultCacheKeyProvider.java index 6a06aa5bc33f..e297691544a2 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/DisabledGlueColumnStatisticsProviderFactory.java +++ b/lib/trino-filesystem/src/main/java/io/trino/filesystem/cache/DefaultCacheKeyProvider.java @@ -11,16 +11,20 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package io.trino.plugin.hive.metastore.glue; +package io.trino.filesystem.cache; -import com.amazonaws.services.glue.AWSGlueAsync; +import io.trino.filesystem.TrinoInputFile; -public class DisabledGlueColumnStatisticsProviderFactory - implements GlueColumnStatisticsProviderFactory +import java.io.IOException; +import java.util.Optional; + +public final class DefaultCacheKeyProvider + implements CacheKeyProvider { @Override - public GlueColumnStatisticsProvider createGlueColumnStatisticsProvider(AWSGlueAsync glueClient, GlueMetastoreStats stats) + public Optional getCacheKey(TrinoInputFile inputFile) + throws IOException { - return new DisabledGlueColumnStatisticsProvider(); + return Optional.of(inputFile.location().path() + "#" + inputFile.lastModified() + "#" + inputFile.length()); } } diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/DefaultHiveMaterializedViewMetadataFactory.java b/lib/trino-filesystem/src/main/java/io/trino/filesystem/cache/DefaultCachingHostAddressProvider.java similarity index 62% rename from plugin/trino-hive/src/main/java/io/trino/plugin/hive/DefaultHiveMaterializedViewMetadataFactory.java rename to lib/trino-filesystem/src/main/java/io/trino/filesystem/cache/DefaultCachingHostAddressProvider.java index 2c732b7ee871..8ed93e70caea 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/DefaultHiveMaterializedViewMetadataFactory.java +++ b/lib/trino-filesystem/src/main/java/io/trino/filesystem/cache/DefaultCachingHostAddressProvider.java @@ -11,16 +11,18 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package io.trino.plugin.hive; +package io.trino.filesystem.cache; -public class DefaultHiveMaterializedViewMetadataFactory - implements HiveMaterializedViewMetadataFactory -{ - private static final HiveMaterializedViewMetadata NONE = new NoneHiveMaterializedViewMetadata(); +import io.trino.spi.HostAddress; + +import java.util.List; +public class DefaultCachingHostAddressProvider + implements CachingHostAddressProvider +{ @Override - public HiveMaterializedViewMetadata create(HiveMetastoreClosure hiveMetastoreClosure) + public List getHosts(String splitPath, List defaultAddresses) { - return NONE; + return defaultAddresses; } } diff --git a/lib/trino-filesystem/src/main/java/io/trino/filesystem/cache/TrinoFileSystemCache.java b/lib/trino-filesystem/src/main/java/io/trino/filesystem/cache/TrinoFileSystemCache.java new file mode 100644 index 000000000000..6f4b3e588788 --- /dev/null +++ b/lib/trino-filesystem/src/main/java/io/trino/filesystem/cache/TrinoFileSystemCache.java @@ -0,0 +1,55 @@ +/* + * 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.filesystem.cache; + +import io.trino.filesystem.Location; +import io.trino.filesystem.TrinoInput; +import io.trino.filesystem.TrinoInputFile; +import io.trino.filesystem.TrinoInputStream; + +import java.io.IOException; +import java.util.Collection; + +public interface TrinoFileSystemCache +{ + /** + * Get the TrinoInput of the TrinoInputFile, potentially using or updating the data cached at key. + */ + TrinoInput cacheInput(TrinoInputFile delegate, String key) + throws IOException; + + /** + * Get the TrinoInputStream of the TrinoInputFile, potentially using or updating the data cached at key. + */ + TrinoInputStream cacheStream(TrinoInputFile delegate, String key) + throws IOException; + + /** + * Get the length of the TrinoInputFile, potentially using or updating the data cached at key. + */ + long cacheLength(TrinoInputFile delegate, String key) + throws IOException; + + /** + * Give a hint to the cache that the cache entry for location should be expired. + */ + void expire(Location location) + throws IOException; + + /** + * Give a hint to the cache that the cache entry for locations should be expired. + */ + void expire(Collection locations) + throws IOException; +} diff --git a/lib/trino-filesystem/src/main/java/io/trino/filesystem/encryption/EncryptionKey.java b/lib/trino-filesystem/src/main/java/io/trino/filesystem/encryption/EncryptionKey.java new file mode 100644 index 000000000000..b56e1c151955 --- /dev/null +++ b/lib/trino-filesystem/src/main/java/io/trino/filesystem/encryption/EncryptionKey.java @@ -0,0 +1,63 @@ +/* + * 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.filesystem.encryption; + +import java.util.Arrays; +import java.util.Objects; +import java.util.concurrent.ThreadLocalRandom; + +import static java.util.Objects.requireNonNull; + +public record EncryptionKey(byte[] key, String algorithm) +{ + public EncryptionKey + { + requireNonNull(algorithm, "algorithm is null"); + requireNonNull(key, "key is null"); + } + + public static EncryptionKey randomAes256() + { + byte[] key = new byte[32]; + ThreadLocalRandom.current().nextBytes(key); + return new EncryptionKey(key, "AES256"); + } + + @Override + public String toString() + { + // We intentionally overwrite toString to hide a key + return algorithm; + } + + @Override + public boolean equals(Object o) + { + if (this == o) { + return true; + } + + if (!(o instanceof EncryptionKey that)) { + return false; + } + return Objects.deepEquals(key, that.key) + && Objects.equals(algorithm, that.algorithm); + } + + @Override + public int hashCode() + { + return Objects.hash(Arrays.hashCode(key), algorithm); + } +} diff --git a/lib/trino-filesystem/src/main/java/io/trino/filesystem/local/LocalFileSystem.java b/lib/trino-filesystem/src/main/java/io/trino/filesystem/local/LocalFileSystem.java index e3e719dc081f..3b22442c2119 100644 --- a/lib/trino-filesystem/src/main/java/io/trino/filesystem/local/LocalFileSystem.java +++ b/lib/trino-filesystem/src/main/java/io/trino/filesystem/local/LocalFileSystem.java @@ -13,6 +13,7 @@ */ package io.trino.filesystem.local; +import com.google.common.collect.ImmutableSet; import io.trino.filesystem.FileIterator; import io.trino.filesystem.Location; import io.trino.filesystem.TrinoFileSystem; @@ -26,9 +27,14 @@ import java.nio.file.SimpleFileVisitor; import java.nio.file.attribute.BasicFileAttributes; import java.util.Optional; +import java.util.Set; +import java.util.stream.Stream; import static com.google.common.base.Preconditions.checkArgument; +import static com.google.common.collect.ImmutableSet.toImmutableSet; import static io.trino.filesystem.local.LocalUtils.handleException; +import static java.nio.file.LinkOption.NOFOLLOW_LINKS; +import static java.util.UUID.randomUUID; /** * A hierarchical file system for testing. @@ -157,6 +163,86 @@ public Optional directoryExists(Location location) return Optional.of(Files.isDirectory(toDirectoryPath(location))); } + @Override + public void createDirectory(Location location) + throws IOException + { + validateLocalLocation(location); + try { + Files.createDirectories(toDirectoryPath(location)); + } + catch (IOException e) { + throw new IOException("Failed to create directory: " + location, e); + } + } + + @Override + public void renameDirectory(Location source, Location target) + throws IOException + { + Path sourcePath = toDirectoryPath(source); + Path targetPath = toDirectoryPath(target); + try { + if (!Files.exists(sourcePath)) { + throw new IOException("Source does not exist: " + source); + } + if (!Files.isDirectory(sourcePath)) { + throw new IOException("Source is not a directory: " + source); + } + + Files.createDirectories(targetPath.getParent()); + + // Do not specify atomic move, as unix overwrites when atomic is enabled + Files.move(sourcePath, targetPath); + } + catch (IOException e) { + throw new IOException("Directory rename from %s to %s failed: %s".formatted(source, target, e.getMessage()), e); + } + } + + @Override + public Set listDirectories(Location location) + throws IOException + { + Path path = toDirectoryPath(location); + if (Files.isRegularFile(path)) { + throw new IOException("Location is a file: " + location); + } + if (!Files.isDirectory(path)) { + return ImmutableSet.of(); + } + try (Stream stream = Files.list(path)) { + return stream + .filter(file -> Files.isDirectory(file, NOFOLLOW_LINKS)) + .map(file -> file.getFileName() + "/") + .map(location::appendPath) + .collect(toImmutableSet()); + } + } + + @Override + public Optional createTemporaryDirectory(Location targetPath, String temporaryPrefix, String relativePrefix) + throws IOException + { + // allow for absolute or relative temporary prefix + Location temporary; + if (temporaryPrefix.startsWith("/")) { + String prefix = temporaryPrefix; + while (prefix.startsWith("/")) { + prefix = prefix.substring(1); + } + temporary = Location.of("local:///").appendPath(prefix); + } + else { + temporary = targetPath.appendPath(temporaryPrefix); + } + + temporary = temporary.appendPath(randomUUID().toString()); + + createDirectory(temporary); + return Optional.of(temporary); + } + private Path toFilePath(Location location) { validateLocalLocation(location); @@ -177,7 +263,7 @@ private Path toDirectoryPath(Location location) private static void validateLocalLocation(Location location) { - checkArgument(location.scheme().equals(Optional.of("local")), "Only 'local' scheme is supported: %s", location); + checkArgument(location.scheme().equals(Optional.of("local")) || location.scheme().equals(Optional.of("file")), "Only 'local' and 'file' scheme is supported: %s", location); checkArgument(location.userInfo().isEmpty(), "Local location cannot contain user info: %s", location); checkArgument(location.host().isEmpty(), "Local location cannot contain a host: %s", location); } diff --git a/lib/trino-filesystem/src/main/java/io/trino/filesystem/local/LocalOutputFile.java b/lib/trino-filesystem/src/main/java/io/trino/filesystem/local/LocalOutputFile.java index 87deb5fd11ff..105aa75f3e60 100644 --- a/lib/trino-filesystem/src/main/java/io/trino/filesystem/local/LocalOutputFile.java +++ b/lib/trino-filesystem/src/main/java/io/trino/filesystem/local/LocalOutputFile.java @@ -58,6 +58,22 @@ public OutputStream create(AggregatedMemoryContext memoryContext) } } + @Override + public void createOrOverwrite(byte[] data) + throws IOException + { + try { + Files.createDirectories(path.getParent()); + OutputStream stream = Files.newOutputStream(path); + try (OutputStream out = new LocalOutputStream(location, stream)) { + out.write(data); + } + } + catch (IOException e) { + throw handleException(location, e); + } + } + @Override public OutputStream createOrOverwrite(AggregatedMemoryContext memoryContext) throws IOException diff --git a/lib/trino-filesystem/src/main/java/io/trino/filesystem/memory/MemoryFileSystem.java b/lib/trino-filesystem/src/main/java/io/trino/filesystem/memory/MemoryFileSystem.java index a9b2aa455c46..7edd544717e6 100644 --- a/lib/trino-filesystem/src/main/java/io/trino/filesystem/memory/MemoryFileSystem.java +++ b/lib/trino-filesystem/src/main/java/io/trino/filesystem/memory/MemoryFileSystem.java @@ -13,6 +13,7 @@ */ package io.trino.filesystem.memory; +import com.google.common.collect.ImmutableSet; import io.airlift.slice.Slice; import io.trino.filesystem.FileEntry; import io.trino.filesystem.FileIterator; @@ -28,6 +29,7 @@ import java.util.Iterator; import java.util.Optional; import java.util.OptionalLong; +import java.util.Set; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentMap; @@ -165,6 +167,47 @@ public Optional directoryExists(Location location) return Optional.empty(); } + @Override + public void createDirectory(Location location) + throws IOException + { + validateMemoryLocation(location); + // memory file system does not have directories + } + + @Override + public void renameDirectory(Location source, Location target) + throws IOException + { + throw new IOException("Memory file system does not support directory renames"); + } + + @Override + public Set listDirectories(Location location) + throws IOException + { + String prefix = toBlobPrefix(location); + ImmutableSet.Builder directories = ImmutableSet.builder(); + for (String key : blobs.keySet()) { + if (key.startsWith(prefix)) { + int index = key.indexOf('/', prefix.length() + 1); + if (index >= 0) { + directories.add(Location.of("memory:///" + key.substring(0, index + 1))); + } + } + } + return directories.build(); + } + + @Override + public Optional createTemporaryDirectory(Location targetPath, String temporaryPrefix, String relativePrefix) + throws IOException + { + validateMemoryLocation(targetPath); + // memory file system does not have directories + return Optional.empty(); + } + private static String toBlobKey(Location location) { validateMemoryLocation(location); diff --git a/lib/trino-filesystem/src/main/java/io/trino/filesystem/memory/MemoryFileSystemCache.java b/lib/trino-filesystem/src/main/java/io/trino/filesystem/memory/MemoryFileSystemCache.java new file mode 100644 index 000000000000..bb03c1569071 --- /dev/null +++ b/lib/trino-filesystem/src/main/java/io/trino/filesystem/memory/MemoryFileSystemCache.java @@ -0,0 +1,205 @@ +/* + * 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.filesystem.memory; + +import com.google.common.annotations.VisibleForTesting; +import com.google.common.cache.Cache; +import com.google.common.cache.Weigher; +import com.google.inject.Inject; +import io.airlift.slice.Slice; +import io.airlift.units.DataSize; +import io.airlift.units.Duration; +import io.trino.cache.EvictableCacheBuilder; +import io.trino.filesystem.Location; +import io.trino.filesystem.TrinoInput; +import io.trino.filesystem.TrinoInputFile; +import io.trino.filesystem.TrinoInputStream; +import io.trino.filesystem.cache.TrinoFileSystemCache; +import org.weakref.jmx.Managed; + +import java.io.FileNotFoundException; +import java.io.IOException; +import java.nio.file.FileAlreadyExistsException; +import java.nio.file.NoSuchFileException; +import java.util.Collection; +import java.util.List; +import java.util.Optional; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicLong; + +import static com.google.common.base.Preconditions.checkArgument; +import static com.google.common.collect.ImmutableList.toImmutableList; +import static io.airlift.slice.SizeOf.estimatedSizeOf; +import static io.airlift.slice.SizeOf.sizeOf; +import static io.airlift.units.DataSize.Unit.GIGABYTE; +import static java.lang.Math.toIntExact; + +public final class MemoryFileSystemCache + implements TrinoFileSystemCache +{ + private final Cache> cache; + private final int maxContentLengthBytes; + private final AtomicLong largeFileSkippedCount = new AtomicLong(); + + @Inject + public MemoryFileSystemCache(MemoryFileSystemCacheConfig config) + { + this(config.getCacheTtl(), config.getMaxSize(), config.getMaxContentLength()); + } + + private MemoryFileSystemCache(Duration expireAfterWrite, DataSize maxSize, DataSize maxContentLength) + { + checkArgument(maxContentLength.compareTo(DataSize.of(1, GIGABYTE)) <= 0, "maxContentLength must be less than or equal to 1GB"); + this.cache = EvictableCacheBuilder.newBuilder() + .maximumWeight(maxSize.toBytes()) + .weigher((Weigher>) (key, value) -> toIntExact(estimatedSizeOf(key) + sizeOf(value, Slice::getRetainedSize))) + .expireAfterWrite(expireAfterWrite.toMillis(), TimeUnit.MILLISECONDS) + .shareNothingWhenDisabled() + .recordStats() + .build(); + this.maxContentLengthBytes = toIntExact(maxContentLength.toBytes()); + } + + @Override + public TrinoInput cacheInput(TrinoInputFile delegate, String key) + throws IOException + { + Optional cachedEntry = getOrLoadFromCache(key, delegate); + if (cachedEntry.isEmpty()) { + largeFileSkippedCount.incrementAndGet(); + return delegate.newInput(); + } + + return new MemoryInput(delegate.location(), cachedEntry.get()); + } + + @Override + public TrinoInputStream cacheStream(TrinoInputFile delegate, String key) + throws IOException + { + Optional cachedEntry = getOrLoadFromCache(key, delegate); + if (cachedEntry.isEmpty()) { + largeFileSkippedCount.incrementAndGet(); + return delegate.newStream(); + } + + return new MemoryInputStream(delegate.location(), cachedEntry.get()); + } + + @Override + public long cacheLength(TrinoInputFile delegate, String key) + throws IOException + { + Optional cachedEntry = getOrLoadFromCache(key, delegate); + if (cachedEntry.isEmpty()) { + largeFileSkippedCount.incrementAndGet(); + return delegate.length(); + } + + return cachedEntry.get().length(); + } + + @Override + public void expire(Location location) + throws IOException + { + List expired = cache.asMap().keySet().stream() + .filter(key -> key.startsWith(location.path())) + .collect(toImmutableList()); + cache.invalidateAll(expired); + } + + @Override + public void expire(Collection locations) + throws IOException + { + List expired = cache.asMap().keySet().stream() + .filter(key -> locations.stream().map(Location::path).anyMatch(key::startsWith)) + .collect(toImmutableList()); + cache.invalidateAll(expired); + } + + @Managed + public void flushCache() + { + cache.invalidateAll(); + } + + @Managed + public long getHitCount() + { + return cache.stats().hitCount(); + } + + @Managed + public long getRequestCount() + { + return cache.stats().requestCount(); + } + + @Managed + public long getLargeFileSkippedCount() + { + return largeFileSkippedCount.get(); + } + + @VisibleForTesting + boolean isCached(String key) + { + Optional cachedEntry = cache.getIfPresent(key); + return cachedEntry != null && cachedEntry.isPresent(); + } + + private Optional getOrLoadFromCache(String key, TrinoInputFile delegate) + throws IOException + { + try { + return cache.get(key, () -> load(delegate)); + } + catch (ExecutionException e) { + throw handleException(delegate.location(), e.getCause()); + } + } + + private Optional load(TrinoInputFile delegate) + throws IOException + { + long fileSize = delegate.length(); + if (fileSize > maxContentLengthBytes) { + return Optional.empty(); + } + try (TrinoInput trinoInput = delegate.newInput()) { + return Optional.of(trinoInput.readTail(toIntExact(fileSize))); + } + } + + private static IOException handleException(Location location, Throwable cause) + throws IOException + { + if (cause instanceof FileNotFoundException || cause instanceof NoSuchFileException) { + throw withCause(new FileNotFoundException(location.toString()), cause); + } + if (cause instanceof FileAlreadyExistsException) { + throw withCause(new FileAlreadyExistsException(location.toString()), cause); + } + throw new IOException(cause.getMessage() + ": " + location, cause); + } + + private static T withCause(T throwable, Throwable cause) + { + throwable.initCause(cause); + return throwable; + } +} diff --git a/lib/trino-filesystem/src/main/java/io/trino/filesystem/memory/MemoryFileSystemCacheConfig.java b/lib/trino-filesystem/src/main/java/io/trino/filesystem/memory/MemoryFileSystemCacheConfig.java new file mode 100644 index 000000000000..4849c716456a --- /dev/null +++ b/lib/trino-filesystem/src/main/java/io/trino/filesystem/memory/MemoryFileSystemCacheConfig.java @@ -0,0 +1,83 @@ +/* + * 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.filesystem.memory; + +import com.google.common.annotations.VisibleForTesting; +import io.airlift.configuration.Config; +import io.airlift.configuration.ConfigDescription; +import io.airlift.units.DataSize; +import io.airlift.units.Duration; +import io.airlift.units.MaxDataSize; +import jakarta.validation.constraints.NotNull; + +import static io.airlift.units.DataSize.Unit.MEGABYTE; +import static java.util.concurrent.TimeUnit.HOURS; + +public class MemoryFileSystemCacheConfig +{ + // Runtime.getRuntime().maxMemory() is not 100% stable and may return slightly different value over JVM lifetime. We use + // constant so default configuration for cache size is stable. + @VisibleForTesting + static final DataSize DEFAULT_CACHE_SIZE = DataSize.succinctBytes(Math.min( + Math.floorDiv(Runtime.getRuntime().maxMemory(), 20L), + DataSize.of(200, MEGABYTE).toBytes())); + + private Duration ttl = new Duration(1, HOURS); + private DataSize maxSize = DEFAULT_CACHE_SIZE; + private DataSize maxContentLength = DataSize.of(8, MEGABYTE); + + @NotNull + public Duration getCacheTtl() + { + return ttl; + } + + @Config("fs.memory-cache.ttl") + @ConfigDescription("Duration to keep files in the cache prior to eviction") + public MemoryFileSystemCacheConfig setCacheTtl(Duration ttl) + { + this.ttl = ttl; + return this; + } + + @NotNull + public DataSize getMaxSize() + { + return maxSize; + } + + @Config("fs.memory-cache.max-size") + @ConfigDescription("Maximum total size of the cache") + public MemoryFileSystemCacheConfig setMaxSize(DataSize maxSize) + { + this.maxSize = maxSize; + return this; + } + + @NotNull + // Avoids humongous allocations with the recommended G1HeapRegionSize of 32MB and prevents a few big files from hogging cache space + @MaxDataSize("15MB") + public DataSize getMaxContentLength() + { + return maxContentLength; + } + + @Config("fs.memory-cache.max-content-length") + @ConfigDescription("Maximum size of file that can be cached") + public MemoryFileSystemCacheConfig setMaxContentLength(DataSize maxContentLength) + { + this.maxContentLength = maxContentLength; + return this; + } +} diff --git a/lib/trino-filesystem/src/main/java/io/trino/filesystem/memory/MemoryFileSystemCacheModule.java b/lib/trino-filesystem/src/main/java/io/trino/filesystem/memory/MemoryFileSystemCacheModule.java new file mode 100644 index 000000000000..aba5bd6f776e --- /dev/null +++ b/lib/trino-filesystem/src/main/java/io/trino/filesystem/memory/MemoryFileSystemCacheModule.java @@ -0,0 +1,40 @@ +/* + * 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.filesystem.memory; + +import com.google.inject.Binder; +import io.airlift.configuration.AbstractConfigurationAwareModule; + +import static com.google.inject.Scopes.SINGLETON; +import static io.airlift.configuration.ConfigBinder.configBinder; + +public class MemoryFileSystemCacheModule + extends AbstractConfigurationAwareModule +{ + private final boolean isCoordinator; + + public MemoryFileSystemCacheModule(boolean isCoordinator) + { + this.isCoordinator = isCoordinator; + } + + @Override + protected void setup(Binder binder) + { + configBinder(binder).bindConfig(MemoryFileSystemCacheConfig.class); + if (isCoordinator) { + binder.bind(MemoryFileSystemCache.class).in(SINGLETON); + } + } +} diff --git a/lib/trino-filesystem/src/main/java/io/trino/filesystem/memory/MemoryFileSystemFactory.java b/lib/trino-filesystem/src/main/java/io/trino/filesystem/memory/MemoryFileSystemFactory.java index 3521be7f6b77..ea582f975e5d 100644 --- a/lib/trino-filesystem/src/main/java/io/trino/filesystem/memory/MemoryFileSystemFactory.java +++ b/lib/trino-filesystem/src/main/java/io/trino/filesystem/memory/MemoryFileSystemFactory.java @@ -23,9 +23,11 @@ public class MemoryFileSystemFactory implements TrinoFileSystemFactory { + private final MemoryFileSystem memoryFileSystem = new MemoryFileSystem(); + @Override public TrinoFileSystem create(ConnectorIdentity identity) { - return new MemoryFileSystem(); + return memoryFileSystem; } } diff --git a/lib/trino-filesystem/src/main/java/io/trino/filesystem/memory/MemoryOutputFile.java b/lib/trino-filesystem/src/main/java/io/trino/filesystem/memory/MemoryOutputFile.java index fd10e1ecfb25..adf4766813b6 100644 --- a/lib/trino-filesystem/src/main/java/io/trino/filesystem/memory/MemoryOutputFile.java +++ b/lib/trino-filesystem/src/main/java/io/trino/filesystem/memory/MemoryOutputFile.java @@ -22,6 +22,7 @@ import java.io.OutputStream; import java.nio.file.FileAlreadyExistsException; +import static io.airlift.slice.Slices.wrappedBuffer; import static java.util.Objects.requireNonNull; class MemoryOutputFile @@ -56,6 +57,20 @@ public OutputStream create(AggregatedMemoryContext memoryContext) return new MemoryOutputStream(location, outputBlob::createBlob); } + @Override + public void createOrOverwrite(byte[] data) + throws IOException + { + outputBlob.overwriteBlob(wrappedBuffer(data)); + } + + @Override + public void createExclusive(byte[] data) + throws IOException + { + outputBlob.createBlob(wrappedBuffer(data)); + } + @Override public OutputStream createOrOverwrite(AggregatedMemoryContext memoryContext) throws IOException diff --git a/lib/trino-filesystem/src/main/java/io/trino/filesystem/s3/S3FileSystemConstants.java b/lib/trino-filesystem/src/main/java/io/trino/filesystem/s3/S3FileSystemConstants.java new file mode 100644 index 000000000000..1155026292a6 --- /dev/null +++ b/lib/trino-filesystem/src/main/java/io/trino/filesystem/s3/S3FileSystemConstants.java @@ -0,0 +1,23 @@ +/* + * 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.filesystem.s3; + +public final class S3FileSystemConstants +{ + public static final String EXTRA_CREDENTIALS_ACCESS_KEY_PROPERTY = "internal$s3_aws_access_key"; + public static final String EXTRA_CREDENTIALS_SECRET_KEY_PROPERTY = "internal$s3_aws_secret_key"; + public static final String EXTRA_CREDENTIALS_SESSION_TOKEN_PROPERTY = "internal$s3_aws_session_token"; + + private S3FileSystemConstants() {} +} diff --git a/lib/trino-filesystem-manager/src/main/java/io/trino/filesystem/manager/SwitchingFileSystem.java b/lib/trino-filesystem/src/main/java/io/trino/filesystem/switching/SwitchingFileSystem.java similarity index 74% rename from lib/trino-filesystem-manager/src/main/java/io/trino/filesystem/manager/SwitchingFileSystem.java rename to lib/trino-filesystem/src/main/java/io/trino/filesystem/switching/SwitchingFileSystem.java index c43010e26eaf..96d2031998a9 100644 --- a/lib/trino-filesystem-manager/src/main/java/io/trino/filesystem/manager/SwitchingFileSystem.java +++ b/lib/trino-filesystem/src/main/java/io/trino/filesystem/switching/SwitchingFileSystem.java @@ -11,9 +11,8 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package io.trino.filesystem.manager; +package io.trino.filesystem.switching; -import com.google.common.collect.ImmutableMap; import io.trino.filesystem.FileIterator; import io.trino.filesystem.Location; import io.trino.filesystem.TrinoFileSystem; @@ -25,8 +24,9 @@ import java.io.IOException; import java.util.Collection; -import java.util.Map; import java.util.Optional; +import java.util.Set; +import java.util.function.Function; import static com.google.common.base.Preconditions.checkArgument; import static java.util.Objects.requireNonNull; @@ -37,20 +37,17 @@ final class SwitchingFileSystem { private final Optional session; private final Optional identity; - private final Optional hdfsFactory; - private final Map factories; + private final Function loader; public SwitchingFileSystem( Optional session, Optional identity, - Optional hdfsFactory, - Map factories) + Function loader) { checkArgument(session.isPresent() != identity.isPresent(), "exactly one of session and identity must be present"); this.session = session; this.identity = identity; - this.hdfsFactory = requireNonNull(hdfsFactory, "hdfsFactory is null"); - this.factories = ImmutableMap.copyOf(requireNonNull(factories, "factories is null")); + this.loader = requireNonNull(loader, "loader is null"); } @Override @@ -82,7 +79,7 @@ public void deleteFile(Location location) public void deleteFiles(Collection locations) throws IOException { - var groups = locations.stream().collect(groupingBy(this::determineFactory)); + var groups = locations.stream().collect(groupingBy(loader)); for (var entry : groups.entrySet()) { createFileSystem(entry.getKey()).deleteFiles(entry.getValue()); } @@ -116,17 +113,37 @@ public Optional directoryExists(Location location) return fileSystem(location).directoryExists(location); } - private TrinoFileSystem fileSystem(Location location) + @Override + public void createDirectory(Location location) + throws IOException + { + fileSystem(location).createDirectory(location); + } + + @Override + public void renameDirectory(Location source, Location target) + throws IOException { - return createFileSystem(determineFactory(location)); + fileSystem(source).renameDirectory(source, target); } - private TrinoFileSystemFactory determineFactory(Location location) + @Override + public Set listDirectories(Location location) + throws IOException + { + return fileSystem(location).listDirectories(location); + } + + @Override + public Optional createTemporaryDirectory(Location targetPath, String temporaryPrefix, String relativePrefix) + throws IOException + { + return fileSystem(targetPath).createTemporaryDirectory(targetPath, temporaryPrefix, relativePrefix); + } + + private TrinoFileSystem fileSystem(Location location) { - return location.scheme() - .map(factories::get) - .or(() -> hdfsFactory) - .orElseThrow(() -> new IllegalArgumentException("No factory for location: " + location)); + return createFileSystem(loader.apply(location)); } private TrinoFileSystem createFileSystem(TrinoFileSystemFactory factory) diff --git a/lib/trino-filesystem-manager/src/main/java/io/trino/filesystem/manager/SwitchingFileSystemFactory.java b/lib/trino-filesystem/src/main/java/io/trino/filesystem/switching/SwitchingFileSystemFactory.java similarity index 66% rename from lib/trino-filesystem-manager/src/main/java/io/trino/filesystem/manager/SwitchingFileSystemFactory.java rename to lib/trino-filesystem/src/main/java/io/trino/filesystem/switching/SwitchingFileSystemFactory.java index 973c0a5baf35..f39ff57dbf2f 100644 --- a/lib/trino-filesystem-manager/src/main/java/io/trino/filesystem/manager/SwitchingFileSystemFactory.java +++ b/lib/trino-filesystem/src/main/java/io/trino/filesystem/switching/SwitchingFileSystemFactory.java @@ -11,40 +11,38 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package io.trino.filesystem.manager; +package io.trino.filesystem.switching; -import com.google.common.collect.ImmutableMap; +import io.trino.filesystem.Location; import io.trino.filesystem.TrinoFileSystem; import io.trino.filesystem.TrinoFileSystemFactory; import io.trino.spi.connector.ConnectorSession; import io.trino.spi.security.ConnectorIdentity; -import java.util.Map; import java.util.Optional; +import java.util.function.Function; import static java.util.Objects.requireNonNull; public class SwitchingFileSystemFactory implements TrinoFileSystemFactory { - private final Optional hdfsFactory; - private final Map factories; + private final Function loader; - public SwitchingFileSystemFactory(Optional hdfsFactory, Map factories) + public SwitchingFileSystemFactory(Function loader) { - this.hdfsFactory = requireNonNull(hdfsFactory, "hdfsFactory is null"); - this.factories = ImmutableMap.copyOf(requireNonNull(factories, "factories is null")); + this.loader = requireNonNull(loader, "loader is null"); } @Override public TrinoFileSystem create(ConnectorSession session) { - return new SwitchingFileSystem(Optional.of(session), Optional.empty(), hdfsFactory, factories); + return new SwitchingFileSystem(Optional.of(session), Optional.empty(), loader); } @Override public TrinoFileSystem create(ConnectorIdentity identity) { - return new SwitchingFileSystem(Optional.empty(), Optional.of(identity), hdfsFactory, factories); + return new SwitchingFileSystem(Optional.empty(), Optional.of(identity), loader); } } diff --git a/lib/trino-filesystem/src/main/java/io/trino/filesystem/tracing/CacheSystemAttributes.java b/lib/trino-filesystem/src/main/java/io/trino/filesystem/tracing/CacheSystemAttributes.java new file mode 100644 index 000000000000..09ee2ce4297f --- /dev/null +++ b/lib/trino-filesystem/src/main/java/io/trino/filesystem/tracing/CacheSystemAttributes.java @@ -0,0 +1,32 @@ +/* + * 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.filesystem.tracing; + +import io.opentelemetry.api.common.AttributeKey; + +import static io.opentelemetry.api.common.AttributeKey.longKey; +import static io.opentelemetry.api.common.AttributeKey.stringKey; + +public class CacheSystemAttributes +{ + private CacheSystemAttributes() {} + + public static final AttributeKey CACHE_KEY = stringKey("trino.cache.key"); + public static final AttributeKey CACHE_FILE_LOCATION = stringKey("trino.cache.file.location"); + public static final AttributeKey CACHE_FILE_LOCATION_COUNT = longKey("trino.cache.file.location_count"); + public static final AttributeKey CACHE_FILE_READ_SIZE = longKey("trino.cache.read_size"); + public static final AttributeKey CACHE_FILE_READ_POSITION = longKey("trino.cache.read_position"); + public static final AttributeKey CACHE_FILE_WRITE_SIZE = longKey("trino.cache.write_size"); + public static final AttributeKey CACHE_FILE_WRITE_POSITION = longKey("trino.cache.write_position"); +} diff --git a/lib/trino-filesystem/src/main/java/io/trino/filesystem/tracing/TracingFileSystem.java b/lib/trino-filesystem/src/main/java/io/trino/filesystem/tracing/TracingFileSystem.java index 96c91dd3cef4..3d9070592be5 100644 --- a/lib/trino-filesystem/src/main/java/io/trino/filesystem/tracing/TracingFileSystem.java +++ b/lib/trino-filesystem/src/main/java/io/trino/filesystem/tracing/TracingFileSystem.java @@ -22,8 +22,10 @@ import io.trino.filesystem.TrinoOutputFile; import java.io.IOException; +import java.time.Instant; import java.util.Collection; import java.util.Optional; +import java.util.Set; import static io.trino.filesystem.tracing.Tracing.withTracing; import static java.util.Objects.requireNonNull; @@ -43,13 +45,13 @@ public TracingFileSystem(Tracer tracer, TrinoFileSystem delegate) @Override public TrinoInputFile newInputFile(Location location) { - return new TracingInputFile(tracer, delegate.newInputFile(location), Optional.empty()); + return new TracingInputFile(tracer, delegate.newInputFile(location), Optional.empty(), Optional.empty()); } @Override public TrinoInputFile newInputFile(Location location, long length) { - return new TracingInputFile(tracer, delegate.newInputFile(location, length), Optional.of(length)); + return new TracingInputFile(tracer, delegate.newInputFile(location, length), Optional.of(length), Optional.empty()); } @Override @@ -58,6 +60,12 @@ public TrinoOutputFile newOutputFile(Location location) return new TracingOutputFile(tracer, delegate.newOutputFile(location)); } + @Override + public TrinoInputFile newInputFile(Location location, long length, Instant instant) + { + return new TracingInputFile(tracer, delegate.newInputFile(location, length), Optional.of(length), Optional.ofNullable(instant)); + } + @Override public void deleteFile(Location location) throws IOException @@ -117,4 +125,44 @@ public Optional directoryExists(Location location) .startSpan(); return withTracing(span, () -> delegate.directoryExists(location)); } + + @Override + public void createDirectory(Location location) + throws IOException + { + Span span = tracer.spanBuilder("FileSystem.createDirectory") + .setAttribute(FileSystemAttributes.FILE_LOCATION, location.toString()) + .startSpan(); + withTracing(span, () -> delegate.createDirectory(location)); + } + + @Override + public void renameDirectory(Location source, Location target) + throws IOException + { + Span span = tracer.spanBuilder("FileSystem.renameDirectory") + .setAttribute(FileSystemAttributes.FILE_LOCATION, source.toString()) + .startSpan(); + withTracing(span, () -> delegate.renameDirectory(source, target)); + } + + @Override + public Set listDirectories(Location location) + throws IOException + { + Span span = tracer.spanBuilder("FileSystem.listDirectories") + .setAttribute(FileSystemAttributes.FILE_LOCATION, location.toString()) + .startSpan(); + return withTracing(span, () -> delegate.listDirectories(location)); + } + + @Override + public Optional createTemporaryDirectory(Location targetPath, String temporaryPrefix, String relativePrefix) + throws IOException + { + Span span = tracer.spanBuilder("FileSystem.createTemporaryDirectory") + .setAttribute(FileSystemAttributes.FILE_LOCATION, targetPath.toString()) + .startSpan(); + return withTracing(span, () -> delegate.createTemporaryDirectory(targetPath, temporaryPrefix, relativePrefix)); + } } diff --git a/lib/trino-filesystem/src/main/java/io/trino/filesystem/tracing/TracingFileSystemCache.java b/lib/trino-filesystem/src/main/java/io/trino/filesystem/tracing/TracingFileSystemCache.java new file mode 100644 index 000000000000..9b61f5307850 --- /dev/null +++ b/lib/trino-filesystem/src/main/java/io/trino/filesystem/tracing/TracingFileSystemCache.java @@ -0,0 +1,102 @@ +/* + * 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.filesystem.tracing; + +import io.opentelemetry.api.trace.Span; +import io.opentelemetry.api.trace.Tracer; +import io.trino.filesystem.Location; +import io.trino.filesystem.TrinoInput; +import io.trino.filesystem.TrinoInputFile; +import io.trino.filesystem.TrinoInputStream; +import io.trino.filesystem.cache.TrinoFileSystemCache; + +import java.io.IOException; +import java.util.Collection; + +import static io.trino.filesystem.tracing.CacheSystemAttributes.CACHE_FILE_LOCATION; +import static io.trino.filesystem.tracing.CacheSystemAttributes.CACHE_FILE_LOCATION_COUNT; +import static io.trino.filesystem.tracing.CacheSystemAttributes.CACHE_KEY; +import static io.trino.filesystem.tracing.Tracing.withTracing; +import static java.util.Objects.requireNonNull; + +public class TracingFileSystemCache + implements TrinoFileSystemCache +{ + private final Tracer tracer; + private final TrinoFileSystemCache delegate; + + public TracingFileSystemCache(Tracer tracer, TrinoFileSystemCache delegate) + { + this.tracer = requireNonNull(tracer, "tracer is null"); + this.delegate = requireNonNull(delegate, "delegate is null"); + } + + @Override + public TrinoInput cacheInput(TrinoInputFile delegate, String key) + throws IOException + { + Span span = tracer.spanBuilder("FileSystemCache.cacheInput") + .setAttribute(CACHE_FILE_LOCATION, delegate.location().toString()) + .setAttribute(CACHE_KEY, key) + .startSpan(); + + return withTracing(span, () -> this.delegate.cacheInput(delegate, key)); + } + + @Override + public TrinoInputStream cacheStream(TrinoInputFile delegate, String key) + throws IOException + { + Span span = tracer.spanBuilder("FileSystemCache.cacheStream") + .setAttribute(CACHE_FILE_LOCATION, delegate.location().toString()) + .setAttribute(CACHE_KEY, key) + .startSpan(); + + return withTracing(span, () -> this.delegate.cacheStream(delegate, key)); + } + + @Override + public long cacheLength(TrinoInputFile delegate, String key) + throws IOException + { + Span span = tracer.spanBuilder("FileSystemCache.cacheLength") + .setAttribute(CACHE_FILE_LOCATION, delegate.location().toString()) + .setAttribute(CACHE_KEY, key) + .startSpan(); + + return withTracing(span, () -> this.delegate.cacheLength(delegate, key)); + } + + @Override + public void expire(Location location) + throws IOException + { + Span span = tracer.spanBuilder("FileSystemCache.expire") + .setAttribute(CACHE_FILE_LOCATION, location.toString()) + .startSpan(); + + withTracing(span, () -> delegate.expire(location)); + } + + @Override + public void expire(Collection locations) + throws IOException + { + Span span = tracer.spanBuilder("FileSystemCache.expire") + .setAttribute(CACHE_FILE_LOCATION_COUNT, (long) locations.size()) + .startSpan(); + + withTracing(span, () -> delegate.expire(locations)); + } +} diff --git a/lib/trino-filesystem/src/main/java/io/trino/filesystem/tracing/TracingInputFile.java b/lib/trino-filesystem/src/main/java/io/trino/filesystem/tracing/TracingInputFile.java index 9e9581082903..26c15cac22a5 100644 --- a/lib/trino-filesystem/src/main/java/io/trino/filesystem/tracing/TracingInputFile.java +++ b/lib/trino-filesystem/src/main/java/io/trino/filesystem/tracing/TracingInputFile.java @@ -34,12 +34,14 @@ final class TracingInputFile private final Tracer tracer; private final TrinoInputFile delegate; private final Optional length; + private boolean isLastModifiedKnown; - public TracingInputFile(Tracer tracer, TrinoInputFile delegate, Optional length) + public TracingInputFile(Tracer tracer, TrinoInputFile delegate, Optional length, Optional lastModified) { this.tracer = requireNonNull(tracer, "tracer is null"); this.delegate = requireNonNull(delegate, "delegate is null"); this.length = requireNonNull(length, "length is null"); + this.isLastModifiedKnown = lastModified.isPresent(); } @Override @@ -83,10 +85,18 @@ public long length() public Instant lastModified() throws IOException { + // skip tracing if lastModified is cached, but delegate anyway + if (isLastModifiedKnown) { + return delegate.lastModified(); + } + Span span = tracer.spanBuilder("InputFile.lastModified") .setAttribute(FileSystemAttributes.FILE_LOCATION, toString()) .startSpan(); - return withTracing(span, delegate::lastModified); + Instant fileLastModified = withTracing(span, delegate::lastModified); + isLastModifiedKnown = true; + + return fileLastModified; } @Override diff --git a/lib/trino-filesystem/src/main/java/io/trino/filesystem/tracing/TracingOutputFile.java b/lib/trino-filesystem/src/main/java/io/trino/filesystem/tracing/TracingOutputFile.java index 2a3941c55266..0676527c3571 100644 --- a/lib/trino-filesystem/src/main/java/io/trino/filesystem/tracing/TracingOutputFile.java +++ b/lib/trino-filesystem/src/main/java/io/trino/filesystem/tracing/TracingOutputFile.java @@ -47,6 +47,26 @@ public OutputStream create() return withTracing(span, () -> delegate.create()); } + @Override + public void createOrOverwrite(byte[] data) + throws IOException + { + Span span = tracer.spanBuilder("OutputFile.createOrOverwrite") + .setAttribute(FileSystemAttributes.FILE_LOCATION, toString()) + .startSpan(); + withTracing(span, () -> delegate.createOrOverwrite(data)); + } + + @Override + public void createExclusive(byte[] data) + throws IOException + { + Span span = tracer.spanBuilder("OutputFile.createExclusive") + .setAttribute(FileSystemAttributes.FILE_LOCATION, toString()) + .startSpan(); + withTracing(span, () -> delegate.createExclusive(data)); + } + @Override public OutputStream createOrOverwrite() throws IOException diff --git a/lib/trino-filesystem/src/test/java/io/trino/filesystem/TestLocations.java b/lib/trino-filesystem/src/test/java/io/trino/filesystem/TestLocations.java index fa339eea5d97..909183b69c50 100644 --- a/lib/trino-filesystem/src/test/java/io/trino/filesystem/TestLocations.java +++ b/lib/trino-filesystem/src/test/java/io/trino/filesystem/TestLocations.java @@ -95,4 +95,17 @@ private static void assertDirectoryLocationEquivalence(String leftLocation, Stri assertThat(areDirectoryLocationsEquivalent(Location.of(rightLocation), Location.of(leftLocation))).as("equivalence of '%s' in relation to '%s'", rightLocation, leftLocation) .isEqualTo(equivalent); } + + @Test + void testIsS3Tables() + { + assertThat(Locations.isS3Tables("s3://e97725d9-dbfb-4334-784sox7edps35ncq16arh546frqa1use2b--table-s3")).isTrue(); + assertThat(Locations.isS3Tables("s3://75fed916-b871-4909-mx9t6iohbseks57q16e5y6nf1c8gguse2b--table-s3")).isTrue(); + + assertThat(Locations.isS3Tables("s3://e97725d9-dbfb-4334-784sox7edps35ncq16arh546frqa1use2b--table-s3/")).isFalse(); + assertThat(Locations.isS3Tables("s3://75fed916-b871-4909-mx9t6iohbseks57q16e5y6nf1c8gguse2b--table-s3/")).isFalse(); + assertThat(Locations.isS3Tables("s3://75fed916-b871-4909/mx9t6iohbseks57q16e5y6nf1c8gguse2b--table-s3")).isFalse(); + assertThat(Locations.isS3Tables("s3://test-bucket")).isFalse(); + assertThat(Locations.isS3Tables("s3://test-bucket/default")).isFalse(); + } } diff --git a/lib/trino-filesystem/src/test/java/io/trino/filesystem/TrackingFileSystemFactory.java b/lib/trino-filesystem/src/test/java/io/trino/filesystem/TrackingFileSystemFactory.java index ea1800f55201..c5ee1a747fbe 100644 --- a/lib/trino-filesystem/src/test/java/io/trino/filesystem/TrackingFileSystemFactory.java +++ b/lib/trino-filesystem/src/test/java/io/trino/filesystem/TrackingFileSystemFactory.java @@ -26,6 +26,7 @@ import java.util.Objects; import java.util.Optional; import java.util.OptionalLong; +import java.util.Set; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.atomic.AtomicInteger; import java.util.function.Consumer; @@ -173,6 +174,34 @@ public Optional directoryExists(Location location) { return delegate.directoryExists(location); } + + @Override + public void createDirectory(Location location) + throws IOException + { + delegate.createDirectory(location); + } + + @Override + public void renameDirectory(Location source, Location target) + throws IOException + { + delegate.renameDirectory(source, target); + } + + @Override + public Set listDirectories(Location location) + throws IOException + { + return delegate.listDirectories(location); + } + + @Override + public Optional createTemporaryDirectory(Location targetPath, String temporaryPrefix, String relativePrefix) + throws IOException + { + return delegate.createTemporaryDirectory(targetPath, temporaryPrefix, relativePrefix); + } } private static class TrackingInputFile @@ -260,6 +289,14 @@ public TrackingOutputFile(TrinoOutputFile delegate, Consumer trac this.tracker = requireNonNull(tracker, "tracker is null"); } + @Override + public void createOrOverwrite(byte[] data) + throws IOException + { + tracker.accept(OUTPUT_FILE_CREATE_OR_OVERWRITE); + delegate.createOrOverwrite(data); + } + @Override public OutputStream create(AggregatedMemoryContext memoryContext) throws IOException diff --git a/lib/trino-hdfs/pom.xml b/lib/trino-hdfs/pom.xml index 11a654d5cd3d..5028acdd1839 100644 --- a/lib/trino-hdfs/pom.xml +++ b/lib/trino-hdfs/pom.xml @@ -76,6 +76,11 @@ failsafe + + io.airlift + bootstrap + + io.airlift concurrent diff --git a/lib/trino-hdfs/src/main/java/io/trino/filesystem/hdfs/HdfsFileIterator.java b/lib/trino-hdfs/src/main/java/io/trino/filesystem/hdfs/HdfsFileIterator.java index b92074af8be0..73f59895edb2 100644 --- a/lib/trino-hdfs/src/main/java/io/trino/filesystem/hdfs/HdfsFileIterator.java +++ b/lib/trino-hdfs/src/main/java/io/trino/filesystem/hdfs/HdfsFileIterator.java @@ -86,6 +86,17 @@ public FileEntry next() blocks.isEmpty() ? Optional.empty() : Optional.of(blocks)); } + static Location listedLocation(Location listingLocation, Path listingPath, Path listedPath) + { + String root = listingPath.toUri().getPath(); + String path = listedPath.toUri().getPath(); + + verify(path.startsWith(root), "iterator path [%s] not a child of listing path [%s] for location [%s]", path, root, listingLocation); + + int index = root.endsWith("/") ? root.length() : root.length() + 1; + return listingLocation.appendPath(path.substring(index)); + } + private static Block toTrinoBlock(BlockLocation location) { try { diff --git a/lib/trino-hdfs/src/main/java/io/trino/filesystem/hdfs/HdfsFileSystem.java b/lib/trino-hdfs/src/main/java/io/trino/filesystem/hdfs/HdfsFileSystem.java index f9a6676de55a..282b8bb7f64e 100644 --- a/lib/trino-hdfs/src/main/java/io/trino/filesystem/hdfs/HdfsFileSystem.java +++ b/lib/trino-hdfs/src/main/java/io/trino/filesystem/hdfs/HdfsFileSystem.java @@ -14,6 +14,7 @@ package io.trino.filesystem.hdfs; import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableSet; import io.airlift.stats.TimeStat; import io.trino.filesystem.FileIterator; import io.trino.filesystem.Location; @@ -27,6 +28,9 @@ import org.apache.hadoop.fs.FileStatus; import org.apache.hadoop.fs.FileSystem; import org.apache.hadoop.fs.Path; +import org.apache.hadoop.fs.permission.FsPermission; +import org.apache.hadoop.fs.viewfs.ViewFileSystem; +import org.apache.hadoop.hdfs.DistributedFileSystem; import java.io.FileNotFoundException; import java.io.IOException; @@ -36,12 +40,18 @@ import java.util.Map; import java.util.Map.Entry; import java.util.Optional; +import java.util.Set; import java.util.UUID; +import java.util.stream.Stream; +import static com.google.common.base.Preconditions.checkArgument; +import static com.google.common.collect.ImmutableSet.toImmutableSet; import static io.trino.filesystem.hdfs.HadoopPaths.hadoopPath; +import static io.trino.filesystem.hdfs.HdfsFileIterator.listedLocation; import static io.trino.hdfs.FileSystemUtils.getRawFileSystem; import static java.lang.String.format; import static java.util.Objects.requireNonNull; +import static java.util.UUID.randomUUID; import static java.util.stream.Collectors.groupingBy; import static java.util.stream.Collectors.mapping; import static java.util.stream.Collectors.toList; @@ -230,6 +240,149 @@ public Optional directoryExists(Location location) }); } + @Override + public void createDirectory(Location location) + throws IOException + { + stats.getCreateDirectoryCalls().newCall(); + Path directory = hadoopPath(location); + FileSystem fileSystem = environment.getFileSystem(context, directory); + + environment.doAs(context.getIdentity(), () -> { + if (!hierarchical(fileSystem, location)) { + return null; + } + Optional permission = environment.getNewDirectoryPermissions(); + try (TimeStat.BlockTimer ignore = stats.getCreateDirectoryCalls().time()) { + if (!fileSystem.mkdirs(directory, permission.orElse(null))) { + throw new IOException("mkdirs failed"); + } + // explicitly set permission since the default umask overrides it on creation + if (permission.isPresent()) { + fileSystem.setPermission(directory, permission.get()); + } + } + catch (IOException e) { + stats.getCreateDirectoryCalls().recordException(e); + throw new IOException("Create directory %s failed: %s".formatted(location, e.getMessage()), e); + } + return null; + }); + } + + @Override + public void renameDirectory(Location source, Location target) + throws IOException + { + stats.getRenameDirectoryCalls().newCall(); + Path sourcePath = hadoopPath(source); + Path targetPath = hadoopPath(target); + FileSystem fileSystem = environment.getFileSystem(context, sourcePath); + + environment.doAs(context.getIdentity(), () -> { + try (TimeStat.BlockTimer ignore = stats.getRenameDirectoryCalls().time()) { + if (!hierarchical(fileSystem, source)) { + throw new IOException("Non-hierarchical file system '%s' does not support directory renames".formatted(fileSystem.getScheme())); + } + if (!fileSystem.getFileStatus(sourcePath).isDirectory()) { + throw new IOException("Source location is not a directory"); + } + if (fileSystem.exists(targetPath)) { + throw new IOException("Target location already exists"); + } + if (!fileSystem.rename(sourcePath, targetPath)) { + throw new IOException("rename failed"); + } + return null; + } + catch (IOException e) { + stats.getRenameDirectoryCalls().recordException(e); + throw new IOException("Directory rename from %s to %s failed: %s".formatted(source, target, e.getMessage()), e); + } + }); + } + + @Override + public Set listDirectories(Location location) + throws IOException + { + stats.getListDirectoriesCalls().newCall(); + Path directory = hadoopPath(location); + FileSystem fileSystem = environment.getFileSystem(context, directory); + return environment.doAs(context.getIdentity(), () -> { + try (TimeStat.BlockTimer ignore = stats.getListDirectoriesCalls().time()) { + FileStatus[] files = fileSystem.listStatus(directory); + if (files.length == 0) { + return ImmutableSet.of(); + } + if (files[0].getPath().equals(directory)) { + throw new IOException("Location is a file, not a directory: " + location); + } + return Stream.of(files) + .filter(FileStatus::isDirectory) + .map(file -> listedLocation(location, directory, file.getPath())) + .map(file -> file.appendSuffix("/")) + .collect(toImmutableSet()); + } + catch (FileNotFoundException e) { + return ImmutableSet.of(); + } + catch (IOException e) { + stats.getListDirectoriesCalls().recordException(e); + throw new IOException("List directories for %s failed: %s".formatted(location, e.getMessage()), e); + } + }); + } + + @Override + public Optional createTemporaryDirectory(Location targetLocation, String temporaryPrefix, String relativePrefix) + throws IOException + { + checkArgument(!relativePrefix.contains("/"), "relativePrefix must not contain slash"); + stats.getCreateTemporaryDirectoryCalls().newCall(); + Path targetPath = hadoopPath(targetLocation); + FileSystem fileSystem = environment.getFileSystem(context, targetPath); + + return environment.doAs(context.getIdentity(), () -> { + try (TimeStat.BlockTimer ignore = stats.getCreateTemporaryDirectoryCalls().time()) { + FileSystem rawFileSystem = getRawFileSystem(fileSystem); + + // use relative temporary directory on ViewFS + String prefix = (rawFileSystem instanceof ViewFileSystem) ? relativePrefix : temporaryPrefix; + + // create a temporary directory on the same file system + Path temporaryRoot = new Path(targetPath, prefix); + Path temporaryPath = new Path(temporaryRoot, randomUUID().toString()); + Location temporaryLocation = Location.of(temporaryPath.toString()); + + if (!hierarchical(fileSystem, temporaryLocation)) { + return Optional.empty(); + } + + // files cannot be moved between encryption zones + if ((rawFileSystem instanceof DistributedFileSystem distributedFileSystem) && + (distributedFileSystem.getEZForPath(targetPath) != null)) { + return Optional.empty(); + } + + Optional permission = environment.getNewDirectoryPermissions(); + if (!fileSystem.mkdirs(temporaryPath, permission.orElse(null))) { + throw new IOException("mkdirs failed for " + temporaryPath); + } + // explicitly set permission since the default umask overrides it on creation + if (permission.isPresent()) { + fileSystem.setPermission(temporaryPath, permission.get()); + } + + return Optional.of(temporaryLocation); + } + catch (IOException e) { + stats.getCreateTemporaryDirectoryCalls().recordException(e); + throw new IOException("Create temporary directory for %s failed: %s".formatted(targetLocation, e.getMessage()), e); + } + }); + } + private boolean hierarchical(FileSystem fileSystem, Location rootLocation) { Boolean knownResult = KNOWN_HIERARCHICAL_FILESYSTEMS.get(fileSystem.getScheme()); diff --git a/lib/trino-hdfs/src/main/java/io/trino/filesystem/hdfs/HdfsFileSystemManager.java b/lib/trino-hdfs/src/main/java/io/trino/filesystem/hdfs/HdfsFileSystemManager.java new file mode 100644 index 000000000000..503b291826e0 --- /dev/null +++ b/lib/trino-hdfs/src/main/java/io/trino/filesystem/hdfs/HdfsFileSystemManager.java @@ -0,0 +1,108 @@ +/* + * 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.filesystem.hdfs; + +import com.google.inject.Injector; +import com.google.inject.Module; +import io.airlift.bootstrap.Bootstrap; +import io.airlift.bootstrap.LifeCycleManager; +import io.opentelemetry.api.OpenTelemetry; +import io.trino.filesystem.TrinoFileSystemFactory; +import io.trino.hdfs.HdfsModule; +import io.trino.hdfs.authentication.HdfsAuthenticationModule; +import io.trino.hdfs.azure.HiveAzureModule; +import io.trino.hdfs.cos.HiveCosModule; +import io.trino.hdfs.gcs.HiveGcsModule; +import io.trino.hdfs.s3.HiveS3Module; +import io.trino.plugin.base.CatalogName; +import io.trino.plugin.base.jmx.ConnectorObjectNameGeneratorModule; +import io.trino.plugin.base.jmx.MBeanServerModule; +import io.trino.spi.NodeManager; +import org.weakref.jmx.guice.MBeanModule; + +import java.util.ArrayList; +import java.util.List; +import java.util.Map; + +import static com.google.common.base.Preconditions.checkState; + +public final class HdfsFileSystemManager +{ + private final Bootstrap bootstrap; + private LifeCycleManager lifecycleManager; + + public HdfsFileSystemManager( + Map config, + boolean azureEnabled, + boolean gcsEnabled, + boolean s3Enabled, + String catalogName, + NodeManager nodeManager, + OpenTelemetry openTelemetry) + { + List modules = new ArrayList<>(); + + modules.add(new MBeanModule()); + modules.add(new MBeanServerModule()); + modules.add(new ConnectorObjectNameGeneratorModule("", "")); + + modules.add(new HdfsFileSystemModule()); + modules.add(new HdfsModule()); + modules.add(new HdfsAuthenticationModule()); + modules.add(new HiveCosModule()); + modules.add(binder -> { + binder.bind(NodeManager.class).toInstance(nodeManager); + binder.bind(OpenTelemetry.class).toInstance(openTelemetry); + binder.bind(CatalogName.class).toInstance(new CatalogName(catalogName)); + }); + + if (azureEnabled) { + modules.add(new HiveAzureModule()); + } + if (gcsEnabled) { + modules.add(new HiveGcsModule()); + } + if (s3Enabled) { + modules.add(new HiveS3Module()); + } + + bootstrap = new Bootstrap(modules) + .doNotInitializeLogging() + .setRequiredConfigurationProperties(Map.of()) + .setOptionalConfigurationProperties(config); + } + + public Map configure() + { + //return bootstrap..configure() + // .stream() + // .collect(toMap(ConfigPropertyMetadata::name, ConfigPropertyMetadata::securitySensitive)); + return Map.of(); + } + + public TrinoFileSystemFactory create() + { + checkState(lifecycleManager == null, "Already created"); + Injector injector = bootstrap.initialize(); + lifecycleManager = injector.getInstance(LifeCycleManager.class); + return injector.getInstance(HdfsFileSystemFactory.class); + } + + public void stop() + { + if (lifecycleManager != null) { + lifecycleManager.stop(); + } + } +} diff --git a/lib/trino-hdfs/src/main/java/io/trino/filesystem/hdfs/HdfsOutputFile.java b/lib/trino-hdfs/src/main/java/io/trino/filesystem/hdfs/HdfsOutputFile.java index 8a187adca0fe..61d32f8d82a7 100644 --- a/lib/trino-hdfs/src/main/java/io/trino/filesystem/hdfs/HdfsOutputFile.java +++ b/lib/trino-hdfs/src/main/java/io/trino/filesystem/hdfs/HdfsOutputFile.java @@ -31,6 +31,7 @@ import static io.trino.filesystem.hdfs.HadoopPaths.hadoopPath; import static io.trino.hdfs.FileSystemUtils.getRawFileSystem; +import static io.trino.memory.context.AggregatedMemoryContext.newSimpleAggregatedMemoryContext; import static java.util.Objects.requireNonNull; class HdfsOutputFile @@ -56,6 +57,24 @@ public OutputStream create(AggregatedMemoryContext memoryContext) return create(false, memoryContext); } + @Override + public void createOrOverwrite(byte[] data) + throws IOException + { + try (OutputStream out = create(true, newSimpleAggregatedMemoryContext())) { + out.write(data); + } + } + + @Override + public void createExclusive(byte[] data) + throws IOException + { + Path file = hadoopPath(location); + FileSystem fileSystem = getRawFileSystem(environment.getFileSystem(context, file)); + throw new UnsupportedOperationException("createExclusive not supported for " + fileSystem); + } + @Override public OutputStream createOrOverwrite(AggregatedMemoryContext memoryContext) throws IOException diff --git a/lib/trino-hdfs/src/main/java/io/trino/hdfs/TrinoHdfsFileSystemStats.java b/lib/trino-hdfs/src/main/java/io/trino/hdfs/TrinoHdfsFileSystemStats.java index ce4dba1a9f70..c5e46aa187e4 100644 --- a/lib/trino-hdfs/src/main/java/io/trino/hdfs/TrinoHdfsFileSystemStats.java +++ b/lib/trino-hdfs/src/main/java/io/trino/hdfs/TrinoHdfsFileSystemStats.java @@ -25,6 +25,10 @@ public final class TrinoHdfsFileSystemStats private final CallStats deleteFileCalls = new CallStats(); private final CallStats deleteDirectoryCalls = new CallStats(); private final CallStats directoryExistsCalls = new CallStats(); + private final CallStats createDirectoryCalls = new CallStats(); + private final CallStats renameDirectoryCalls = new CallStats(); + private final CallStats listDirectoriesCalls = new CallStats(); + private final CallStats createTemporaryDirectoryCalls = new CallStats(); @Managed @Nested @@ -74,4 +78,32 @@ public CallStats getDirectoryExistsCalls() { return directoryExistsCalls; } + + @Managed + @Nested + public CallStats getCreateDirectoryCalls() + { + return createDirectoryCalls; + } + + @Managed + @Nested + public CallStats getRenameDirectoryCalls() + { + return renameDirectoryCalls; + } + + @Managed + @Nested + public CallStats getListDirectoriesCalls() + { + return listDirectoriesCalls; + } + + @Managed + @Nested + public CallStats getCreateTemporaryDirectoryCalls() + { + return createTemporaryDirectoryCalls; + } } diff --git a/lib/trino-hive-formats/pom.xml b/lib/trino-hive-formats/pom.xml index dd8a1d3f88da..a81bc69d79a5 100644 --- a/lib/trino-hive-formats/pom.xml +++ b/lib/trino-hive-formats/pom.xml @@ -95,6 +95,13 @@ avro + + org.gaul modernizer-maven-annotations @@ -194,13 +201,6 @@ test - - org.apache.commons - commons-lang3 - 3.12.0 - test - - org.assertj assertj-core diff --git a/lib/trino-hive-formats/src/main/java/io/trino/hive/formats/line/LineReaderFactory.java b/lib/trino-hive-formats/src/main/java/io/trino/hive/formats/line/LineReaderFactory.java index 38590fabefd6..6ff4ea849b5a 100644 --- a/lib/trino-hive-formats/src/main/java/io/trino/hive/formats/line/LineReaderFactory.java +++ b/lib/trino-hive-formats/src/main/java/io/trino/hive/formats/line/LineReaderFactory.java @@ -16,10 +16,11 @@ import io.trino.filesystem.TrinoInputFile; import java.io.IOException; +import java.util.Set; public interface LineReaderFactory { - String getHiveOutputFormatClassName(); + Set getHiveInputFormatClassNames(); LineBuffer createLineBuffer(); diff --git a/lib/trino-hive-formats/src/main/java/io/trino/hive/formats/line/sequence/SequenceFileReaderFactory.java b/lib/trino-hive-formats/src/main/java/io/trino/hive/formats/line/sequence/SequenceFileReaderFactory.java index 782ae5435fbe..efb5df53a11f 100644 --- a/lib/trino-hive-formats/src/main/java/io/trino/hive/formats/line/sequence/SequenceFileReaderFactory.java +++ b/lib/trino-hive-formats/src/main/java/io/trino/hive/formats/line/sequence/SequenceFileReaderFactory.java @@ -13,6 +13,7 @@ */ package io.trino.hive.formats.line.sequence; +import com.google.common.collect.ImmutableSet; import io.trino.filesystem.TrinoInputFile; import io.trino.hive.formats.line.FooterAwareLineReader; import io.trino.hive.formats.line.LineBuffer; @@ -20,6 +21,7 @@ import io.trino.hive.formats.line.LineReaderFactory; import java.io.IOException; +import java.util.Set; import static com.google.common.base.Preconditions.checkArgument; @@ -36,9 +38,9 @@ public SequenceFileReaderFactory(int initialLineBufferSize, int maxLineLength) } @Override - public String getHiveOutputFormatClassName() + public Set getHiveInputFormatClassNames() { - return "org.apache.hadoop.mapred.SequenceFileInputFormat"; + return ImmutableSet.of("org.apache.hadoop.mapred.SequenceFileInputFormat"); } @Override diff --git a/lib/trino-hive-formats/src/main/java/io/trino/hive/formats/line/text/TextLineReaderFactory.java b/lib/trino-hive-formats/src/main/java/io/trino/hive/formats/line/text/TextLineReaderFactory.java index 28752042f738..0cdc1f4c5311 100644 --- a/lib/trino-hive-formats/src/main/java/io/trino/hive/formats/line/text/TextLineReaderFactory.java +++ b/lib/trino-hive-formats/src/main/java/io/trino/hive/formats/line/text/TextLineReaderFactory.java @@ -13,6 +13,7 @@ */ package io.trino.hive.formats.line.text; +import com.google.common.collect.ImmutableSet; import io.trino.filesystem.TrinoInputFile; import io.trino.hive.formats.compression.Codec; import io.trino.hive.formats.compression.CompressionKind; @@ -25,6 +26,7 @@ import java.io.IOException; import java.io.InputStream; import java.util.Optional; +import java.util.Set; import static com.google.common.base.Preconditions.checkArgument; @@ -43,9 +45,9 @@ public TextLineReaderFactory(int fileBufferSize, int initialLineBufferSize, int } @Override - public String getHiveOutputFormatClassName() + public Set getHiveInputFormatClassNames() { - return "org.apache.hadoop.mapred.TextInputFormat"; + return ImmutableSet.of("org.apache.hadoop.mapred.TextInputFormat"); } @Override diff --git a/lib/trino-orc/src/main/java/io/trino/orc/reader/ReaderUtils.java b/lib/trino-orc/src/main/java/io/trino/orc/reader/ReaderUtils.java index 31c4800cdf69..8f46643c4ae3 100644 --- a/lib/trino-orc/src/main/java/io/trino/orc/reader/ReaderUtils.java +++ b/lib/trino-orc/src/main/java/io/trino/orc/reader/ReaderUtils.java @@ -15,11 +15,14 @@ import io.trino.orc.OrcColumn; import io.trino.orc.OrcCorruptionException; +import io.trino.spi.block.Block; +import io.trino.spi.block.DictionaryBlock; import io.trino.spi.type.Type; import java.util.function.Predicate; import static java.lang.Math.max; +import static java.util.Objects.requireNonNull; final class ReaderUtils { @@ -147,4 +150,40 @@ public static void convertLengthVectorToOffsetVector(int[] vector) currentLength = nextLength; } } + + static Block toNotNullSupressedBlock(int positionCount, boolean[] rowIsNull, Block fieldBlock) + { + requireNonNull(rowIsNull, "rowIsNull is null"); + requireNonNull(fieldBlock, "fieldBlock is null"); + + // find an existing position in the block that is null + int nullIndex = -1; + if (fieldBlock.mayHaveNull()) { + for (int position = 0; position < fieldBlock.getPositionCount(); position++) { + if (fieldBlock.isNull(position)) { + nullIndex = position; + break; + } + } + } + // if there are no null positions, append a null to the end of the block + if (nullIndex == -1) { + nullIndex = fieldBlock.getPositionCount(); + fieldBlock = fieldBlock.copyWithAppendedNull(); + } + + // create a dictionary that maps null positions to the null index + int[] dictionaryIds = new int[positionCount]; + int nullSuppressedPosition = 0; + for (int position = 0; position < positionCount; position++) { + if (rowIsNull[position]) { + dictionaryIds[position] = nullIndex; + } + else { + dictionaryIds[position] = nullSuppressedPosition; + nullSuppressedPosition++; + } + } + return DictionaryBlock.create(positionCount, fieldBlock, dictionaryIds); + } } diff --git a/lib/trino-orc/src/main/java/io/trino/orc/reader/StructColumnReader.java b/lib/trino-orc/src/main/java/io/trino/orc/reader/StructColumnReader.java index 16238b7a0a1b..e193369bd210 100644 --- a/lib/trino-orc/src/main/java/io/trino/orc/reader/StructColumnReader.java +++ b/lib/trino-orc/src/main/java/io/trino/orc/reader/StructColumnReader.java @@ -49,6 +49,7 @@ import static io.airlift.slice.SizeOf.instanceSize; import static io.trino.orc.metadata.Stream.StreamKind.PRESENT; import static io.trino.orc.reader.ColumnReaders.createColumnReader; +import static io.trino.orc.reader.ReaderUtils.toNotNullSupressedBlock; import static io.trino.orc.reader.ReaderUtils.verifyStreamType; import static io.trino.orc.stream.MissingInputStreamSource.missingStreamSource; import static java.util.Locale.ENGLISH; @@ -151,19 +152,19 @@ public Block readBlock() Block[] blocks; if (presentStream == null) { - blocks = getBlocksForType(nextBatchSize); + blocks = getBlocksForType(nextBatchSize, nextBatchSize, null); } else { nullVector = new boolean[nextBatchSize]; int nullValues = presentStream.getUnsetBits(nextBatchSize, nullVector); if (nullValues != nextBatchSize) { - blocks = getBlocksForType(nextBatchSize - nullValues); + blocks = getBlocksForType(nextBatchSize, nextBatchSize - nullValues, nullVector); } else { List typeParameters = type.getTypeParameters(); blocks = new Block[typeParameters.size()]; for (int i = 0; i < typeParameters.size(); i++) { - blocks[i] = typeParameters.get(i).createBlockBuilder(null, 0).build(); + blocks[i] = RunLengthEncodedBlock.create(type.getFields().get(i).getType(), null, nextBatchSize); } } } @@ -234,7 +235,8 @@ public String toString() .toString(); } - private Block[] getBlocksForType(int positionCount) + private Block[] getBlocksForType(int positionCount, int nonNullCount, boolean[] nullVector) + throws IOException { Block[] blocks = new Block[fieldNames.size()]; @@ -243,8 +245,12 @@ private Block[] getBlocksForType(int positionCount) ColumnReader columnReader = structFields.get(fieldName); if (columnReader != null) { - columnReader.prepareNextRead(positionCount); - blocks[i] = blockFactory.createBlock(positionCount, columnReader::readBlock, true); + columnReader.prepareNextRead(nonNullCount); + Block block = columnReader.readBlock(); + if (nullVector != null) { + block = toNotNullSupressedBlock(positionCount, nullVector, block); + } + blocks[i] = block; } else { blocks[i] = RunLengthEncodedBlock.create(type.getFields().get(i).getType(), null, positionCount); diff --git a/lib/trino-parquet/pom.xml b/lib/trino-parquet/pom.xml index ea888e03b3da..2124e945633c 100644 --- a/lib/trino-parquet/pom.xml +++ b/lib/trino-parquet/pom.xml @@ -16,6 +16,17 @@ + + com.fasterxml.jackson.core + jackson-core + + + + com.google.errorprone + error_prone_annotations + true + + com.google.guava guava diff --git a/lib/trino-parquet/src/main/java/io/trino/parquet/BloomFilterStore.java b/lib/trino-parquet/src/main/java/io/trino/parquet/BloomFilterStore.java index 85c3dc616eb4..1afc665c1494 100644 --- a/lib/trino-parquet/src/main/java/io/trino/parquet/BloomFilterStore.java +++ b/lib/trino-parquet/src/main/java/io/trino/parquet/BloomFilterStore.java @@ -16,6 +16,8 @@ import com.google.common.collect.ImmutableMap; import io.airlift.slice.BasicSliceInput; import io.airlift.slice.Slice; +import io.trino.parquet.metadata.BlockMetadata; +import io.trino.parquet.metadata.ColumnChunkMetadata; import io.trino.spi.predicate.Domain; import io.trino.spi.predicate.TupleDomain; import org.apache.parquet.column.ColumnDescriptor; @@ -23,8 +25,6 @@ import org.apache.parquet.column.values.bloomfilter.BloomFilter; import org.apache.parquet.format.BloomFilterHeader; import org.apache.parquet.format.Util; -import org.apache.parquet.hadoop.metadata.BlockMetaData; -import org.apache.parquet.hadoop.metadata.ColumnChunkMetaData; import org.apache.parquet.hadoop.metadata.ColumnPath; import org.apache.parquet.io.ParquetDecodingException; @@ -49,14 +49,14 @@ public class BloomFilterStore private final ParquetDataSource dataSource; private final Map bloomFilterOffsets; - public BloomFilterStore(ParquetDataSource dataSource, BlockMetaData block, Set columnsFiltered) + public BloomFilterStore(ParquetDataSource dataSource, BlockMetadata block, Set columnsFiltered) { this.dataSource = requireNonNull(dataSource, "dataSource is null"); requireNonNull(block, "block is null"); requireNonNull(columnsFiltered, "columnsFiltered is null"); ImmutableMap.Builder bloomFilterOffsetBuilder = ImmutableMap.builder(); - for (ColumnChunkMetaData column : block.getColumns()) { + for (ColumnChunkMetadata column : block.columns()) { ColumnPath path = column.getPath(); if (hasBloomFilter(column) && columnsFiltered.contains(path)) { bloomFilterOffsetBuilder.put(path, column.getBloomFilterOffset()); @@ -98,7 +98,7 @@ public Optional getBloomFilter(ColumnPath columnPath) public static Optional getBloomFilterStore( ParquetDataSource dataSource, - BlockMetaData blockMetadata, + BlockMetadata blockMetadata, TupleDomain parquetTupleDomain, ParquetReaderOptions options) { @@ -106,7 +106,7 @@ public static Optional getBloomFilterStore( return Optional.empty(); } - boolean hasBloomFilter = blockMetadata.getColumns().stream().anyMatch(BloomFilterStore::hasBloomFilter); + boolean hasBloomFilter = blockMetadata.columns().stream().anyMatch(BloomFilterStore::hasBloomFilter); if (!hasBloomFilter) { return Optional.empty(); } @@ -120,7 +120,7 @@ public static Optional getBloomFilterStore( return Optional.of(new BloomFilterStore(dataSource, blockMetadata, columnsFilteredPaths)); } - public static boolean hasBloomFilter(ColumnChunkMetaData columnMetaData) + public static boolean hasBloomFilter(ColumnChunkMetadata columnMetaData) { return columnMetaData.getBloomFilterOffset() > 0; } diff --git a/lib/trino-parquet/src/main/java/io/trino/parquet/Column.java b/lib/trino-parquet/src/main/java/io/trino/parquet/Column.java new file mode 100644 index 000000000000..a6c703cafb90 --- /dev/null +++ b/lib/trino-parquet/src/main/java/io/trino/parquet/Column.java @@ -0,0 +1,25 @@ +/* + * 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.parquet; + +import static java.util.Objects.requireNonNull; + +public record Column(String name, Field field) +{ + public Column + { + requireNonNull(name, "name is null"); + requireNonNull(field, "field is null"); + } +} diff --git a/lib/trino-parquet/src/main/java/io/trino/parquet/ColumnStatisticsValidation.java b/lib/trino-parquet/src/main/java/io/trino/parquet/ColumnStatisticsValidation.java index 71e29b6c1134..d68bc2980d96 100644 --- a/lib/trino-parquet/src/main/java/io/trino/parquet/ColumnStatisticsValidation.java +++ b/lib/trino-parquet/src/main/java/io/trino/parquet/ColumnStatisticsValidation.java @@ -18,7 +18,7 @@ import io.trino.spi.block.Block; import io.trino.spi.block.ColumnarArray; import io.trino.spi.block.ColumnarMap; -import io.trino.spi.block.ColumnarRow; +import io.trino.spi.block.RowBlock; import io.trino.spi.type.ArrayType; import io.trino.spi.type.MapType; import io.trino.spi.type.RowType; @@ -86,13 +86,8 @@ else if (type instanceof MapType) { mergedColumnStatistics = columnStatistics.merge(addMapBlock(columnarMap)); } else if (type instanceof RowType) { - ColumnarRow columnarRow = ColumnarRow.toColumnarRow(block); - ImmutableList.Builder fieldsBuilder = ImmutableList.builder(); - for (int index = 0; index < columnarRow.getFieldCount(); index++) { - fieldsBuilder.add(columnarRow.getField(index)); - } - fields = fieldsBuilder.build(); - mergedColumnStatistics = columnStatistics.merge(addRowBlock(columnarRow)); + fields = RowBlock.getNullSuppressedRowFieldsFromBlock(block); + mergedColumnStatistics = columnStatistics.merge(addRowBlock(block)); } else { throw new TrinoException(NOT_SUPPORTED, format("Unsupported type: %s", type)); @@ -148,7 +143,7 @@ private static ColumnStatistics addArrayBlock(ColumnarArray block) return new ColumnStatistics(nonLeafValuesCount, nonLeafValuesCount); } - private static ColumnStatistics addRowBlock(ColumnarRow block) + private static ColumnStatistics addRowBlock(Block block) { if (!block.mayHaveNull()) { return new ColumnStatistics(0, 0); diff --git a/lib/trino-parquet/src/main/java/io/trino/parquet/ParquetCompressionUtils.java b/lib/trino-parquet/src/main/java/io/trino/parquet/ParquetCompressionUtils.java index d482f627caac..dcc852cdf97d 100644 --- a/lib/trino-parquet/src/main/java/io/trino/parquet/ParquetCompressionUtils.java +++ b/lib/trino-parquet/src/main/java/io/trino/parquet/ParquetCompressionUtils.java @@ -40,6 +40,26 @@ public final class ParquetCompressionUtils private ParquetCompressionUtils() {} + public static Slice decompress(ParquetDataSourceId dataSourceId, CompressionCodec codec, Slice input, int uncompressedSize) + throws IOException + { + requireNonNull(input, "input is null"); + + if (input.length() == 0) { + return EMPTY_SLICE; + } + + return switch (codec) { + case UNCOMPRESSED -> input; + case GZIP -> decompressGzip(input, uncompressedSize); + case SNAPPY -> decompressSnappy(input, uncompressedSize); + case LZO -> decompressLZO(input, uncompressedSize); + case LZ4 -> decompressLz4(input, uncompressedSize); + case ZSTD -> decompressZstd(input, uncompressedSize); + case BROTLI, LZ4_RAW -> throw new ParquetCorruptionException(dataSourceId, "Codec not supported in Parquet: %s", codec); + }; + } + public static Slice decompress(CompressionCodec codec, Slice input, int uncompressedSize) throws IOException { diff --git a/lib/trino-parquet/src/main/java/io/trino/parquet/ParquetMetadataConverter.java b/lib/trino-parquet/src/main/java/io/trino/parquet/ParquetMetadataConverter.java new file mode 100644 index 000000000000..f73ee9ffe14a --- /dev/null +++ b/lib/trino-parquet/src/main/java/io/trino/parquet/ParquetMetadataConverter.java @@ -0,0 +1,575 @@ +/* + * 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.parquet; + +import io.trino.parquet.metadata.IndexReference; +import org.apache.parquet.column.EncodingStats; +import org.apache.parquet.column.statistics.BinaryStatistics; +import org.apache.parquet.format.BoundaryOrder; +import org.apache.parquet.format.BsonType; +import org.apache.parquet.format.ColumnChunk; +import org.apache.parquet.format.ColumnIndex; +import org.apache.parquet.format.ConvertedType; +import org.apache.parquet.format.DateType; +import org.apache.parquet.format.DecimalType; +import org.apache.parquet.format.Encoding; +import org.apache.parquet.format.EnumType; +import org.apache.parquet.format.IntType; +import org.apache.parquet.format.JsonType; +import org.apache.parquet.format.ListType; +import org.apache.parquet.format.LogicalType; +import org.apache.parquet.format.MapType; +import org.apache.parquet.format.MicroSeconds; +import org.apache.parquet.format.MilliSeconds; +import org.apache.parquet.format.NanoSeconds; +import org.apache.parquet.format.NullType; +import org.apache.parquet.format.OffsetIndex; +import org.apache.parquet.format.PageEncodingStats; +import org.apache.parquet.format.PageLocation; +import org.apache.parquet.format.SchemaElement; +import org.apache.parquet.format.Statistics; +import org.apache.parquet.format.StringType; +import org.apache.parquet.format.TimeType; +import org.apache.parquet.format.TimeUnit; +import org.apache.parquet.format.TimestampType; +import org.apache.parquet.format.Type; +import org.apache.parquet.format.UUIDType; +import org.apache.parquet.internal.column.columnindex.BinaryTruncator; +import org.apache.parquet.internal.column.columnindex.ColumnIndexBuilder; +import org.apache.parquet.internal.column.columnindex.OffsetIndexBuilder; +import org.apache.parquet.io.api.Binary; +import org.apache.parquet.schema.ColumnOrder.ColumnOrderName; +import org.apache.parquet.schema.LogicalTypeAnnotation; +import org.apache.parquet.schema.LogicalTypeAnnotation.LogicalTypeAnnotationVisitor; +import org.apache.parquet.schema.PrimitiveType; +import org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName; + +import java.util.Arrays; +import java.util.List; +import java.util.Optional; + +import static org.apache.parquet.CorruptStatistics.shouldIgnoreStatistics; +import static org.apache.parquet.schema.LogicalTypeAnnotation.BsonLogicalTypeAnnotation; +import static org.apache.parquet.schema.LogicalTypeAnnotation.DateLogicalTypeAnnotation; +import static org.apache.parquet.schema.LogicalTypeAnnotation.DecimalLogicalTypeAnnotation; +import static org.apache.parquet.schema.LogicalTypeAnnotation.EnumLogicalTypeAnnotation; +import static org.apache.parquet.schema.LogicalTypeAnnotation.IntLogicalTypeAnnotation; +import static org.apache.parquet.schema.LogicalTypeAnnotation.IntervalLogicalTypeAnnotation; +import static org.apache.parquet.schema.LogicalTypeAnnotation.JsonLogicalTypeAnnotation; +import static org.apache.parquet.schema.LogicalTypeAnnotation.ListLogicalTypeAnnotation; +import static org.apache.parquet.schema.LogicalTypeAnnotation.MapKeyValueTypeAnnotation; +import static org.apache.parquet.schema.LogicalTypeAnnotation.MapLogicalTypeAnnotation; +import static org.apache.parquet.schema.LogicalTypeAnnotation.StringLogicalTypeAnnotation; +import static org.apache.parquet.schema.LogicalTypeAnnotation.TimeLogicalTypeAnnotation; +import static org.apache.parquet.schema.LogicalTypeAnnotation.TimestampLogicalTypeAnnotation; +import static org.apache.parquet.schema.LogicalTypeAnnotation.UUIDLogicalTypeAnnotation; +import static org.apache.parquet.schema.LogicalTypeAnnotation.bsonType; +import static org.apache.parquet.schema.LogicalTypeAnnotation.dateType; +import static org.apache.parquet.schema.LogicalTypeAnnotation.decimalType; +import static org.apache.parquet.schema.LogicalTypeAnnotation.enumType; +import static org.apache.parquet.schema.LogicalTypeAnnotation.float16Type; +import static org.apache.parquet.schema.LogicalTypeAnnotation.intType; +import static org.apache.parquet.schema.LogicalTypeAnnotation.jsonType; +import static org.apache.parquet.schema.LogicalTypeAnnotation.listType; +import static org.apache.parquet.schema.LogicalTypeAnnotation.mapType; +import static org.apache.parquet.schema.LogicalTypeAnnotation.stringType; +import static org.apache.parquet.schema.LogicalTypeAnnotation.timeType; +import static org.apache.parquet.schema.LogicalTypeAnnotation.timestampType; +import static org.apache.parquet.schema.LogicalTypeAnnotation.uuidType; + +// based on org.apache.parquet.format.converter.ParquetMetadataConverter +public final class ParquetMetadataConverter +{ + public static final long MAX_STATS_SIZE = 4096; + + private ParquetMetadataConverter() {} + + public static LogicalTypeAnnotation getLogicalTypeAnnotation(ConvertedType type, SchemaElement element) + { + return switch (type) { + case UTF8 -> stringType(); + case MAP -> mapType(); + case MAP_KEY_VALUE -> MapKeyValueTypeAnnotation.getInstance(); + case LIST -> listType(); + case ENUM -> enumType(); + case DECIMAL -> { + int scale = (element == null) ? 0 : element.scale; + int precision = (element == null) ? 0 : element.precision; + yield decimalType(scale, precision); + } + case DATE -> dateType(); + case TIME_MILLIS -> timeType(true, LogicalTypeAnnotation.TimeUnit.MILLIS); + case TIME_MICROS -> timeType(true, LogicalTypeAnnotation.TimeUnit.MICROS); + case TIMESTAMP_MILLIS -> timestampType(true, LogicalTypeAnnotation.TimeUnit.MILLIS); + case TIMESTAMP_MICROS -> timestampType(true, LogicalTypeAnnotation.TimeUnit.MICROS); + case INTERVAL -> IntervalLogicalTypeAnnotation.getInstance(); + case INT_8 -> intType(8, true); + case INT_16 -> intType(16, true); + case INT_32 -> intType(32, true); + case INT_64 -> intType(64, true); + case UINT_8 -> intType(8, false); + case UINT_16 -> intType(16, false); + case UINT_32 -> intType(32, false); + case UINT_64 -> intType(64, false); + case JSON -> jsonType(); + case BSON -> bsonType(); + }; + } + + public static LogicalTypeAnnotation getLogicalTypeAnnotation(LogicalType type) + { + return switch (type.getSetField()) { + case MAP -> mapType(); + case BSON -> bsonType(); + case DATE -> dateType(); + case ENUM -> enumType(); + case JSON -> jsonType(); + case LIST -> listType(); + case TIME -> { + TimeType time = type.getTIME(); + yield timeType(time.isAdjustedToUTC, convertTimeUnit(time.unit)); + } + case STRING -> stringType(); + case DECIMAL -> { + DecimalType decimal = type.getDECIMAL(); + yield decimalType(decimal.scale, decimal.precision); + } + case INTEGER -> { + IntType integer = type.getINTEGER(); + yield intType(integer.bitWidth, integer.isSigned); + } + case UNKNOWN -> null; + case TIMESTAMP -> { + TimestampType timestamp = type.getTIMESTAMP(); + yield timestampType(timestamp.isAdjustedToUTC, convertTimeUnit(timestamp.unit)); + } + case UUID -> uuidType(); + case FLOAT16 -> float16Type(); + }; + } + + public static LogicalType convertToLogicalType(LogicalTypeAnnotation annotation) + { + return annotation.accept(new LogicalTypeConverterVisitor()).orElse(null); + } + + public static PrimitiveTypeName getPrimitive(Type type) + { + return switch (type) { + case BYTE_ARRAY -> PrimitiveTypeName.BINARY; + case INT64 -> PrimitiveTypeName.INT64; + case INT32 -> PrimitiveTypeName.INT32; + case BOOLEAN -> PrimitiveTypeName.BOOLEAN; + case FLOAT -> PrimitiveTypeName.FLOAT; + case DOUBLE -> PrimitiveTypeName.DOUBLE; + case INT96 -> PrimitiveTypeName.INT96; + case FIXED_LEN_BYTE_ARRAY -> PrimitiveTypeName.FIXED_LEN_BYTE_ARRAY; + }; + } + + public static Encoding getEncoding(org.apache.parquet.column.Encoding encoding) + { + return Encoding.valueOf(encoding.name()); + } + + public static org.apache.parquet.column.Encoding getEncoding(Encoding encoding) + { + return org.apache.parquet.column.Encoding.valueOf(encoding.name()); + } + + public static EncodingStats convertEncodingStats(List stats) + { + if (stats == null) { + return null; + } + + EncodingStats.Builder builder = new EncodingStats.Builder(); + for (PageEncodingStats stat : stats) { + switch (stat.getPage_type()) { + case DATA_PAGE_V2: + builder.withV2Pages(); + // fall through + case DATA_PAGE: + builder.addDataEncoding(getEncoding(stat.getEncoding()), stat.getCount()); + break; + case DICTIONARY_PAGE: + builder.addDictEncoding(getEncoding(stat.getEncoding()), stat.getCount()); + break; + default: + // ignore + } + } + return builder.build(); + } + + public static org.apache.parquet.internal.column.columnindex.ColumnIndex fromParquetColumnIndex(PrimitiveType type, + ColumnIndex parquetColumnIndex) + { + if (!isMinMaxStatsSupported(type)) { + return null; + } + return ColumnIndexBuilder.build(type, + fromParquetBoundaryOrder(parquetColumnIndex.getBoundary_order()), + parquetColumnIndex.getNull_pages(), + parquetColumnIndex.getNull_counts(), + parquetColumnIndex.getMin_values(), + parquetColumnIndex.getMax_values()); + } + + public static org.apache.parquet.internal.column.columnindex.OffsetIndex fromParquetOffsetIndex(OffsetIndex parquetOffsetIndex) + { + OffsetIndexBuilder builder = OffsetIndexBuilder.getBuilder(); + for (PageLocation pageLocation : parquetOffsetIndex.getPage_locations()) { + builder.add(pageLocation.getOffset(), pageLocation.getCompressed_page_size(), pageLocation.getFirst_row_index()); + } + return builder.build(); + } + + public static boolean isMinMaxStatsSupported(PrimitiveType type) + { + return type.columnOrder().getColumnOrderName() == ColumnOrderName.TYPE_DEFINED_ORDER; + } + + public static Statistics toParquetStatistics(org.apache.parquet.column.statistics.Statistics stats, int truncateLength) + { + Statistics formatStats = new Statistics(); + if (!stats.isEmpty() && withinLimit(stats, truncateLength)) { + formatStats.setNull_count(stats.getNumNulls()); + if (stats.hasNonNullValue()) { + byte[] min; + byte[] max; + boolean isMinValueExact = true; + boolean isMaxValueExact = true; + + if (stats instanceof BinaryStatistics && truncateLength != Integer.MAX_VALUE) { + BinaryTruncator truncator = BinaryTruncator.getTruncator(stats.type()); + byte[] originalMin = stats.getMinBytes(); + byte[] originalMax = stats.getMaxBytes(); + min = truncateMin(truncator, truncateLength, originalMin); + max = truncateMax(truncator, truncateLength, originalMax); + isMinValueExact = originalMin.length == min.length; + isMaxValueExact = originalMax.length == max.length; + } + else { + min = stats.getMinBytes(); + max = stats.getMaxBytes(); + } + // Fill the former min-max statistics only if the comparison logic is + // signed so the logic of V1 and V2 stats are the same (which is + // trivially true for equal min-max values) + if (sortOrder(stats.type()) == SortOrder.SIGNED || Arrays.equals(min, max)) { + formatStats.setMin(min); + formatStats.setMax(max); + } + + if (isMinMaxStatsSupported(stats.type()) || Arrays.equals(min, max)) { + formatStats.setMin_value(min); + formatStats.setMax_value(max); + //formatStats.setIs_min_value_exact(isMinValueExact); + //formatStats.setIs_max_value_exact(isMaxValueExact); + } + } + } + return formatStats; + } + + public static org.apache.parquet.column.statistics.Statistics fromParquetStatistics(String createdBy, Statistics statistics, PrimitiveType type) + { + org.apache.parquet.column.statistics.Statistics.Builder statsBuilder = + org.apache.parquet.column.statistics.Statistics.getBuilderForReading(type); + if (statistics != null) { + if (statistics.isSetMin_value() && statistics.isSetMax_value()) { + byte[] min = statistics.min_value.array(); + byte[] max = statistics.max_value.array(); + if (isMinMaxStatsSupported(type) || Arrays.equals(min, max)) { + statsBuilder.withMin(min); + statsBuilder.withMax(max); + } + } + else { + boolean isSet = statistics.isSetMax() && statistics.isSetMin(); + boolean maxEqualsMin = isSet && Arrays.equals(statistics.getMin(), statistics.getMax()); + boolean sortOrdersMatch = SortOrder.SIGNED == sortOrder(type); + if (isSet && !shouldIgnoreStatistics(createdBy, type.getPrimitiveTypeName()) && (sortOrdersMatch || maxEqualsMin)) { + statsBuilder.withMin(statistics.min.array()); + statsBuilder.withMax(statistics.max.array()); + } + } + + if (statistics.isSetNull_count()) { + statsBuilder.withNumNulls(statistics.null_count); + } + } + return statsBuilder.build(); + } + + public static IndexReference toColumnIndexReference(ColumnChunk columnChunk) + { + if (columnChunk.isSetColumn_index_offset() && columnChunk.isSetColumn_index_length()) { + return new IndexReference(columnChunk.getColumn_index_offset(), columnChunk.getColumn_index_length()); + } + return null; + } + + public static IndexReference toOffsetIndexReference(ColumnChunk columnChunk) + { + if (columnChunk.isSetOffset_index_offset() && columnChunk.isSetOffset_index_length()) { + return new IndexReference(columnChunk.getOffset_index_offset(), columnChunk.getOffset_index_length()); + } + return null; + } + + public enum SortOrder + { + SIGNED, + UNSIGNED, + UNKNOWN + } + + private static SortOrder sortOrder(PrimitiveType primitive) + { + LogicalTypeAnnotation annotation = primitive.getLogicalTypeAnnotation(); + if (annotation == null) { + return defaultSortOrder(primitive.getPrimitiveTypeName()); + } + return annotation.accept(new SortOrderVisitor()) + .orElse(defaultSortOrder(primitive.getPrimitiveTypeName())); + } + + private static SortOrder defaultSortOrder(PrimitiveTypeName primitive) + { + return switch (primitive) { + case BOOLEAN, INT32, INT64, FLOAT, DOUBLE -> SortOrder.SIGNED; + case BINARY, FIXED_LEN_BYTE_ARRAY -> SortOrder.UNSIGNED; + default -> SortOrder.UNKNOWN; + }; + } + + private static LogicalTypeAnnotation.TimeUnit convertTimeUnit(TimeUnit unit) + { + return switch (unit.getSetField()) { + case MICROS -> LogicalTypeAnnotation.TimeUnit.MICROS; + case MILLIS -> LogicalTypeAnnotation.TimeUnit.MILLIS; + case NANOS -> LogicalTypeAnnotation.TimeUnit.NANOS; + }; + } + + private static org.apache.parquet.internal.column.columnindex.BoundaryOrder fromParquetBoundaryOrder(BoundaryOrder boundaryOrder) + { + return switch (boundaryOrder) { + case ASCENDING -> org.apache.parquet.internal.column.columnindex.BoundaryOrder.ASCENDING; + case DESCENDING -> org.apache.parquet.internal.column.columnindex.BoundaryOrder.DESCENDING; + case UNORDERED -> org.apache.parquet.internal.column.columnindex.BoundaryOrder.UNORDERED; + }; + } + + private static boolean withinLimit(org.apache.parquet.column.statistics.Statistics stats, int truncateLength) + { + if (stats.isSmallerThan(MAX_STATS_SIZE)) { + return true; + } + + return (stats instanceof BinaryStatistics binaryStats) && + binaryStats.isSmallerThanWithTruncation(MAX_STATS_SIZE, truncateLength); + } + + private static byte[] truncateMin(BinaryTruncator truncator, int truncateLength, byte[] input) + { + return truncator.truncateMin(Binary.fromConstantByteArray(input), truncateLength).getBytes(); + } + + private static byte[] truncateMax(BinaryTruncator truncator, int truncateLength, byte[] input) + { + return truncator.truncateMax(Binary.fromConstantByteArray(input), truncateLength).getBytes(); + } + + private static class LogicalTypeConverterVisitor + implements LogicalTypeAnnotationVisitor + { + @Override + public Optional visit(StringLogicalTypeAnnotation type) + { + return Optional.of(LogicalType.STRING(new StringType())); + } + + @Override + public Optional visit(MapLogicalTypeAnnotation type) + { + return Optional.of(LogicalType.MAP(new MapType())); + } + + @Override + public Optional visit(ListLogicalTypeAnnotation type) + { + return Optional.of(LogicalType.LIST(new ListType())); + } + + @Override + public Optional visit(EnumLogicalTypeAnnotation type) + { + return Optional.of(LogicalType.ENUM(new EnumType())); + } + + @Override + public Optional visit(DecimalLogicalTypeAnnotation type) + { + return Optional.of(LogicalType.DECIMAL(new DecimalType(type.getScale(), type.getPrecision()))); + } + + @Override + public Optional visit(DateLogicalTypeAnnotation type) + { + return Optional.of(LogicalType.DATE(new DateType())); + } + + @Override + public Optional visit(TimeLogicalTypeAnnotation type) + { + return Optional.of(LogicalType.TIME(new TimeType(type.isAdjustedToUTC(), convertUnit(type.getUnit())))); + } + + @Override + public Optional visit(TimestampLogicalTypeAnnotation type) + { + return Optional.of(LogicalType.TIMESTAMP(new TimestampType(type.isAdjustedToUTC(), convertUnit(type.getUnit())))); + } + + @Override + public Optional visit(IntLogicalTypeAnnotation type) + { + return Optional.of(LogicalType.INTEGER(new IntType((byte) type.getBitWidth(), type.isSigned()))); + } + + @Override + public Optional visit(JsonLogicalTypeAnnotation type) + { + return Optional.of(LogicalType.JSON(new JsonType())); + } + + @Override + public Optional visit(BsonLogicalTypeAnnotation type) + { + return Optional.of(LogicalType.BSON(new BsonType())); + } + + @Override + public Optional visit(UUIDLogicalTypeAnnotation type) + { + return Optional.of(LogicalType.UUID(new UUIDType())); + } + + @Override + public Optional visit(IntervalLogicalTypeAnnotation type) + { + return Optional.of(LogicalType.UNKNOWN(new NullType())); + } + + static TimeUnit convertUnit(LogicalTypeAnnotation.TimeUnit unit) + { + return switch (unit) { + case MICROS -> TimeUnit.MICROS(new MicroSeconds()); + case MILLIS -> TimeUnit.MILLIS(new MilliSeconds()); + case NANOS -> TimeUnit.NANOS(new NanoSeconds()); + }; + } + } + + private static class SortOrderVisitor + implements LogicalTypeAnnotationVisitor + { + @Override + public Optional visit(IntLogicalTypeAnnotation intLogicalType) + { + return Optional.of(intLogicalType.isSigned() ? SortOrder.SIGNED : SortOrder.UNSIGNED); + } + + @Override + public Optional visit(IntervalLogicalTypeAnnotation intervalLogicalType) + { + return Optional.of(SortOrder.UNKNOWN); + } + + @Override + public Optional visit(DateLogicalTypeAnnotation dateLogicalType) + { + return Optional.of(SortOrder.SIGNED); + } + + @Override + public Optional visit(EnumLogicalTypeAnnotation enumLogicalType) + { + return Optional.of(SortOrder.UNSIGNED); + } + + @Override + public Optional visit(BsonLogicalTypeAnnotation bsonLogicalType) + { + return Optional.of(SortOrder.UNSIGNED); + } + + @Override + public Optional visit(UUIDLogicalTypeAnnotation uuidLogicalType) + { + return Optional.of(SortOrder.UNSIGNED); + } + + @Override + public Optional visit(JsonLogicalTypeAnnotation jsonLogicalType) + { + return Optional.of(SortOrder.UNSIGNED); + } + + @Override + public Optional visit(StringLogicalTypeAnnotation stringLogicalType) + { + return Optional.of(SortOrder.UNSIGNED); + } + + @Override + public Optional visit(DecimalLogicalTypeAnnotation decimalLogicalType) + { + return Optional.of(SortOrder.UNKNOWN); + } + + @Override + public Optional visit(MapKeyValueTypeAnnotation mapKeyValueLogicalType) + { + return Optional.of(SortOrder.UNKNOWN); + } + + @Override + public Optional visit(MapLogicalTypeAnnotation mapLogicalType) + { + return Optional.of(SortOrder.UNKNOWN); + } + + @Override + public Optional visit(ListLogicalTypeAnnotation listLogicalType) + { + return Optional.of(SortOrder.UNKNOWN); + } + + @Override + public Optional visit(TimeLogicalTypeAnnotation timeLogicalType) + { + return Optional.of(SortOrder.SIGNED); + } + + @Override + public Optional visit(TimestampLogicalTypeAnnotation timestampLogicalType) + { + return Optional.of(SortOrder.SIGNED); + } + } +} diff --git a/lib/trino-parquet/src/main/java/io/trino/parquet/ParquetReaderOptions.java b/lib/trino-parquet/src/main/java/io/trino/parquet/ParquetReaderOptions.java index 56896a6f35f9..e0e7d4418fbb 100644 --- a/lib/trino-parquet/src/main/java/io/trino/parquet/ParquetReaderOptions.java +++ b/lib/trino-parquet/src/main/java/io/trino/parquet/ParquetReaderOptions.java @@ -25,6 +25,7 @@ public class ParquetReaderOptions private static final int DEFAULT_MAX_READ_BLOCK_ROW_COUNT = 8 * 1024; private static final DataSize DEFAULT_MAX_MERGE_DISTANCE = DataSize.of(1, MEGABYTE); private static final DataSize DEFAULT_MAX_BUFFER_SIZE = DataSize.of(8, MEGABYTE); + private static final DataSize DEFAULT_SMALL_FILE_THRESHOLD = DataSize.of(3, MEGABYTE); private final boolean ignoreStatistics; private final DataSize maxReadBlockSize; @@ -32,9 +33,8 @@ public class ParquetReaderOptions private final DataSize maxMergeDistance; private final DataSize maxBufferSize; private final boolean useColumnIndex; - private final boolean useBatchColumnReaders; - private final boolean useBatchNestedColumnReaders; private final boolean useBloomFilter; + private final DataSize smallFileThreshold; public ParquetReaderOptions() { @@ -44,9 +44,8 @@ public ParquetReaderOptions() maxMergeDistance = DEFAULT_MAX_MERGE_DISTANCE; maxBufferSize = DEFAULT_MAX_BUFFER_SIZE; useColumnIndex = true; - useBatchColumnReaders = true; - useBatchNestedColumnReaders = true; useBloomFilter = true; + smallFileThreshold = DEFAULT_SMALL_FILE_THRESHOLD; } private ParquetReaderOptions( @@ -56,9 +55,8 @@ private ParquetReaderOptions( DataSize maxMergeDistance, DataSize maxBufferSize, boolean useColumnIndex, - boolean useBatchColumnReaders, - boolean useBatchNestedColumnReaders, - boolean useBloomFilter) + boolean useBloomFilter, + DataSize smallFileThreshold) { this.ignoreStatistics = ignoreStatistics; this.maxReadBlockSize = requireNonNull(maxReadBlockSize, "maxReadBlockSize is null"); @@ -67,9 +65,8 @@ private ParquetReaderOptions( this.maxMergeDistance = requireNonNull(maxMergeDistance, "maxMergeDistance is null"); this.maxBufferSize = requireNonNull(maxBufferSize, "maxBufferSize is null"); this.useColumnIndex = useColumnIndex; - this.useBatchColumnReaders = useBatchColumnReaders; - this.useBatchNestedColumnReaders = useBatchNestedColumnReaders; this.useBloomFilter = useBloomFilter; + this.smallFileThreshold = requireNonNull(smallFileThreshold, "smallFileThreshold is null"); } public boolean isIgnoreStatistics() @@ -92,16 +89,6 @@ public boolean isUseColumnIndex() return useColumnIndex; } - public boolean useBatchColumnReaders() - { - return useBatchColumnReaders; - } - - public boolean useBatchNestedColumnReaders() - { - return useBatchNestedColumnReaders; - } - public boolean useBloomFilter() { return useBloomFilter; @@ -117,6 +104,11 @@ public int getMaxReadBlockRowCount() return maxReadBlockRowCount; } + public DataSize getSmallFileThreshold() + { + return smallFileThreshold; + } + public ParquetReaderOptions withIgnoreStatistics(boolean ignoreStatistics) { return new ParquetReaderOptions( @@ -126,9 +118,8 @@ public ParquetReaderOptions withIgnoreStatistics(boolean ignoreStatistics) maxMergeDistance, maxBufferSize, useColumnIndex, - useBatchColumnReaders, - useBatchNestedColumnReaders, - useBloomFilter); + useBloomFilter, + smallFileThreshold); } public ParquetReaderOptions withMaxReadBlockSize(DataSize maxReadBlockSize) @@ -140,9 +131,8 @@ public ParquetReaderOptions withMaxReadBlockSize(DataSize maxReadBlockSize) maxMergeDistance, maxBufferSize, useColumnIndex, - useBatchColumnReaders, - useBatchNestedColumnReaders, - useBloomFilter); + useBloomFilter, + smallFileThreshold); } public ParquetReaderOptions withMaxReadBlockRowCount(int maxReadBlockRowCount) @@ -154,9 +144,8 @@ public ParquetReaderOptions withMaxReadBlockRowCount(int maxReadBlockRowCount) maxMergeDistance, maxBufferSize, useColumnIndex, - useBatchColumnReaders, - useBatchNestedColumnReaders, - useBloomFilter); + useBloomFilter, + smallFileThreshold); } public ParquetReaderOptions withMaxMergeDistance(DataSize maxMergeDistance) @@ -168,9 +157,8 @@ public ParquetReaderOptions withMaxMergeDistance(DataSize maxMergeDistance) maxMergeDistance, maxBufferSize, useColumnIndex, - useBatchColumnReaders, - useBatchNestedColumnReaders, - useBloomFilter); + useBloomFilter, + smallFileThreshold); } public ParquetReaderOptions withMaxBufferSize(DataSize maxBufferSize) @@ -182,9 +170,8 @@ public ParquetReaderOptions withMaxBufferSize(DataSize maxBufferSize) maxMergeDistance, maxBufferSize, useColumnIndex, - useBatchColumnReaders, - useBatchNestedColumnReaders, - useBloomFilter); + useBloomFilter, + smallFileThreshold); } public ParquetReaderOptions withUseColumnIndex(boolean useColumnIndex) @@ -196,26 +183,11 @@ public ParquetReaderOptions withUseColumnIndex(boolean useColumnIndex) maxMergeDistance, maxBufferSize, useColumnIndex, - useBatchColumnReaders, - useBatchNestedColumnReaders, - useBloomFilter); + useBloomFilter, + smallFileThreshold); } - public ParquetReaderOptions withBatchColumnReaders(boolean useBatchColumnReaders) - { - return new ParquetReaderOptions( - ignoreStatistics, - maxReadBlockSize, - maxReadBlockRowCount, - maxMergeDistance, - maxBufferSize, - useColumnIndex, - useBatchColumnReaders, - useBatchNestedColumnReaders, - useBloomFilter); - } - - public ParquetReaderOptions withBatchNestedColumnReaders(boolean useBatchNestedColumnReaders) + public ParquetReaderOptions withBloomFilter(boolean useBloomFilter) { return new ParquetReaderOptions( ignoreStatistics, @@ -224,12 +196,11 @@ public ParquetReaderOptions withBatchNestedColumnReaders(boolean useBatchNestedC maxMergeDistance, maxBufferSize, useColumnIndex, - useBatchColumnReaders, - useBatchNestedColumnReaders, - useBloomFilter); + useBloomFilter, + smallFileThreshold); } - public ParquetReaderOptions withBloomFilter(boolean useBloomFilter) + public ParquetReaderOptions withSmallFileThreshold(DataSize smallFileThreshold) { return new ParquetReaderOptions( ignoreStatistics, @@ -238,8 +209,7 @@ public ParquetReaderOptions withBloomFilter(boolean useBloomFilter) maxMergeDistance, maxBufferSize, useColumnIndex, - useBatchColumnReaders, - useBatchNestedColumnReaders, - useBloomFilter); + useBloomFilter, + smallFileThreshold); } } diff --git a/lib/trino-parquet/src/main/java/io/trino/parquet/ParquetReaderUtils.java b/lib/trino-parquet/src/main/java/io/trino/parquet/ParquetReaderUtils.java index 92bcae249737..42844392428c 100644 --- a/lib/trino-parquet/src/main/java/io/trino/parquet/ParquetReaderUtils.java +++ b/lib/trino-parquet/src/main/java/io/trino/parquet/ParquetReaderUtils.java @@ -16,6 +16,7 @@ import com.google.common.collect.ImmutableSet; import com.google.common.collect.Sets; import io.airlift.slice.Slice; +import io.trino.parquet.metadata.ColumnChunkMetadata; import io.trino.parquet.reader.SimpleSliceInputStream; import org.apache.parquet.bytes.ByteBufferInputStream; import org.apache.parquet.column.Encoding; @@ -290,4 +291,25 @@ public static boolean isOnlyDictionaryEncodingPages(ColumnChunkMetaData columnMe return false; } + + @SuppressWarnings("deprecation") + public static boolean isOnlyDictionaryEncodingPages(ColumnChunkMetadata columnMetaData) + { + // Files written with newer versions of Parquet libraries (e.g. parquet-mr 1.9.0) will have EncodingStats available + // Otherwise, fallback to v1 logic + EncodingStats stats = columnMetaData.getEncodingStats(); + if (stats != null) { + return stats.hasDictionaryPages() && !stats.hasNonDictionaryEncodedPages(); + } + + Set encodings = columnMetaData.getEncodings(); + if (encodings.contains(PLAIN_DICTIONARY)) { + // PLAIN_DICTIONARY was present, which means at least one page was + // dictionary-encoded and 1.0 encodings are used + // The only other allowed encodings are RLE and BIT_PACKED which are used for repetition or definition levels + return Sets.difference(encodings, ImmutableSet.of(PLAIN_DICTIONARY, RLE, BIT_PACKED)).isEmpty(); + } + + return false; + } } diff --git a/lib/trino-parquet/src/main/java/io/trino/parquet/ParquetTypeUtils.java b/lib/trino-parquet/src/main/java/io/trino/parquet/ParquetTypeUtils.java index 4797f6526a92..f8d2d84af201 100644 --- a/lib/trino-parquet/src/main/java/io/trino/parquet/ParquetTypeUtils.java +++ b/lib/trino-parquet/src/main/java/io/trino/parquet/ParquetTypeUtils.java @@ -35,13 +35,12 @@ import java.math.BigInteger; import java.util.Arrays; -import java.util.HashMap; import java.util.List; import java.util.Locale; import java.util.Map; import java.util.Optional; -import static com.google.common.base.Preconditions.checkArgument; +import static com.google.common.collect.ImmutableMap.toImmutableMap; import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED; import static java.lang.String.format; import static org.apache.parquet.schema.Type.Repetition.OPTIONAL; @@ -107,50 +106,17 @@ public static ColumnIO getArrayElementColumn(ColumnIO columnIO) public static Map, ColumnDescriptor> getDescriptors(MessageType fileSchema, MessageType requestedSchema) { - Map, ColumnDescriptor> descriptorsByPath = new HashMap<>(); - List columns = getColumns(fileSchema, requestedSchema); - for (String[] paths : fileSchema.getPaths()) { - List columnPath = Arrays.asList(paths); - getDescriptor(columns, columnPath) - .ifPresent(columnDescriptor -> descriptorsByPath.put(columnPath, columnDescriptor)); - } - return descriptorsByPath; - } - - public static Optional getDescriptor(List columns, List path) - { - checkArgument(path.size() >= 1, "Parquet nested path should have at least one component"); - int index = getPathIndex(columns, path); - if (index == -1) { - return Optional.empty(); - } - PrimitiveColumnIO columnIO = columns.get(index); - return Optional.of(columnIO.getColumnDescriptor()); - } - - private static int getPathIndex(List columns, List path) - { - int maxLevel = path.size(); - int index = -1; - for (int columnIndex = 0; columnIndex < columns.size(); columnIndex++) { - ColumnIO[] fields = columns.get(columnIndex).getPath(); - if (fields.length <= maxLevel) { - continue; - } - if (fields[maxLevel].getName().equalsIgnoreCase(path.get(maxLevel - 1))) { - boolean match = true; - for (int level = 0; level < maxLevel - 1; level++) { - if (!fields[level + 1].getName().equalsIgnoreCase(path.get(level))) { - match = false; - } - } - - if (match) { - index = columnIndex; - } - } - } - return index; + // io.trino.parquet.reader.MetadataReader.readFooter performs lower casing of all column names in fileSchema. + // requestedSchema also contains lower cased columns because of being derived from fileSchema. + // io.trino.parquet.ParquetTypeUtils.getParquetTypeByName takes care of case-insensitive matching if needed. + // Therefore, we don't need to repeat case-insensitive matching here. + return getColumns(fileSchema, requestedSchema) + .stream() + .collect(toImmutableMap( + columnIO -> Arrays.asList(columnIO.getFieldPath()), + PrimitiveColumnIO::getColumnDescriptor, + // Same column name may occur more than once when the file is written by case-sensitive tools + (oldValue, ignore) -> oldValue)); } @SuppressWarnings("deprecation") diff --git a/lib/trino-parquet/src/main/java/io/trino/parquet/ParquetWriteValidation.java b/lib/trino-parquet/src/main/java/io/trino/parquet/ParquetWriteValidation.java index 83a712efb19f..daa123da3154 100644 --- a/lib/trino-parquet/src/main/java/io/trino/parquet/ParquetWriteValidation.java +++ b/lib/trino-parquet/src/main/java/io/trino/parquet/ParquetWriteValidation.java @@ -18,6 +18,10 @@ import io.airlift.slice.Slice; import io.airlift.slice.Slices; import io.airlift.slice.XxHash64; +import io.trino.parquet.metadata.ColumnChunkMetadata; +import io.trino.parquet.metadata.IndexReference; +import io.trino.parquet.metadata.PrunedBlockMetadata; +import io.trino.parquet.reader.RowGroupInfo; import io.trino.spi.Page; import io.trino.spi.block.Block; import io.trino.spi.type.Type; @@ -27,10 +31,7 @@ import org.apache.parquet.format.ColumnMetaData; import org.apache.parquet.format.RowGroup; import org.apache.parquet.format.converter.ParquetMetadataConverter; -import org.apache.parquet.hadoop.metadata.BlockMetaData; -import org.apache.parquet.hadoop.metadata.ColumnChunkMetaData; import org.apache.parquet.hadoop.metadata.ColumnPath; -import org.apache.parquet.internal.hadoop.metadata.IndexReference; import org.apache.parquet.schema.MessageType; import org.apache.parquet.schema.PrimitiveType; @@ -50,6 +51,7 @@ import static io.airlift.slice.SizeOf.instanceSize; import static io.airlift.slice.SizeOf.sizeOf; import static io.trino.parquet.ColumnStatisticsValidation.ColumnStatistics; +import static io.trino.parquet.ParquetMetadataConverter.getPrimitive; import static io.trino.parquet.ParquetValidationUtils.validateParquet; import static io.trino.parquet.ParquetWriteValidation.IndexReferenceValidation.fromIndexReference; import static java.util.Objects.requireNonNull; @@ -126,17 +128,17 @@ public void validateColumns(ParquetDataSourceId dataSourceId, MessageType schema } } - public void validateBlocksMetadata(ParquetDataSourceId dataSourceId, List blocksMetaData) + public void validateBlocksMetadata(ParquetDataSourceId dataSourceId, List rowGroupInfos) throws ParquetCorruptionException { validateParquet( - blocksMetaData.size() == rowGroups.size(), + rowGroupInfos.size() == rowGroups.size(), dataSourceId, "Number of row groups %d did not match %d", - blocksMetaData.size(), + rowGroupInfos.size(), rowGroups.size()); - for (int rowGroupIndex = 0; rowGroupIndex < blocksMetaData.size(); rowGroupIndex++) { - BlockMetaData block = blocksMetaData.get(rowGroupIndex); + for (int rowGroupIndex = 0; rowGroupIndex < rowGroupInfos.size(); rowGroupIndex++) { + PrunedBlockMetadata block = rowGroupInfos.get(rowGroupIndex).prunedBlockMetadata(); RowGroup rowGroup = rowGroups.get(rowGroupIndex); validateParquet( block.getRowCount() == rowGroup.getNum_rows(), @@ -146,7 +148,7 @@ public void validateBlocksMetadata(ParquetDataSourceId dataSourceId, List columnChunkMetaData = block.getColumns(); + List columnChunkMetaData = block.getColumns(); validateParquet( columnChunkMetaData.size() == rowGroup.getColumnsSize(), dataSourceId, @@ -156,7 +158,7 @@ public void validateBlocksMetadata(ParquetDataSourceId dataSourceId, List actualColumnStatistics) + public void validateRowGroupStatistics(ParquetDataSourceId dataSourceId, PrunedBlockMetadata blockMetaData, List actualColumnStatistics) throws ParquetCorruptionException { - List columnChunks = blockMetaData.getColumns(); + List columnChunks = blockMetaData.getColumns(); checkArgument( columnChunks.size() == actualColumnStatistics.size(), "Column chunk metadata count %s did not match column fields count %s", @@ -367,7 +369,7 @@ public void validateRowGroupStatistics(ParquetDataSourceId dataSourceId, BlockMe actualColumnStatistics.size()); for (int columnIndex = 0; columnIndex < columnChunks.size(); columnIndex++) { - ColumnChunkMetaData columnMetaData = columnChunks.get(columnIndex); + ColumnChunkMetadata columnMetaData = columnChunks.get(columnIndex); ColumnStatistics columnStatistics = actualColumnStatistics.get(columnIndex); long expectedValuesCount = columnMetaData.getValueCount(); validateParquet( diff --git a/lib/trino-parquet/src/main/java/io/trino/parquet/VariantField.java b/lib/trino-parquet/src/main/java/io/trino/parquet/VariantField.java new file mode 100644 index 000000000000..f9b65baf4ca9 --- /dev/null +++ b/lib/trino-parquet/src/main/java/io/trino/parquet/VariantField.java @@ -0,0 +1,56 @@ +/* + * 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.parquet; + +import io.trino.spi.type.Type; + +import static com.google.common.base.MoreObjects.toStringHelper; +import static java.util.Objects.requireNonNull; + +public class VariantField + extends Field +{ + private final Field value; + private final Field metadata; + + public VariantField(Type type, int repetitionLevel, int definitionLevel, boolean required, Field value, Field metadata) + { + super(type, repetitionLevel, definitionLevel, required); + this.value = requireNonNull(value, "value is null"); + this.metadata = requireNonNull(metadata, "metadata is null"); + } + + public Field getValue() + { + return value; + } + + public Field getMetadata() + { + return metadata; + } + + @Override + public String toString() + { + return toStringHelper(this) + .add("type", getType()) + .add("repetitionLevel", getRepetitionLevel()) + .add("definitionLevel", getDefinitionLevel()) + .add("required", isRequired()) + .add("value", value) + .add("metadata", getMetadata()) + .toString(); + } +} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveMaterializedViewPropertiesProvider.java b/lib/trino-parquet/src/main/java/io/trino/parquet/metadata/BlockMetadata.java similarity index 73% rename from plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveMaterializedViewPropertiesProvider.java rename to lib/trino-parquet/src/main/java/io/trino/parquet/metadata/BlockMetadata.java index 829465375eab..6896bde83a62 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveMaterializedViewPropertiesProvider.java +++ b/lib/trino-parquet/src/main/java/io/trino/parquet/metadata/BlockMetadata.java @@ -11,14 +11,14 @@ * See the License for the specific language governing permissions and * limitations under the License. */ - -package io.trino.plugin.hive; - -import io.trino.spi.session.PropertyMetadata; +package io.trino.parquet.metadata; import java.util.List; -public interface HiveMaterializedViewPropertiesProvider +public record BlockMetadata(long rowCount, List columns) { - List> getMaterializedViewProperties(); + public long getStartingPos() + { + return columns().get(0).getStartingPos(); + } } diff --git a/lib/trino-parquet/src/main/java/io/trino/parquet/metadata/ColumnChunkMetadata.java b/lib/trino-parquet/src/main/java/io/trino/parquet/metadata/ColumnChunkMetadata.java new file mode 100644 index 000000000000..381260829869 --- /dev/null +++ b/lib/trino-parquet/src/main/java/io/trino/parquet/metadata/ColumnChunkMetadata.java @@ -0,0 +1,203 @@ +/* + * 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.parquet.metadata; + +import org.apache.parquet.column.Encoding; +import org.apache.parquet.column.EncodingStats; +import org.apache.parquet.column.statistics.Statistics; +import org.apache.parquet.hadoop.metadata.ColumnPath; +import org.apache.parquet.hadoop.metadata.CompressionCodecName; +import org.apache.parquet.schema.PrimitiveType; +import org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName; + +import java.util.Set; + +public abstract class ColumnChunkMetadata +{ + protected int rowGroupOrdinal = -1; + + public static ColumnChunkMetadata get( + ColumnPath path, + PrimitiveType type, + CompressionCodecName codec, + EncodingStats encodingStats, + Set encodings, + Statistics statistics, + long firstDataPage, + long dictionaryPageOffset, + long valueCount, + long totalSize, + long totalUncompressedSize) + { + if (positiveLongFitsInAnInt(firstDataPage) + && positiveLongFitsInAnInt(dictionaryPageOffset) + && positiveLongFitsInAnInt(valueCount) + && positiveLongFitsInAnInt(totalSize) + && positiveLongFitsInAnInt(totalUncompressedSize)) { + return new IntColumnChunkMetadata( + path, type, codec, + encodingStats, encodings, + statistics, + firstDataPage, + dictionaryPageOffset, + valueCount, + totalSize, + totalUncompressedSize); + } + return new LongColumnChunkMetadata( + path, type, codec, + encodingStats, encodings, + statistics, + firstDataPage, + dictionaryPageOffset, + valueCount, + totalSize, + totalUncompressedSize); + } + + public void setRowGroupOrdinal(int rowGroupOrdinal) + { + this.rowGroupOrdinal = rowGroupOrdinal; + } + + public int getRowGroupOrdinal() + { + return rowGroupOrdinal; + } + + public long getStartingPos() + { + decryptIfNeeded(); + long dictionaryPageOffset = getDictionaryPageOffset(); + long firstDataPageOffset = getFirstDataPageOffset(); + if (dictionaryPageOffset > 0 && dictionaryPageOffset < firstDataPageOffset) { + return dictionaryPageOffset; + } + return firstDataPageOffset; + } + + protected static boolean positiveLongFitsInAnInt(long value) + { + return (value >= 0) && (value + Integer.MIN_VALUE <= Integer.MAX_VALUE); + } + + EncodingStats encodingStats; + + ColumnChunkProperties properties; + + private IndexReference columnIndexReference; + private IndexReference offsetIndexReference; + + private long bloomFilterOffset = -1; + + protected ColumnChunkMetadata(ColumnChunkProperties columnChunkProperties) + { + this(null, columnChunkProperties); + } + + protected ColumnChunkMetadata(EncodingStats encodingStats, ColumnChunkProperties columnChunkProperties) + { + this.encodingStats = encodingStats; + this.properties = columnChunkProperties; + } + + protected void decryptIfNeeded() {} + + public CompressionCodecName getCodec() + { + decryptIfNeeded(); + return properties.codec(); + } + + public ColumnPath getPath() + { + return properties.path(); + } + + public PrimitiveTypeName getType() + { + decryptIfNeeded(); + return properties.type().getPrimitiveTypeName(); + } + + public PrimitiveType getPrimitiveType() + { + decryptIfNeeded(); + return properties.type(); + } + + public abstract long getFirstDataPageOffset(); + + public abstract long getDictionaryPageOffset(); + + public abstract long getValueCount(); + + public abstract long getTotalUncompressedSize(); + + public abstract long getTotalSize(); + + public abstract Statistics getStatistics(); + + public IndexReference getColumnIndexReference() + { + decryptIfNeeded(); + return columnIndexReference; + } + + public void setColumnIndexReference(IndexReference indexReference) + { + this.columnIndexReference = indexReference; + } + + public IndexReference getOffsetIndexReference() + { + decryptIfNeeded(); + return offsetIndexReference; + } + + public void setOffsetIndexReference(IndexReference offsetIndexReference) + { + this.offsetIndexReference = offsetIndexReference; + } + + public void setBloomFilterOffset(long bloomFilterOffset) + { + this.bloomFilterOffset = bloomFilterOffset; + } + + public long getBloomFilterOffset() + { + decryptIfNeeded(); + return bloomFilterOffset; + } + + public Set getEncodings() + { + decryptIfNeeded(); + return properties.encodings(); + } + + public EncodingStats getEncodingStats() + { + decryptIfNeeded(); + return encodingStats; + } + + @Override + public String toString() + { + decryptIfNeeded(); + return "ColumnMetaData{" + properties.toString() + ", " + getFirstDataPageOffset() + "}"; + } +} diff --git a/lib/trino-parquet/src/main/java/io/trino/parquet/metadata/ColumnChunkProperties.java b/lib/trino-parquet/src/main/java/io/trino/parquet/metadata/ColumnChunkProperties.java new file mode 100644 index 000000000000..37e94eed03f5 --- /dev/null +++ b/lib/trino-parquet/src/main/java/io/trino/parquet/metadata/ColumnChunkProperties.java @@ -0,0 +1,23 @@ +/* + * 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.parquet.metadata; + +import org.apache.parquet.column.Encoding; +import org.apache.parquet.hadoop.metadata.ColumnPath; +import org.apache.parquet.hadoop.metadata.CompressionCodecName; +import org.apache.parquet.schema.PrimitiveType; + +import java.util.Set; + +public record ColumnChunkProperties(ColumnPath path, PrimitiveType type, CompressionCodecName codec, Set encodings) {} diff --git a/lib/trino-parquet/src/main/java/io/trino/parquet/metadata/FileMetadata.java b/lib/trino-parquet/src/main/java/io/trino/parquet/metadata/FileMetadata.java new file mode 100644 index 000000000000..3950d1c2dd07 --- /dev/null +++ b/lib/trino-parquet/src/main/java/io/trino/parquet/metadata/FileMetadata.java @@ -0,0 +1,56 @@ +/* + * 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.parquet.metadata; + +import org.apache.parquet.schema.MessageType; + +import java.util.Map; + +import static java.util.Collections.unmodifiableMap; +import static java.util.Objects.requireNonNull; + +public final class FileMetadata +{ + private final MessageType schema; + private final Map keyValueMetaData; + private final String createdBy; + + public FileMetadata(MessageType schema, Map keyValueMetaData, String createdBy) + { + this.schema = requireNonNull(schema, "schema cannot be null"); + this.keyValueMetaData = unmodifiableMap(requireNonNull(keyValueMetaData, "keyValueMetaData cannot be null")); + this.createdBy = createdBy; + } + + public MessageType getSchema() + { + return schema; + } + + @Override + public String toString() + { + return "FileMetaData{schema: " + schema + ", metadata: " + keyValueMetaData + "}"; + } + + public Map getKeyValueMetaData() + { + return keyValueMetaData; + } + + public String getCreatedBy() + { + return createdBy; + } +} diff --git a/lib/trino-parquet/src/main/java/io/trino/parquet/metadata/IndexReference.java b/lib/trino-parquet/src/main/java/io/trino/parquet/metadata/IndexReference.java new file mode 100644 index 000000000000..a2918fc03171 --- /dev/null +++ b/lib/trino-parquet/src/main/java/io/trino/parquet/metadata/IndexReference.java @@ -0,0 +1,36 @@ +/* + * 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.parquet.metadata; + +public class IndexReference +{ + private final long offset; + private final int length; + + public IndexReference(long offset, int length) + { + this.offset = offset; + this.length = length; + } + + public long getOffset() + { + return offset; + } + + public int getLength() + { + return length; + } +} diff --git a/lib/trino-parquet/src/main/java/io/trino/parquet/metadata/IntColumnChunkMetadata.java b/lib/trino-parquet/src/main/java/io/trino/parquet/metadata/IntColumnChunkMetadata.java new file mode 100644 index 000000000000..046d34c174bc --- /dev/null +++ b/lib/trino-parquet/src/main/java/io/trino/parquet/metadata/IntColumnChunkMetadata.java @@ -0,0 +1,105 @@ +/* + * 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.parquet.metadata; + +import org.apache.parquet.column.Encoding; +import org.apache.parquet.column.EncodingStats; +import org.apache.parquet.column.statistics.Statistics; +import org.apache.parquet.hadoop.metadata.ColumnPath; +import org.apache.parquet.hadoop.metadata.CompressionCodecName; +import org.apache.parquet.schema.PrimitiveType; + +import java.util.Set; + +class IntColumnChunkMetadata + extends ColumnChunkMetadata +{ + private final int firstDataPage; + private final int dictionaryPageOffset; + private final int valueCount; + private final int totalSize; + private final int totalUncompressedSize; + private final Statistics statistics; + + IntColumnChunkMetadata( + ColumnPath path, + PrimitiveType type, + CompressionCodecName codec, + EncodingStats encodingStats, + Set encodings, + Statistics statistics, + long firstDataPage, + long dictionaryPageOffset, + long valueCount, + long totalSize, + long totalUncompressedSize) + { + super(encodingStats, new ColumnChunkProperties(path, type, codec, encodings)); + this.firstDataPage = positiveLongToInt(firstDataPage); + this.dictionaryPageOffset = positiveLongToInt(dictionaryPageOffset); + this.valueCount = positiveLongToInt(valueCount); + this.totalSize = positiveLongToInt(totalSize); + this.totalUncompressedSize = positiveLongToInt(totalUncompressedSize); + this.statistics = statistics; + } + + private int positiveLongToInt(long value) + { + if (!ColumnChunkMetadata.positiveLongFitsInAnInt(value)) { + throw new IllegalArgumentException("value should be positive and fit in an int: " + value); + } + return (int) (value + Integer.MIN_VALUE); + } + + private long intToPositiveLong(int value) + { + return (long) value - Integer.MIN_VALUE; + } + + @Override + public long getFirstDataPageOffset() + { + return intToPositiveLong(firstDataPage); + } + + @Override + public long getDictionaryPageOffset() + { + return intToPositiveLong(dictionaryPageOffset); + } + + @Override + public long getValueCount() + { + return intToPositiveLong(valueCount); + } + + @Override + public long getTotalUncompressedSize() + { + return intToPositiveLong(totalUncompressedSize); + } + + @Override + public long getTotalSize() + { + return intToPositiveLong(totalSize); + } + + @Override + public Statistics getStatistics() + { + return statistics; + } +} diff --git a/lib/trino-parquet/src/main/java/io/trino/parquet/metadata/LongColumnChunkMetadata.java b/lib/trino-parquet/src/main/java/io/trino/parquet/metadata/LongColumnChunkMetadata.java new file mode 100644 index 000000000000..d0ca0089b3ca --- /dev/null +++ b/lib/trino-parquet/src/main/java/io/trino/parquet/metadata/LongColumnChunkMetadata.java @@ -0,0 +1,92 @@ +/* + * 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.parquet.metadata; + +import org.apache.parquet.column.Encoding; +import org.apache.parquet.column.EncodingStats; +import org.apache.parquet.column.statistics.Statistics; +import org.apache.parquet.hadoop.metadata.ColumnPath; +import org.apache.parquet.hadoop.metadata.CompressionCodecName; +import org.apache.parquet.schema.PrimitiveType; + +import java.util.Set; + +class LongColumnChunkMetadata + extends ColumnChunkMetadata +{ + private final long firstDataPageOffset; + private final long dictionaryPageOffset; + private final long valueCount; + private final long totalSize; + private final long totalUncompressedSize; + private final Statistics statistics; + + LongColumnChunkMetadata( + ColumnPath path, + PrimitiveType type, + CompressionCodecName codec, + EncodingStats encodingStats, + Set encodings, + Statistics statistics, + long firstDataPageOffset, + long dictionaryPageOffset, + long valueCount, + long totalSize, + long totalUncompressedSize) + { + super(encodingStats, new ColumnChunkProperties(path, type, codec, encodings)); + this.firstDataPageOffset = firstDataPageOffset; + this.dictionaryPageOffset = dictionaryPageOffset; + this.valueCount = valueCount; + this.totalSize = totalSize; + this.totalUncompressedSize = totalUncompressedSize; + this.statistics = statistics; + } + + @Override + public long getFirstDataPageOffset() + { + return firstDataPageOffset; + } + + @Override + public long getDictionaryPageOffset() + { + return dictionaryPageOffset; + } + + @Override + public long getValueCount() + { + return valueCount; + } + + @Override + public long getTotalUncompressedSize() + { + return totalUncompressedSize; + } + + @Override + public long getTotalSize() + { + return totalSize; + } + + @Override + public Statistics getStatistics() + { + return statistics; + } +} diff --git a/lib/trino-parquet/src/main/java/io/trino/parquet/metadata/ParquetMetadata.java b/lib/trino-parquet/src/main/java/io/trino/parquet/metadata/ParquetMetadata.java new file mode 100644 index 000000000000..ff00292b1b7e --- /dev/null +++ b/lib/trino-parquet/src/main/java/io/trino/parquet/metadata/ParquetMetadata.java @@ -0,0 +1,264 @@ +/* + * 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.parquet.metadata; + +import com.google.common.annotations.VisibleForTesting; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import io.airlift.log.Logger; +import io.trino.parquet.ParquetCorruptionException; +import io.trino.parquet.ParquetDataSourceId; +import io.trino.parquet.reader.MetadataReader; +import org.apache.parquet.column.Encoding; +import org.apache.parquet.format.ColumnChunk; +import org.apache.parquet.format.ColumnMetaData; +import org.apache.parquet.format.FileMetaData; +import org.apache.parquet.format.KeyValue; +import org.apache.parquet.format.RowGroup; +import org.apache.parquet.format.SchemaElement; +import org.apache.parquet.hadoop.metadata.ColumnPath; +import org.apache.parquet.hadoop.metadata.CompressionCodecName; +import org.apache.parquet.schema.LogicalTypeAnnotation; +import org.apache.parquet.schema.MessageType; +import org.apache.parquet.schema.PrimitiveType; +import org.apache.parquet.schema.Type; +import org.apache.parquet.schema.Types; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashSet; +import java.util.Iterator; +import java.util.List; +import java.util.Locale; +import java.util.Map; +import java.util.Optional; +import java.util.Set; + +import static com.google.common.base.MoreObjects.toStringHelper; +import static com.google.common.collect.ImmutableMap.toImmutableMap; +import static io.trino.parquet.ParquetMetadataConverter.convertEncodingStats; +import static io.trino.parquet.ParquetMetadataConverter.getEncoding; +import static io.trino.parquet.ParquetMetadataConverter.getLogicalTypeAnnotation; +import static io.trino.parquet.ParquetMetadataConverter.getPrimitive; +import static io.trino.parquet.ParquetMetadataConverter.toColumnIndexReference; +import static io.trino.parquet.ParquetMetadataConverter.toOffsetIndexReference; +import static io.trino.parquet.ParquetValidationUtils.validateParquet; +import static java.util.Objects.requireNonNull; + +public class ParquetMetadata +{ + private static final Logger log = Logger.get(ParquetMetadata.class); + + private final FileMetaData parquetMetadata; + private final ParquetDataSourceId dataSourceId; + private final FileMetadata fileMetadata; + + public ParquetMetadata(FileMetaData parquetMetadata, ParquetDataSourceId dataSourceId) + throws ParquetCorruptionException + { + this.fileMetadata = new FileMetadata( + readMessageType(parquetMetadata, dataSourceId), + keyValueMetaData(parquetMetadata), + parquetMetadata.getCreated_by()); + this.parquetMetadata = parquetMetadata; + this.dataSourceId = requireNonNull(dataSourceId, "dataSourceId is null"); + } + + public FileMetadata getFileMetaData() + { + return fileMetadata; + } + + @Override + public String toString() + { + return toStringHelper(this) + .add("parquetMetadata", parquetMetadata) + .toString(); + } + + public List getBlocks() + throws ParquetCorruptionException + { + return getBlocks(0, Long.MAX_VALUE); + } + + public List getBlocks(long splitStart, long splitLength) + throws ParquetCorruptionException + { + List schema = parquetMetadata.getSchema(); + validateParquet(!schema.isEmpty(), dataSourceId, "Schema is empty"); + + MessageType messageType = readParquetSchema(schema); + List blocks = new ArrayList<>(); + List rowGroups = parquetMetadata.getRow_groups(); + if (rowGroups != null) { + for (RowGroup rowGroup : rowGroups) { + if (rowGroup.isSetFile_offset()) { + long rowGroupStart = rowGroup.getFile_offset(); + boolean splitContainsRowGroup = splitStart <= rowGroupStart && rowGroupStart < splitStart + splitLength; + if (!splitContainsRowGroup) { + continue; + } + } + + List columns = rowGroup.getColumns(); + validateParquet(!columns.isEmpty(), dataSourceId, "No columns in row group: %s", rowGroup); + String filePath = columns.get(0).getFile_path(); + ImmutableList.Builder columnMetadataBuilder = ImmutableList.builderWithExpectedSize(columns.size()); + for (ColumnChunk columnChunk : columns) { + validateParquet( + (filePath == null && columnChunk.getFile_path() == null) + || (filePath != null && filePath.equals(columnChunk.getFile_path())), + dataSourceId, + "all column chunks of the same row group must be in the same file"); + ColumnMetaData metaData = columnChunk.meta_data; + String[] path = metaData.path_in_schema.stream() + .map(value -> value.toLowerCase(Locale.ENGLISH)) + .toArray(String[]::new); + ColumnPath columnPath = ColumnPath.get(path); + PrimitiveType primitiveType = messageType.getType(columnPath.toArray()).asPrimitiveType(); + ColumnChunkMetadata column = ColumnChunkMetadata.get( + columnPath, + primitiveType, + CompressionCodecName.fromParquet(metaData.codec), + convertEncodingStats(metaData.encoding_stats), + readEncodings(metaData.encodings), + MetadataReader.readStats(Optional.ofNullable(parquetMetadata.getCreated_by()), Optional.ofNullable(metaData.statistics), primitiveType), + metaData.data_page_offset, + metaData.dictionary_page_offset, + metaData.num_values, + metaData.total_compressed_size, + metaData.total_uncompressed_size); + column.setColumnIndexReference(toColumnIndexReference(columnChunk)); + column.setOffsetIndexReference(toOffsetIndexReference(columnChunk)); + column.setBloomFilterOffset(metaData.bloom_filter_offset); + columnMetadataBuilder.add(column); + } + blocks.add(new BlockMetadata(rowGroup.getNum_rows(), columnMetadataBuilder.build())); + } + } + + return blocks; + } + + @VisibleForTesting + public FileMetaData getParquetMetadata() + { + return parquetMetadata; + } + + private static MessageType readParquetSchema(List schema) + { + Iterator schemaIterator = schema.iterator(); + SchemaElement rootSchema = schemaIterator.next(); + Types.MessageTypeBuilder builder = Types.buildMessage(); + readTypeSchema(builder, schemaIterator, rootSchema.getNum_children()); + return builder.named(rootSchema.name); + } + + private static void readTypeSchema(Types.GroupBuilder builder, Iterator schemaIterator, int typeCount) + { + for (int i = 0; i < typeCount; i++) { + SchemaElement element = schemaIterator.next(); + Types.Builder typeBuilder; + if (element.type == null) { + typeBuilder = builder.group(Type.Repetition.valueOf(element.repetition_type.name())); + readTypeSchema((Types.GroupBuilder) typeBuilder, schemaIterator, element.num_children); + } + else { + Types.PrimitiveBuilder primitiveBuilder = builder.primitive(getPrimitive(element.type), Type.Repetition.valueOf(element.repetition_type.name())); + if (element.isSetType_length()) { + primitiveBuilder.length(element.type_length); + } + if (element.isSetPrecision()) { + primitiveBuilder.precision(element.precision); + } + if (element.isSetScale()) { + primitiveBuilder.scale(element.scale); + } + typeBuilder = primitiveBuilder; + } + + // Reading of element.logicalType and element.converted_type corresponds to parquet-mr's code at + // https://github.com/apache/parquet-mr/blob/apache-parquet-1.12.0/parquet-hadoop/src/main/java/org/apache/parquet/format/converter/ParquetMetadataConverter.java#L1568-L1582 + LogicalTypeAnnotation annotationFromLogicalType = null; + if (element.isSetLogicalType()) { + annotationFromLogicalType = getLogicalTypeAnnotation(element.logicalType); + typeBuilder.as(annotationFromLogicalType); + } + if (element.isSetConverted_type()) { + LogicalTypeAnnotation annotationFromConvertedType = getLogicalTypeAnnotation(element.converted_type, element); + if (annotationFromLogicalType != null) { + // Both element.logicalType and element.converted_type set + if (annotationFromLogicalType.toOriginalType() == annotationFromConvertedType.toOriginalType()) { + // element.converted_type matches element.logicalType, even though annotationFromLogicalType may differ from annotationFromConvertedType + // Following parquet-mr behavior, we favor LogicalTypeAnnotation derived from element.logicalType, as potentially containing more information. + } + else { + // Following parquet-mr behavior, issue warning and let converted_type take precedence. + log.warn("Converted type and logical type metadata map to different OriginalType (convertedType: %s, logical type: %s). Using value in converted type.", + element.converted_type, element.logicalType); + // parquet-mr reads only OriginalType from converted_type. We retain full LogicalTypeAnnotation + // 1. for compatibility, as previous Trino reader code would read LogicalTypeAnnotation from element.converted_type and some additional fields. + // 2. so that we override LogicalTypeAnnotation annotation read from element.logicalType in case of mismatch detected. + typeBuilder.as(annotationFromConvertedType); + } + } + else { + // parquet-mr reads only OriginalType from converted_type. We retain full LogicalTypeAnnotation for compatibility, as previous + // Trino reader code would read LogicalTypeAnnotation from element.converted_type and some additional fields. + typeBuilder.as(annotationFromConvertedType); + } + } + + if (element.isSetField_id()) { + typeBuilder.id(element.field_id); + } + typeBuilder.named(element.name.toLowerCase(Locale.ENGLISH)); + } + } + + private static Set readEncodings(List encodings) + { + Set columnEncodings = new HashSet<>(); + for (org.apache.parquet.format.Encoding encoding : encodings) { + columnEncodings.add(getEncoding(encoding)); + } + return Collections.unmodifiableSet(columnEncodings); + } + + private static MessageType readMessageType(FileMetaData parquetMetadata, ParquetDataSourceId dataSourceId) + throws ParquetCorruptionException + { + List schema = parquetMetadata.getSchema(); + validateParquet(!schema.isEmpty(), dataSourceId, "Schema is empty"); + + Iterator schemaIterator = schema.iterator(); + SchemaElement rootSchema = schemaIterator.next(); + Types.MessageTypeBuilder builder = Types.buildMessage(); + readTypeSchema(builder, schemaIterator, rootSchema.getNum_children()); + return builder.named(rootSchema.name); + } + + private static Map keyValueMetaData(FileMetaData parquetMetadata) + { + if (parquetMetadata.getKey_value_metadata() == null) { + return ImmutableMap.of(); + } + return parquetMetadata.getKey_value_metadata() + .stream() + .collect(toImmutableMap(KeyValue::getKey, KeyValue::getValue, (ignore, second) -> second)); + } +} diff --git a/lib/trino-parquet/src/main/java/io/trino/parquet/metadata/PrunedBlockMetadata.java b/lib/trino-parquet/src/main/java/io/trino/parquet/metadata/PrunedBlockMetadata.java new file mode 100644 index 000000000000..4ceeadf9686f --- /dev/null +++ b/lib/trino-parquet/src/main/java/io/trino/parquet/metadata/PrunedBlockMetadata.java @@ -0,0 +1,98 @@ +/* + * 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.parquet.metadata; + +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import io.trino.parquet.ParquetCorruptionException; +import io.trino.parquet.ParquetDataSourceId; +import org.apache.parquet.column.ColumnDescriptor; + +import java.util.List; +import java.util.Map; +import java.util.Set; + +import static com.google.common.base.MoreObjects.toStringHelper; +import static com.google.common.collect.ImmutableMap.toImmutableMap; +import static java.util.Arrays.asList; +import static java.util.function.Function.identity; + +public final class PrunedBlockMetadata +{ + /** + * Stores only the necessary columns metadata from BlockMetadata and indexes them by path for efficient look-ups + */ + public static PrunedBlockMetadata createPrunedColumnsMetadata(BlockMetadata blockMetadata, ParquetDataSourceId dataSourceId, Map, ColumnDescriptor> descriptorsByPath) + throws ParquetCorruptionException + { + Set> requiredPaths = descriptorsByPath.keySet(); + Map, ColumnChunkMetadata> columnMetadataByPath = blockMetadata.columns().stream() + .collect(toImmutableMap( + column -> asList(column.getPath().toArray()), + identity(), + // Same column name may occur more than once when the file is written by case-sensitive tools + (oldValue, ignore) -> oldValue)); + ImmutableMap.Builder, ColumnChunkMetadata> columnMetadataByPathBuilder = ImmutableMap.builderWithExpectedSize(requiredPaths.size()); + for (Map.Entry, ColumnDescriptor> entry : descriptorsByPath.entrySet()) { + List requiredPath = entry.getKey(); + ColumnDescriptor columnDescriptor = entry.getValue(); + ColumnChunkMetadata columnChunkMetadata = columnMetadataByPath.get(requiredPath); + if (columnChunkMetadata == null) { + throw new ParquetCorruptionException(dataSourceId, "Metadata is missing for column: %s", columnDescriptor); + } + columnMetadataByPathBuilder.put(requiredPath, columnChunkMetadata); + } + return new PrunedBlockMetadata(blockMetadata.rowCount(), dataSourceId, columnMetadataByPathBuilder.buildOrThrow()); + } + + private final long rowCount; + private final ParquetDataSourceId dataSourceId; + private final Map, ColumnChunkMetadata> columnMetadataByPath; + + private PrunedBlockMetadata(long rowCount, ParquetDataSourceId dataSourceId, Map, ColumnChunkMetadata> columnMetadataByPath) + { + this.rowCount = rowCount; + this.dataSourceId = dataSourceId; + this.columnMetadataByPath = columnMetadataByPath; + } + + public long getRowCount() + { + return rowCount; + } + + public List getColumns() + { + return ImmutableList.copyOf(columnMetadataByPath.values()); + } + + public ColumnChunkMetadata getColumnChunkMetaData(ColumnDescriptor columnDescriptor) + throws ParquetCorruptionException + { + ColumnChunkMetadata columnChunkMetadata = columnMetadataByPath.get(asList(columnDescriptor.getPath())); + if (columnChunkMetadata == null) { + throw new ParquetCorruptionException(dataSourceId, "Metadata is missing for column: %s", columnDescriptor); + } + return columnChunkMetadata; + } + + @Override + public String toString() + { + return toStringHelper(this) + .add("rowCount", rowCount) + .add("columnMetadataByPath", columnMetadataByPath) + .toString(); + } +} diff --git a/lib/trino-parquet/src/main/java/io/trino/parquet/predicate/PredicateUtils.java b/lib/trino-parquet/src/main/java/io/trino/parquet/predicate/PredicateUtils.java index 9deb275b32f1..978e915a7a56 100644 --- a/lib/trino-parquet/src/main/java/io/trino/parquet/predicate/PredicateUtils.java +++ b/lib/trino-parquet/src/main/java/io/trino/parquet/predicate/PredicateUtils.java @@ -20,8 +20,16 @@ import io.airlift.slice.SliceInput; import io.trino.parquet.BloomFilterStore; import io.trino.parquet.DictionaryPage; +import io.trino.parquet.ParquetCorruptionException; import io.trino.parquet.ParquetDataSource; +import io.trino.parquet.ParquetDataSourceId; import io.trino.parquet.ParquetEncoding; +import io.trino.parquet.ParquetReaderOptions; +import io.trino.parquet.metadata.BlockMetadata; +import io.trino.parquet.metadata.ColumnChunkMetadata; +import io.trino.parquet.metadata.ParquetMetadata; +import io.trino.parquet.metadata.PrunedBlockMetadata; +import io.trino.parquet.reader.RowGroupInfo; import io.trino.spi.predicate.TupleDomain; import io.trino.spi.type.DecimalType; import io.trino.spi.type.Type; @@ -32,8 +40,6 @@ import org.apache.parquet.format.PageHeader; import org.apache.parquet.format.PageType; import org.apache.parquet.format.Util; -import org.apache.parquet.hadoop.metadata.BlockMetaData; -import org.apache.parquet.hadoop.metadata.ColumnChunkMetaData; import org.apache.parquet.internal.column.columnindex.OffsetIndex; import org.apache.parquet.internal.filter2.columnindex.ColumnIndexStore; import org.apache.parquet.io.ParquetDecodingException; @@ -49,9 +55,12 @@ import java.util.Optional; import java.util.Set; +import static io.trino.parquet.BloomFilterStore.getBloomFilterStore; import static io.trino.parquet.ParquetCompressionUtils.decompress; import static io.trino.parquet.ParquetReaderUtils.isOnlyDictionaryEncodingPages; import static io.trino.parquet.ParquetTypeUtils.getParquetEncoding; +import static io.trino.parquet.metadata.PrunedBlockMetadata.createPrunedColumnsMetadata; +import static io.trino.parquet.reader.TrinoColumnIndexStore.getColumnIndexStore; import static io.trino.spi.type.BigintType.BIGINT; import static io.trino.spi.type.DateType.DATE; import static io.trino.spi.type.IntegerType.INTEGER; @@ -126,7 +135,7 @@ public static TupleDomainParquetPredicate buildPredicate( public static boolean predicateMatches( TupleDomainParquetPredicate parquetPredicate, - BlockMetaData block, + PrunedBlockMetadata columnsMetadata, ParquetDataSource dataSource, Map, ColumnDescriptor> descriptorsByPath, TupleDomain parquetTupleDomain, @@ -136,11 +145,11 @@ public static boolean predicateMatches( int domainCompactionThreshold) throws IOException { - if (block.getRowCount() == 0) { + if (columnsMetadata.getRowCount() == 0) { return false; } - Map> columnStatistics = getStatistics(block, descriptorsByPath); - Map columnValueCounts = getColumnValueCounts(block, descriptorsByPath); + Map> columnStatistics = getStatistics(columnsMetadata, descriptorsByPath); + Map columnValueCounts = getColumnValueCounts(columnsMetadata, descriptorsByPath); Optional> candidateColumns = parquetPredicate.getIndexLookupCandidates(columnValueCounts, columnStatistics, dataSource.getId()); if (candidateColumns.isEmpty()) { return false; @@ -164,52 +173,95 @@ public static boolean predicateMatches( return dictionaryPredicatesMatch( indexPredicate, - block, + columnsMetadata, dataSource, descriptorsByPath, ImmutableSet.copyOf(candidateColumns.get()), columnIndexStore); } - private static Map> getStatistics(BlockMetaData blockMetadata, Map, ColumnDescriptor> descriptorsByPath) + public static List getFilteredRowGroups( + long splitStart, + long splitLength, + ParquetDataSource dataSource, + ParquetMetadata parquetMetadata, + List> parquetTupleDomains, + List parquetPredicates, + Map, ColumnDescriptor> descriptorsByPath, + DateTimeZone timeZone, + int domainCompactionThreshold, + ParquetReaderOptions options) + throws IOException { - ImmutableMap.Builder> statistics = ImmutableMap.builder(); - for (ColumnChunkMetaData columnMetaData : blockMetadata.getColumns()) { + long fileRowCount = 0; + ImmutableList.Builder rowGroupInfoBuilder = ImmutableList.builder(); + for (BlockMetadata block : parquetMetadata.getBlocks(splitStart, splitLength)) { + long blockStart = block.getStartingPos(); + boolean splitContainsBlock = splitStart <= blockStart && blockStart < splitStart + splitLength; + if (splitContainsBlock) { + for (int i = 0; i < parquetTupleDomains.size(); i++) { + TupleDomain parquetTupleDomain = parquetTupleDomains.get(i); + TupleDomainParquetPredicate parquetPredicate = parquetPredicates.get(i); + Optional columnIndex = getColumnIndexStore(dataSource, block, descriptorsByPath, parquetTupleDomain, options); + Optional bloomFilterStore = getBloomFilterStore(dataSource, block, parquetTupleDomain, options); + PrunedBlockMetadata columnsMetadata = createPrunedColumnsMetadata(block, dataSource.getId(), descriptorsByPath); + if (predicateMatches( + parquetPredicate, + columnsMetadata, + dataSource, + descriptorsByPath, + parquetTupleDomain, + columnIndex, + bloomFilterStore, + timeZone, + domainCompactionThreshold)) { + rowGroupInfoBuilder.add(new RowGroupInfo(columnsMetadata, fileRowCount, columnIndex)); + break; + } + } + } + fileRowCount += block.rowCount(); + } + return rowGroupInfoBuilder.build(); + } + + private static Map> getStatistics(PrunedBlockMetadata columnsMetadata, Map, ColumnDescriptor> descriptorsByPath) + throws ParquetCorruptionException + { + ImmutableMap.Builder> statistics = ImmutableMap.builderWithExpectedSize(descriptorsByPath.size()); + for (ColumnDescriptor descriptor : descriptorsByPath.values()) { + ColumnChunkMetadata columnMetaData = columnsMetadata.getColumnChunkMetaData(descriptor); Statistics columnStatistics = columnMetaData.getStatistics(); if (columnStatistics != null) { - ColumnDescriptor descriptor = descriptorsByPath.get(Arrays.asList(columnMetaData.getPath().toArray())); - if (descriptor != null) { - statistics.put(descriptor, columnStatistics); - } + statistics.put(descriptor, columnStatistics); } } return statistics.buildOrThrow(); } - private static Map getColumnValueCounts(BlockMetaData blockMetadata, Map, ColumnDescriptor> descriptorsByPath) + private static Map getColumnValueCounts(PrunedBlockMetadata columnsMetadata, Map, ColumnDescriptor> descriptorsByPath) + throws ParquetCorruptionException { - ImmutableMap.Builder columnValueCounts = ImmutableMap.builder(); - for (ColumnChunkMetaData columnMetaData : blockMetadata.getColumns()) { - ColumnDescriptor descriptor = descriptorsByPath.get(Arrays.asList(columnMetaData.getPath().toArray())); - if (descriptor != null) { - columnValueCounts.put(descriptor, columnMetaData.getValueCount()); - } + ImmutableMap.Builder columnValueCounts = ImmutableMap.builderWithExpectedSize(descriptorsByPath.size()); + for (ColumnDescriptor descriptor : descriptorsByPath.values()) { + ColumnChunkMetadata columnMetaData = columnsMetadata.getColumnChunkMetaData(descriptor); + columnValueCounts.put(descriptor, columnMetaData.getValueCount()); } return columnValueCounts.buildOrThrow(); } private static boolean dictionaryPredicatesMatch( TupleDomainParquetPredicate parquetPredicate, - BlockMetaData blockMetadata, + PrunedBlockMetadata columnsMetadata, ParquetDataSource dataSource, Map, ColumnDescriptor> descriptorsByPath, Set candidateColumns, Optional columnIndexStore) throws IOException { - for (ColumnChunkMetaData columnMetaData : blockMetadata.getColumns()) { - ColumnDescriptor descriptor = descriptorsByPath.get(Arrays.asList(columnMetaData.getPath().toArray())); - if (descriptor == null || !candidateColumns.contains(descriptor)) { + for (ColumnDescriptor descriptor : descriptorsByPath.values()) { + ColumnChunkMetadata columnMetaData = columnsMetadata.getColumnChunkMetaData(descriptor); + if (!candidateColumns.contains(descriptor)) { continue; } if (isOnlyDictionaryEncodingPages(columnMetaData)) { @@ -229,7 +281,7 @@ private static boolean dictionaryPredicatesMatch( private static Optional readDictionaryPage( ParquetDataSource dataSource, - ColumnChunkMetaData columnMetaData, + ColumnChunkMetadata columnMetaData, Optional columnIndexStore) throws IOException { @@ -254,10 +306,10 @@ private static Optional readDictionaryPage( } // Get the dictionary page header and the dictionary in single read Slice buffer = dataSource.readFully(columnMetaData.getStartingPos(), dictionaryPageSize); - return readPageHeaderWithData(buffer.getInput()).map(data -> decodeDictionaryPage(data, columnMetaData)); + return readPageHeaderWithData(buffer.getInput()).map(data -> decodeDictionaryPage(dataSource.getId(), data, columnMetaData)); } - private static Optional getDictionaryPageSize(ColumnIndexStore columnIndexStore, ColumnChunkMetaData columnMetaData) + private static Optional getDictionaryPageSize(ColumnIndexStore columnIndexStore, ColumnChunkMetadata columnMetaData) { OffsetIndex offsetIndex = columnIndexStore.getOffsetIndex(columnMetaData.getPath()); if (offsetIndex == null) { @@ -293,7 +345,7 @@ private static Optional readPageHeaderWithData(SliceInput in inputStream.readSlice(pageHeader.getCompressed_page_size()))); } - private static DictionaryPage decodeDictionaryPage(PageHeaderWithData pageHeaderWithData, ColumnChunkMetaData chunkMetaData) + private static DictionaryPage decodeDictionaryPage(ParquetDataSourceId dataSourceId, PageHeaderWithData pageHeaderWithData, ColumnChunkMetadata chunkMetaData) { PageHeader pageHeader = pageHeaderWithData.pageHeader(); DictionaryPageHeader dicHeader = pageHeader.getDictionary_page_header(); @@ -302,7 +354,7 @@ private static DictionaryPage decodeDictionaryPage(PageHeaderWithData pageHeader Slice compressedData = pageHeaderWithData.compressedData(); try { - return new DictionaryPage(decompress(chunkMetaData.getCodec().getParquetCompressionCodec(), compressedData, pageHeader.getUncompressed_page_size()), dictionarySize, encoding); + return new DictionaryPage(decompress(dataSourceId, chunkMetaData.getCodec().getParquetCompressionCodec(), compressedData, pageHeader.getUncompressed_page_size()), dictionarySize, encoding); } catch (IOException e) { throw new ParquetDecodingException("Could not decode the dictionary for " + chunkMetaData.getPath(), e); diff --git a/lib/trino-parquet/src/main/java/io/trino/parquet/reader/ColumnReaderFactory.java b/lib/trino-parquet/src/main/java/io/trino/parquet/reader/ColumnReaderFactory.java index 2a457f97d302..0c9cc68f0882 100644 --- a/lib/trino-parquet/src/main/java/io/trino/parquet/reader/ColumnReaderFactory.java +++ b/lib/trino-parquet/src/main/java/io/trino/parquet/reader/ColumnReaderFactory.java @@ -17,11 +17,9 @@ import io.trino.memory.context.LocalMemoryContext; import io.trino.parquet.ParquetReaderOptions; import io.trino.parquet.PrimitiveField; -import io.trino.parquet.reader.decoders.ValueDecoder; import io.trino.parquet.reader.decoders.ValueDecoders; import io.trino.parquet.reader.flat.ColumnAdapter; import io.trino.parquet.reader.flat.FlatColumnReader; -import io.trino.parquet.reader.flat.FlatDefinitionLevelDecoder; import io.trino.spi.TrinoException; import io.trino.spi.type.AbstractIntType; import io.trino.spi.type.AbstractLongType; @@ -50,11 +48,13 @@ import static io.trino.parquet.ParquetEncoding.PLAIN; import static io.trino.parquet.ParquetTypeUtils.createDecimalType; import static io.trino.parquet.reader.decoders.ValueDecoder.ValueDecodersProvider; +import static io.trino.parquet.reader.decoders.ValueDecoder.createLevelsDecoder; import static io.trino.parquet.reader.flat.BinaryColumnAdapter.BINARY_ADAPTER; import static io.trino.parquet.reader.flat.ByteColumnAdapter.BYTE_ADAPTER; import static io.trino.parquet.reader.flat.DictionaryDecoder.DictionaryDecoderProvider; import static io.trino.parquet.reader.flat.DictionaryDecoder.getDictionaryDecoder; import static io.trino.parquet.reader.flat.Fixed12ColumnAdapter.FIXED12_ADAPTER; +import static io.trino.parquet.reader.flat.FlatDefinitionLevelDecoder.getFlatDefinitionLevelDecoder; import static io.trino.parquet.reader.flat.Int128ColumnAdapter.INT128_ADAPTER; import static io.trino.parquet.reader.flat.IntColumnAdapter.INT_ADAPTER; import static io.trino.parquet.reader.flat.LongColumnAdapter.LONG_ADAPTER; @@ -76,7 +76,6 @@ import static java.util.Objects.requireNonNull; import static org.apache.parquet.schema.LogicalTypeAnnotation.TimeUnit.MICROS; import static org.apache.parquet.schema.LogicalTypeAnnotation.TimeUnit.MILLIS; -import static org.apache.parquet.schema.LogicalTypeAnnotation.TimeUnit.NANOS; import static org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName.BINARY; import static org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName.FIXED_LEN_BYTE_ARRAY; import static org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName.FLOAT; @@ -86,15 +85,13 @@ public final class ColumnReaderFactory { + private static final int PREFERRED_BIT_WIDTH = getVectorBitSize(); + private final DateTimeZone timeZone; - private final boolean useBatchColumnReaders; - private final boolean useBatchNestedColumnReaders; public ColumnReaderFactory(DateTimeZone timeZone, ParquetReaderOptions readerOptions) { this.timeZone = requireNonNull(timeZone, "dateTimeZone is null"); - this.useBatchColumnReaders = readerOptions.useBatchColumnReaders(); - this.useBatchNestedColumnReaders = readerOptions.useBatchNestedColumnReaders(); } public ColumnReader create(PrimitiveField field, AggregatedMemoryContext aggregatedMemoryContext) @@ -103,241 +100,184 @@ public ColumnReader create(PrimitiveField field, AggregatedMemoryContext aggrega PrimitiveTypeName primitiveType = field.getDescriptor().getPrimitiveType().getPrimitiveTypeName(); LogicalTypeAnnotation annotation = field.getDescriptor().getPrimitiveType().getLogicalTypeAnnotation(); LocalMemoryContext memoryContext = aggregatedMemoryContext.newLocalMemoryContext(ColumnReader.class.getSimpleName()); - if (useBatchedColumnReaders(field)) { - ValueDecoders valueDecoders = new ValueDecoders(field); - if (BOOLEAN.equals(type) && primitiveType == PrimitiveTypeName.BOOLEAN) { - return createColumnReader(field, valueDecoders::getBooleanDecoder, BYTE_ADAPTER, memoryContext); + ValueDecoders valueDecoders = new ValueDecoders(field); + if (BOOLEAN.equals(type) && primitiveType == PrimitiveTypeName.BOOLEAN) { + return createColumnReader(field, valueDecoders::getBooleanDecoder, BYTE_ADAPTER, memoryContext); + } + if (TINYINT.equals(type) && isIntegerOrDecimalPrimitive(primitiveType)) { + if (isZeroScaleShortDecimalAnnotation(annotation)) { + return createColumnReader(field, valueDecoders::getShortDecimalToByteDecoder, BYTE_ADAPTER, memoryContext); } - if (TINYINT.equals(type) && isIntegerOrDecimalPrimitive(primitiveType)) { - if (isZeroScaleShortDecimalAnnotation(annotation)) { - return createColumnReader(field, valueDecoders::getShortDecimalToByteDecoder, BYTE_ADAPTER, memoryContext); - } - if (!isIntegerAnnotationAndPrimitive(annotation, primitiveType)) { - throw unsupportedException(type, field); - } - return createColumnReader(field, valueDecoders::getByteDecoder, BYTE_ADAPTER, memoryContext); + if (!isIntegerAnnotationAndPrimitive(annotation, primitiveType)) { + throw unsupportedException(type, field); } - if (SMALLINT.equals(type) && isIntegerOrDecimalPrimitive(primitiveType)) { - if (isZeroScaleShortDecimalAnnotation(annotation)) { - return createColumnReader(field, valueDecoders::getShortDecimalToShortDecoder, SHORT_ADAPTER, memoryContext); - } - if (!isIntegerAnnotationAndPrimitive(annotation, primitiveType)) { - throw unsupportedException(type, field); - } - return createColumnReader(field, valueDecoders::getShortDecoder, SHORT_ADAPTER, memoryContext); + return createColumnReader(field, valueDecoders::getByteDecoder, BYTE_ADAPTER, memoryContext); + } + if (SMALLINT.equals(type) && isIntegerOrDecimalPrimitive(primitiveType)) { + if (isZeroScaleShortDecimalAnnotation(annotation)) { + return createColumnReader(field, valueDecoders::getShortDecimalToShortDecoder, SHORT_ADAPTER, memoryContext); } - if (DATE.equals(type) && primitiveType == INT32) { - if (annotation == null || annotation instanceof DateLogicalTypeAnnotation) { - return createColumnReader(field, valueDecoders::getIntDecoder, INT_ADAPTER, memoryContext); - } + if (!isIntegerAnnotationAndPrimitive(annotation, primitiveType)) { throw unsupportedException(type, field); } - if (type instanceof AbstractIntType && isIntegerOrDecimalPrimitive(primitiveType)) { - if (isZeroScaleShortDecimalAnnotation(annotation)) { - return createColumnReader(field, valueDecoders::getShortDecimalToIntDecoder, INT_ADAPTER, memoryContext); - } - if (!isIntegerAnnotationAndPrimitive(annotation, primitiveType)) { - throw unsupportedException(type, field); - } + return createColumnReader(field, valueDecoders::getShortDecoder, SHORT_ADAPTER, memoryContext); + } + if (DATE.equals(type) && primitiveType == INT32) { + if (annotation == null || annotation instanceof DateLogicalTypeAnnotation) { return createColumnReader(field, valueDecoders::getIntDecoder, INT_ADAPTER, memoryContext); } - if (type instanceof TimeType) { - if (!(annotation instanceof TimeLogicalTypeAnnotation timeAnnotation)) { - throw unsupportedException(type, field); - } - if (primitiveType == INT64 && timeAnnotation.getUnit() == MICROS) { - return createColumnReader(field, valueDecoders::getTimeMicrosDecoder, LONG_ADAPTER, memoryContext); - } - if (primitiveType == INT32 && timeAnnotation.getUnit() == MILLIS) { - return createColumnReader(field, valueDecoders::getTimeMillisDecoder, LONG_ADAPTER, memoryContext); - } + throw unsupportedException(type, field); + } + if (type instanceof AbstractIntType && isIntegerOrDecimalPrimitive(primitiveType)) { + if (isZeroScaleShortDecimalAnnotation(annotation)) { + return createColumnReader(field, valueDecoders::getShortDecimalToIntDecoder, INT_ADAPTER, memoryContext); + } + if (!isIntegerAnnotationAndPrimitive(annotation, primitiveType)) { throw unsupportedException(type, field); } - if (BIGINT.equals(type) && primitiveType == INT64 - && (annotation instanceof TimestampLogicalTypeAnnotation || annotation instanceof TimeLogicalTypeAnnotation)) { - return createColumnReader(field, valueDecoders::getLongDecoder, LONG_ADAPTER, memoryContext); + return createColumnReader(field, valueDecoders::getIntDecoder, INT_ADAPTER, memoryContext); + } + if (type instanceof TimeType) { + if (!(annotation instanceof TimeLogicalTypeAnnotation timeAnnotation)) { + throw unsupportedException(type, field); + } + if (primitiveType == INT64 && timeAnnotation.getUnit() == MICROS) { + return createColumnReader(field, valueDecoders::getTimeMicrosDecoder, LONG_ADAPTER, memoryContext); + } + if (primitiveType == INT32 && timeAnnotation.getUnit() == MILLIS) { + return createColumnReader(field, valueDecoders::getTimeMillisDecoder, LONG_ADAPTER, memoryContext); + } + throw unsupportedException(type, field); + } + if (BIGINT.equals(type) && primitiveType == INT64 + && (annotation instanceof TimestampLogicalTypeAnnotation || annotation instanceof TimeLogicalTypeAnnotation)) { + return createColumnReader(field, valueDecoders::getLongDecoder, LONG_ADAPTER, memoryContext); + } + if (type instanceof AbstractLongType && isIntegerOrDecimalPrimitive(primitiveType)) { + if (isZeroScaleShortDecimalAnnotation(annotation)) { + return createColumnReader(field, valueDecoders::getShortDecimalDecoder, LONG_ADAPTER, memoryContext); } - if (type instanceof AbstractLongType && isIntegerOrDecimalPrimitive(primitiveType)) { - if (isZeroScaleShortDecimalAnnotation(annotation)) { - return createColumnReader(field, valueDecoders::getShortDecimalDecoder, LONG_ADAPTER, memoryContext); - } - if (!isIntegerAnnotationAndPrimitive(annotation, primitiveType)) { - throw unsupportedException(type, field); - } - if (primitiveType == INT32) { - return createColumnReader(field, valueDecoders::getInt32ToLongDecoder, LONG_ADAPTER, memoryContext); - } - if (primitiveType == INT64) { - return createColumnReader(field, valueDecoders::getLongDecoder, LONG_ADAPTER, memoryContext); - } + if (!isIntegerAnnotationAndPrimitive(annotation, primitiveType)) { + throw unsupportedException(type, field); + } + if (primitiveType == INT32) { + return createColumnReader(field, valueDecoders::getInt32ToLongDecoder, LONG_ADAPTER, memoryContext); + } + if (primitiveType == INT64) { + return createColumnReader(field, valueDecoders::getLongDecoder, LONG_ADAPTER, memoryContext); } - if (REAL.equals(type) && primitiveType == FLOAT) { - return createColumnReader(field, valueDecoders::getRealDecoder, INT_ADAPTER, memoryContext); + } + if (REAL.equals(type) && primitiveType == FLOAT) { + return createColumnReader(field, valueDecoders::getRealDecoder, INT_ADAPTER, memoryContext); + } + if (DOUBLE.equals(type)) { + if (primitiveType == PrimitiveTypeName.DOUBLE) { + return createColumnReader(field, valueDecoders::getDoubleDecoder, LONG_ADAPTER, memoryContext); } - if (DOUBLE.equals(type)) { - if (primitiveType == PrimitiveTypeName.DOUBLE) { - return createColumnReader(field, valueDecoders::getDoubleDecoder, LONG_ADAPTER, memoryContext); - } - if (primitiveType == FLOAT) { - return createColumnReader(field, valueDecoders::getFloatToDoubleDecoder, LONG_ADAPTER, memoryContext); - } + if (primitiveType == FLOAT) { + return createColumnReader(field, valueDecoders::getFloatToDoubleDecoder, LONG_ADAPTER, memoryContext); } - if (type instanceof TimestampType timestampType && primitiveType == INT96) { - if (timestampType.isShort()) { - return createColumnReader( - field, - (encoding) -> valueDecoders.getInt96ToShortTimestampDecoder(encoding, timeZone), - LONG_ADAPTER, - memoryContext); - } + } + if (type instanceof TimestampType timestampType && primitiveType == INT96) { + if (timestampType.isShort()) { return createColumnReader( field, - (encoding) -> valueDecoders.getInt96ToLongTimestampDecoder(encoding, timeZone), - FIXED12_ADAPTER, + (encoding) -> valueDecoders.getInt96ToShortTimestampDecoder(encoding, timeZone), + LONG_ADAPTER, memoryContext); } - if (type instanceof TimestampWithTimeZoneType timestampWithTimeZoneType && primitiveType == INT96) { - if (timestampWithTimeZoneType.isShort()) { - return createColumnReader(field, valueDecoders::getInt96ToShortTimestampWithTimeZoneDecoder, LONG_ADAPTER, memoryContext); - } + return createColumnReader( + field, + (encoding) -> valueDecoders.getInt96ToLongTimestampDecoder(encoding, timeZone), + FIXED12_ADAPTER, + memoryContext); + } + if (type instanceof TimestampWithTimeZoneType timestampWithTimeZoneType && primitiveType == INT96) { + if (timestampWithTimeZoneType.isShort()) { + return createColumnReader(field, valueDecoders::getInt96ToShortTimestampWithTimeZoneDecoder, LONG_ADAPTER, memoryContext); + } + return createColumnReader(field, valueDecoders::getInt96ToLongTimestampWithTimeZoneDecoder, FIXED12_ADAPTER, memoryContext); + } + if (type instanceof TimestampType timestampType && primitiveType == INT64) { + if (!(annotation instanceof TimestampLogicalTypeAnnotation timestampAnnotation)) { throw unsupportedException(type, field); } - if (type instanceof TimestampType timestampType && primitiveType == INT64) { - if (!(annotation instanceof TimestampLogicalTypeAnnotation timestampAnnotation)) { - throw unsupportedException(type, field); - } - if (timestampType.isShort()) { - return switch (timestampAnnotation.getUnit()) { - case MILLIS -> createColumnReader(field, valueDecoders::getInt64TimestampMillsToShortTimestampDecoder, LONG_ADAPTER, memoryContext); - case MICROS -> createColumnReader(field, valueDecoders::getInt64TimestampMicrosToShortTimestampDecoder, LONG_ADAPTER, memoryContext); - case NANOS -> createColumnReader(field, valueDecoders::getInt64TimestampNanosToShortTimestampDecoder, LONG_ADAPTER, memoryContext); - }; - } + DateTimeZone readTimeZone = timestampAnnotation.isAdjustedToUTC() ? timeZone : DateTimeZone.UTC; + if (timestampType.isShort()) { return switch (timestampAnnotation.getUnit()) { - case MILLIS -> createColumnReader(field, valueDecoders::getInt64TimestampMillisToLongTimestampDecoder, FIXED12_ADAPTER, memoryContext); - case MICROS -> createColumnReader(field, valueDecoders::getInt64TimestampMicrosToLongTimestampDecoder, FIXED12_ADAPTER, memoryContext); - case NANOS -> createColumnReader(field, valueDecoders::getInt64TimestampNanosToLongTimestampDecoder, FIXED12_ADAPTER, memoryContext); + case MILLIS -> createColumnReader(field, encoding -> valueDecoders.getInt64TimestampMillisToShortTimestampDecoder(encoding, readTimeZone), LONG_ADAPTER, memoryContext); + case MICROS -> createColumnReader(field, encoding -> valueDecoders.getInt64TimestampMicrosToShortTimestampDecoder(encoding, readTimeZone), LONG_ADAPTER, memoryContext); + case NANOS -> createColumnReader(field, encoding -> valueDecoders.getInt64TimestampNanosToShortTimestampDecoder(encoding, readTimeZone), LONG_ADAPTER, memoryContext); }; } - if (type instanceof TimestampWithTimeZoneType timestampWithTimeZoneType && primitiveType == INT64) { - if (!(annotation instanceof TimestampLogicalTypeAnnotation timestampAnnotation)) { - throw unsupportedException(type, field); - } - if (timestampWithTimeZoneType.isShort()) { - return switch (timestampAnnotation.getUnit()) { - case MILLIS -> createColumnReader(field, valueDecoders::getInt64TimestampMillsToShortTimestampWithTimeZoneDecoder, LONG_ADAPTER, memoryContext); - case MICROS -> createColumnReader(field, valueDecoders::getInt64TimestampMicrosToShortTimestampWithTimeZoneDecoder, LONG_ADAPTER, memoryContext); - case NANOS -> throw unsupportedException(type, field); - }; - } + return switch (timestampAnnotation.getUnit()) { + case MILLIS -> createColumnReader(field, encoding -> valueDecoders.getInt64TimestampMillisToLongTimestampDecoder(encoding, readTimeZone), FIXED12_ADAPTER, memoryContext); + case MICROS -> createColumnReader(field, encoding -> valueDecoders.getInt64TimestampMicrosToLongTimestampDecoder(encoding, readTimeZone), FIXED12_ADAPTER, memoryContext); + case NANOS -> createColumnReader(field, encoding -> valueDecoders.getInt64TimestampNanosToLongTimestampDecoder(encoding, readTimeZone), FIXED12_ADAPTER, memoryContext); + }; + } + if (type instanceof TimestampWithTimeZoneType timestampWithTimeZoneType && primitiveType == INT64) { + if (!(annotation instanceof TimestampLogicalTypeAnnotation timestampAnnotation)) { + throw unsupportedException(type, field); + } + if (timestampWithTimeZoneType.isShort()) { return switch (timestampAnnotation.getUnit()) { - case MILLIS, NANOS -> throw unsupportedException(type, field); - case MICROS -> createColumnReader(field, valueDecoders::getInt64TimestampMicrosToLongTimestampWithTimeZoneDecoder, FIXED12_ADAPTER, memoryContext); + case MILLIS -> createColumnReader(field, valueDecoders::getInt64TimestampMillsToShortTimestampWithTimeZoneDecoder, LONG_ADAPTER, memoryContext); + case MICROS -> createColumnReader(field, valueDecoders::getInt64TimestampMicrosToShortTimestampWithTimeZoneDecoder, LONG_ADAPTER, memoryContext); + case NANOS -> throw unsupportedException(type, field); }; } - if (type instanceof DecimalType decimalType && decimalType.isShort() - && isIntegerOrDecimalPrimitive(primitiveType)) { - if (primitiveType == INT32 && isIntegerAnnotation(annotation)) { - return createColumnReader(field, valueDecoders::getInt32ToShortDecimalDecoder, LONG_ADAPTER, memoryContext); - } - if (!(annotation instanceof DecimalLogicalTypeAnnotation decimalAnnotation)) { - throw unsupportedException(type, field); - } - if (isDecimalRescaled(decimalAnnotation, decimalType)) { - return createColumnReader(field, valueDecoders::getRescaledShortDecimalDecoder, LONG_ADAPTER, memoryContext); - } - return createColumnReader(field, valueDecoders::getShortDecimalDecoder, LONG_ADAPTER, memoryContext); - } - if (type instanceof DecimalType decimalType && !decimalType.isShort() - && isIntegerOrDecimalPrimitive(primitiveType)) { - if (!(annotation instanceof DecimalLogicalTypeAnnotation decimalAnnotation)) { - throw unsupportedException(type, field); - } - if (isDecimalRescaled(decimalAnnotation, decimalType)) { - return createColumnReader(field, valueDecoders::getRescaledLongDecimalDecoder, INT128_ADAPTER, memoryContext); - } - return createColumnReader(field, valueDecoders::getLongDecimalDecoder, INT128_ADAPTER, memoryContext); + return switch (timestampAnnotation.getUnit()) { + case MILLIS, NANOS -> throw unsupportedException(type, field); + case MICROS -> createColumnReader(field, valueDecoders::getInt64TimestampMicrosToLongTimestampWithTimeZoneDecoder, FIXED12_ADAPTER, memoryContext); + }; + } + if (type instanceof DecimalType decimalType && decimalType.isShort() + && isIntegerOrDecimalPrimitive(primitiveType)) { + if (primitiveType == INT32 && isIntegerAnnotation(annotation)) { + return createColumnReader(field, valueDecoders::getInt32ToShortDecimalDecoder, LONG_ADAPTER, memoryContext); } - if (type instanceof VarcharType varcharType && !varcharType.isUnbounded() && primitiveType == BINARY) { - return createColumnReader(field, valueDecoders::getBoundedVarcharBinaryDecoder, BINARY_ADAPTER, memoryContext); + if (!(annotation instanceof DecimalLogicalTypeAnnotation decimalAnnotation)) { + throw unsupportedException(type, field); } - if (type instanceof CharType && primitiveType == BINARY) { - return createColumnReader(field, valueDecoders::getCharBinaryDecoder, BINARY_ADAPTER, memoryContext); + if (isDecimalRescaled(decimalAnnotation, decimalType)) { + return createColumnReader(field, valueDecoders::getRescaledShortDecimalDecoder, LONG_ADAPTER, memoryContext); } - if (type instanceof AbstractVariableWidthType && primitiveType == BINARY) { - return createColumnReader(field, valueDecoders::getBinaryDecoder, BINARY_ADAPTER, memoryContext); + return createColumnReader(field, valueDecoders::getShortDecimalDecoder, LONG_ADAPTER, memoryContext); + } + if (type instanceof DecimalType decimalType && !decimalType.isShort() + && isIntegerOrDecimalPrimitive(primitiveType)) { + if (!(annotation instanceof DecimalLogicalTypeAnnotation decimalAnnotation)) { + throw unsupportedException(type, field); } - if ((VARBINARY.equals(type) || VARCHAR.equals(type)) && primitiveType == FIXED_LEN_BYTE_ARRAY) { - return createColumnReader(field, valueDecoders::getFixedWidthBinaryDecoder, BINARY_ADAPTER, memoryContext); + if (isDecimalRescaled(decimalAnnotation, decimalType)) { + return createColumnReader(field, valueDecoders::getRescaledLongDecimalDecoder, INT128_ADAPTER, memoryContext); } - if (UUID.equals(type) && primitiveType == FIXED_LEN_BYTE_ARRAY) { - // Iceberg 0.11.1 writes UUID as FIXED_LEN_BYTE_ARRAY without logical type annotation (see https://github.com/apache/iceberg/pull/2913) - // To support such files, we bet on the logical type to be UUID based on the Trino UUID type check. - if (annotation == null || isLogicalUuid(annotation)) { - return createColumnReader(field, valueDecoders::getUuidDecoder, INT128_ADAPTER, memoryContext); - } + return createColumnReader(field, valueDecoders::getLongDecimalDecoder, INT128_ADAPTER, memoryContext); + } + if (type instanceof VarcharType varcharType && !varcharType.isUnbounded() && primitiveType == BINARY) { + return createColumnReader(field, valueDecoders::getBoundedVarcharBinaryDecoder, BINARY_ADAPTER, memoryContext); + } + if (type instanceof CharType && primitiveType == BINARY) { + return createColumnReader(field, valueDecoders::getCharBinaryDecoder, BINARY_ADAPTER, memoryContext); + } + if (type instanceof AbstractVariableWidthType && primitiveType == BINARY) { + return createColumnReader(field, valueDecoders::getBinaryDecoder, BINARY_ADAPTER, memoryContext); + } + if ((VARBINARY.equals(type) || VARCHAR.equals(type)) && primitiveType == FIXED_LEN_BYTE_ARRAY) { + if (annotation instanceof DecimalLogicalTypeAnnotation) { throw unsupportedException(type, field); } - throw new TrinoException( - NOT_SUPPORTED, - format("Reading Trino column (%s) from Parquet column (%s) is not supported by optimized parquet reader", type, field.getDescriptor())); + return createColumnReader(field, valueDecoders::getFixedWidthBinaryDecoder, BINARY_ADAPTER, memoryContext); } - - return switch (primitiveType) { - case BOOLEAN -> new BooleanColumnReader(field); - case INT32 -> createDecimalColumnReader(field).orElseGet(() -> { - if (type instanceof DecimalType decimalType && decimalType.isShort()) { - return new Int32ShortDecimalColumnReader(field); - } - return new IntColumnReader(field); - }); - case INT64 -> { - if (annotation instanceof TimeLogicalTypeAnnotation timeAnnotation) { - if (field.getType() instanceof TimeType && timeAnnotation.getUnit() == MICROS) { - yield new TimeMicrosColumnReader(field); - } - else if (BIGINT.equals(field.getType())) { - yield new LongColumnReader(field); - } - throw unsupportedException(type, field); - } - if (annotation instanceof TimestampLogicalTypeAnnotation timestampAnnotation) { - if (timestampAnnotation.getUnit() == MILLIS) { - yield new Int64TimestampMillisColumnReader(field); - } - if (timestampAnnotation.getUnit() == MICROS) { - yield new TimestampMicrosColumnReader(field); - } - if (timestampAnnotation.getUnit() == NANOS) { - yield new Int64TimestampNanosColumnReader(field); - } - throw unsupportedException(type, field); - } - yield createDecimalColumnReader(field).orElse(new LongColumnReader(field)); + if (UUID.equals(type) && primitiveType == FIXED_LEN_BYTE_ARRAY) { + // Iceberg 0.11.1 writes UUID as FIXED_LEN_BYTE_ARRAY without logical type annotation (see https://github.com/apache/iceberg/pull/2913) + // To support such files, we bet on the logical type to be UUID based on the Trino UUID type check. + if (annotation == null || isLogicalUuid(annotation)) { + return createColumnReader(field, valueDecoders::getUuidDecoder, INT128_ADAPTER, memoryContext); } - case INT96 -> new TimestampColumnReader(field, timeZone); - case FLOAT -> new FloatColumnReader(field); - case DOUBLE -> new DoubleColumnReader(field); - case BINARY -> createDecimalColumnReader(field).orElse(new BinaryColumnReader(field)); - case FIXED_LEN_BYTE_ARRAY -> { - Optional decimalColumnReader = createDecimalColumnReader(field); - if (decimalColumnReader.isPresent()) { - yield decimalColumnReader.get(); - } - if (isLogicalUuid(annotation)) { - yield new UuidColumnReader(field); - } - if (VARBINARY.equals(type) || VARCHAR.equals(type)) { - yield new BinaryColumnReader(field); - } - if (annotation == null) { - // Iceberg 0.11.1 writes UUID as FIXED_LEN_BYTE_ARRAY without logical type annotation (see https://github.com/apache/iceberg/pull/2913) - // To support such files, we bet on the type to be UUID, which gets verified later, when reading the column data. - yield new UuidColumnReader(field); - } - throw unsupportedException(type, field); - } - }; + } + throw unsupportedException(type, field); } private static ColumnReader createColumnReader( @@ -355,7 +295,7 @@ private static ColumnReader createColumnReader( return new FlatColumnReader<>( field, decodersProvider, - FlatDefinitionLevelDecoder::getFlatDefinitionLevelDecoder, + maxDefinitionLevel -> getFlatDefinitionLevelDecoder(maxDefinitionLevel), dictionaryDecoderProvider, columnAdapter, memoryContext); @@ -363,23 +303,15 @@ private static ColumnReader createColumnReader( return new NestedColumnReader<>( field, decodersProvider, - ValueDecoder::createLevelsDecoder, + maxLevel -> createLevelsDecoder(maxLevel), dictionaryDecoderProvider, columnAdapter, memoryContext); } - private boolean useBatchedColumnReaders(PrimitiveField field) - { - if (isFlatColumn(field)) { - return useBatchColumnReaders; - } - return useBatchColumnReaders && useBatchNestedColumnReaders; - } - private static boolean isFlatColumn(PrimitiveField field) { - return field.getDescriptor().getPath().length == 1; + return field.getDescriptor().getPath().length == 1 && field.getRepetitionLevel() == 0; } private static boolean isLogicalUuid(LogicalTypeAnnotation annotation) @@ -428,7 +360,7 @@ private static boolean isIntegerOrDecimalPrimitive(PrimitiveTypeName primitiveTy return primitiveType == INT32 || primitiveType == INT64 || primitiveType == BINARY || primitiveType == FIXED_LEN_BYTE_ARRAY; } - private static boolean isIntegerAnnotationAndPrimitive(LogicalTypeAnnotation typeAnnotation, PrimitiveTypeName primitiveType) + public static boolean isIntegerAnnotationAndPrimitive(LogicalTypeAnnotation typeAnnotation, PrimitiveTypeName primitiveType) { return isIntegerAnnotation(typeAnnotation) && (primitiveType == INT32 || primitiveType == INT64); } @@ -437,4 +369,22 @@ private static TrinoException unsupportedException(Type type, PrimitiveField fie { return new TrinoException(NOT_SUPPORTED, format("Unsupported Trino column type (%s) for Parquet column (%s)", type, field.getDescriptor())); } + + private static boolean isVectorizedDecodingSupported() + { + // Performance gains with vectorized decoding are validated only when the hardware platform provides at least 256 bit width registers + // Graviton 2 machines return false here, whereas x86 and Graviton 3 machines return true + return PREFERRED_BIT_WIDTH >= 256; + } + + private static int getVectorBitSize() + { + try { + Class clazz = Class.forName("jdk.incubator.vector.VectorShape"); + return (int) clazz.getMethod("vectorBitSize").invoke(clazz.getMethod("preferredShape").invoke(null)); + } + catch (Throwable e) { + return -1; + } + } } diff --git a/lib/trino-parquet/src/main/java/io/trino/parquet/reader/MetadataReader.java b/lib/trino-parquet/src/main/java/io/trino/parquet/reader/MetadataReader.java index f085d8f5e4e3..2109ed378d69 100644 --- a/lib/trino-parquet/src/main/java/io/trino/parquet/reader/MetadataReader.java +++ b/lib/trino-parquet/src/main/java/io/trino/parquet/reader/MetadataReader.java @@ -20,44 +20,20 @@ import io.trino.parquet.ParquetDataSource; import io.trino.parquet.ParquetDataSourceId; import io.trino.parquet.ParquetWriteValidation; +import io.trino.parquet.metadata.FileMetadata; +import io.trino.parquet.metadata.ParquetMetadata; import org.apache.parquet.CorruptStatistics; import org.apache.parquet.column.statistics.BinaryStatistics; -import org.apache.parquet.format.ColumnChunk; -import org.apache.parquet.format.ColumnMetaData; -import org.apache.parquet.format.Encoding; import org.apache.parquet.format.FileMetaData; -import org.apache.parquet.format.KeyValue; -import org.apache.parquet.format.RowGroup; -import org.apache.parquet.format.SchemaElement; import org.apache.parquet.format.Statistics; -import org.apache.parquet.format.Type; import org.apache.parquet.format.converter.ParquetMetadataConverter; -import org.apache.parquet.hadoop.metadata.BlockMetaData; -import org.apache.parquet.hadoop.metadata.ColumnChunkMetaData; -import org.apache.parquet.hadoop.metadata.ColumnPath; -import org.apache.parquet.hadoop.metadata.CompressionCodecName; -import org.apache.parquet.hadoop.metadata.ParquetMetadata; -import org.apache.parquet.internal.hadoop.metadata.IndexReference; import org.apache.parquet.schema.LogicalTypeAnnotation; -import org.apache.parquet.schema.MessageType; import org.apache.parquet.schema.PrimitiveType; -import org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName; -import org.apache.parquet.schema.Type.Repetition; -import org.apache.parquet.schema.Types; import java.io.IOException; import java.io.InputStream; -import java.util.ArrayList; import java.util.Arrays; -import java.util.Collections; -import java.util.HashMap; -import java.util.HashSet; -import java.util.Iterator; -import java.util.List; -import java.util.Locale; -import java.util.Map; import java.util.Optional; -import java.util.Set; import static io.trino.parquet.ParquetValidationUtils.validateParquet; import static java.lang.Boolean.FALSE; @@ -65,7 +41,6 @@ import static java.lang.Math.min; import static java.lang.Math.toIntExact; import static org.apache.parquet.format.Util.readFileMetaData; -import static org.apache.parquet.format.converter.ParquetMetadataConverterUtil.getLogicalTypeAnnotation; public final class MetadataReader { @@ -116,138 +91,9 @@ public static ParquetMetadata readFooter(ParquetDataSource dataSource, Optional< InputStream metadataStream = buffer.slice(buffer.length() - completeFooterSize, metadataLength).getInput(); FileMetaData fileMetaData = readFileMetaData(metadataStream); - List schema = fileMetaData.getSchema(); - validateParquet(!schema.isEmpty(), "Empty Parquet schema in file: %s", dataSource.getId()); - - MessageType messageType = readParquetSchema(schema); - List blocks = new ArrayList<>(); - List rowGroups = fileMetaData.getRow_groups(); - if (rowGroups != null) { - for (RowGroup rowGroup : rowGroups) { - BlockMetaData blockMetaData = new BlockMetaData(); - blockMetaData.setRowCount(rowGroup.getNum_rows()); - blockMetaData.setTotalByteSize(rowGroup.getTotal_byte_size()); - List columns = rowGroup.getColumns(); - validateParquet(!columns.isEmpty(), "No columns in row group: %s", rowGroup); - String filePath = columns.get(0).getFile_path(); - for (ColumnChunk columnChunk : columns) { - validateParquet( - (filePath == null && columnChunk.getFile_path() == null) - || (filePath != null && filePath.equals(columnChunk.getFile_path())), - "all column chunks of the same row group must be in the same file"); - ColumnMetaData metaData = columnChunk.meta_data; - String[] path = metaData.path_in_schema.stream() - .map(value -> value.toLowerCase(Locale.ENGLISH)) - .toArray(String[]::new); - ColumnPath columnPath = ColumnPath.get(path); - PrimitiveType primitiveType = messageType.getType(columnPath.toArray()).asPrimitiveType(); - ColumnChunkMetaData column = ColumnChunkMetaData.get( - columnPath, - primitiveType, - CompressionCodecName.fromParquet(metaData.codec), - PARQUET_METADATA_CONVERTER.convertEncodingStats(metaData.encoding_stats), - readEncodings(metaData.encodings), - readStats(Optional.ofNullable(fileMetaData.getCreated_by()), Optional.ofNullable(metaData.statistics), primitiveType), - metaData.data_page_offset, - metaData.dictionary_page_offset, - metaData.num_values, - metaData.total_compressed_size, - metaData.total_uncompressed_size); - column.setColumnIndexReference(toColumnIndexReference(columnChunk)); - column.setOffsetIndexReference(toOffsetIndexReference(columnChunk)); - column.setBloomFilterOffset(metaData.bloom_filter_offset); - blockMetaData.addColumn(column); - } - blockMetaData.setPath(filePath); - blocks.add(blockMetaData); - } - } - - Map keyValueMetaData = new HashMap<>(); - List keyValueList = fileMetaData.getKey_value_metadata(); - if (keyValueList != null) { - for (KeyValue keyValue : keyValueList) { - keyValueMetaData.put(keyValue.key, keyValue.value); - } - } - org.apache.parquet.hadoop.metadata.FileMetaData parquetFileMetadata = new org.apache.parquet.hadoop.metadata.FileMetaData( - messageType, - keyValueMetaData, - fileMetaData.getCreated_by()); - validateFileMetadata(dataSource.getId(), parquetFileMetadata, parquetWriteValidation); - return new ParquetMetadata(parquetFileMetadata, blocks); - } - - private static MessageType readParquetSchema(List schema) - { - Iterator schemaIterator = schema.iterator(); - SchemaElement rootSchema = schemaIterator.next(); - Types.MessageTypeBuilder builder = Types.buildMessage(); - readTypeSchema(builder, schemaIterator, rootSchema.getNum_children()); - return builder.named(rootSchema.name); - } - - private static void readTypeSchema(Types.GroupBuilder builder, Iterator schemaIterator, int typeCount) - { - ParquetMetadataConverter parquetMetadataConverter = new ParquetMetadataConverter(); - for (int i = 0; i < typeCount; i++) { - SchemaElement element = schemaIterator.next(); - Types.Builder typeBuilder; - if (element.type == null) { - typeBuilder = builder.group(Repetition.valueOf(element.repetition_type.name())); - readTypeSchema((Types.GroupBuilder) typeBuilder, schemaIterator, element.num_children); - } - else { - Types.PrimitiveBuilder primitiveBuilder = builder.primitive(getTypeName(element.type), Repetition.valueOf(element.repetition_type.name())); - if (element.isSetType_length()) { - primitiveBuilder.length(element.type_length); - } - if (element.isSetPrecision()) { - primitiveBuilder.precision(element.precision); - } - if (element.isSetScale()) { - primitiveBuilder.scale(element.scale); - } - typeBuilder = primitiveBuilder; - } - - // Reading of element.logicalType and element.converted_type corresponds to parquet-mr's code at - // https://github.com/apache/parquet-mr/blob/apache-parquet-1.12.0/parquet-hadoop/src/main/java/org/apache/parquet/format/converter/ParquetMetadataConverter.java#L1568-L1582 - LogicalTypeAnnotation annotationFromLogicalType = null; - if (element.isSetLogicalType()) { - annotationFromLogicalType = getLogicalTypeAnnotation(parquetMetadataConverter, element.logicalType); - typeBuilder.as(annotationFromLogicalType); - } - if (element.isSetConverted_type()) { - LogicalTypeAnnotation annotationFromConvertedType = getLogicalTypeAnnotation(parquetMetadataConverter, element.converted_type, element); - if (annotationFromLogicalType != null) { - // Both element.logicalType and element.converted_type set - if (annotationFromLogicalType.toOriginalType() == annotationFromConvertedType.toOriginalType()) { - // element.converted_type matches element.logicalType, even though annotationFromLogicalType may differ from annotationFromConvertedType - // Following parquet-mr behavior, we favor LogicalTypeAnnotation derived from element.logicalType, as potentially containing more information. - } - else { - // Following parquet-mr behavior, issue warning and let converted_type take precedence. - log.warn("Converted type and logical type metadata map to different OriginalType (convertedType: %s, logical type: %s). Using value in converted type.", - element.converted_type, element.logicalType); - // parquet-mr reads only OriginalType from converted_type. We retain full LogicalTypeAnnotation - // 1. for compatibility, as previous Trino reader code would read LogicalTypeAnnotation from element.converted_type and some additional fields. - // 2. so that we override LogicalTypeAnnotation annotation read from element.logicalType in case of mismatch detected. - typeBuilder.as(annotationFromConvertedType); - } - } - else { - // parquet-mr reads only OriginalType from converted_type. We retain full LogicalTypeAnnotation for compatibility, as previous - // Trino reader code would read LogicalTypeAnnotation from element.converted_type and some additional fields. - typeBuilder.as(annotationFromConvertedType); - } - } - - if (element.isSetField_id()) { - typeBuilder.id(element.field_id); - } - typeBuilder.named(element.name.toLowerCase(Locale.ENGLISH)); - } + ParquetMetadata parquetMetadata = new ParquetMetadata(fileMetaData, dataSource.getId()); + validateFileMetadata(dataSource.getId(), parquetMetadata.getFileMetaData(), parquetWriteValidation); + return parquetMetadata; } public static org.apache.parquet.column.statistics.Statistics readStats(Optional fileCreatedBy, Optional statisticsFromFile, PrimitiveType type) @@ -343,55 +189,7 @@ private static int commonPrefix(byte[] a, byte[] b) return commonPrefixLength; } - private static Set readEncodings(List encodings) - { - Set columnEncodings = new HashSet<>(); - for (Encoding encoding : encodings) { - columnEncodings.add(org.apache.parquet.column.Encoding.valueOf(encoding.name())); - } - return Collections.unmodifiableSet(columnEncodings); - } - - private static PrimitiveTypeName getTypeName(Type type) - { - switch (type) { - case BYTE_ARRAY: - return PrimitiveTypeName.BINARY; - case INT64: - return PrimitiveTypeName.INT64; - case INT32: - return PrimitiveTypeName.INT32; - case BOOLEAN: - return PrimitiveTypeName.BOOLEAN; - case FLOAT: - return PrimitiveTypeName.FLOAT; - case DOUBLE: - return PrimitiveTypeName.DOUBLE; - case INT96: - return PrimitiveTypeName.INT96; - case FIXED_LEN_BYTE_ARRAY: - return PrimitiveTypeName.FIXED_LEN_BYTE_ARRAY; - } - throw new IllegalArgumentException("Unknown type " + type); - } - - private static IndexReference toColumnIndexReference(ColumnChunk columnChunk) - { - if (columnChunk.isSetColumn_index_offset() && columnChunk.isSetColumn_index_length()) { - return new IndexReference(columnChunk.getColumn_index_offset(), columnChunk.getColumn_index_length()); - } - return null; - } - - private static IndexReference toOffsetIndexReference(ColumnChunk columnChunk) - { - if (columnChunk.isSetOffset_index_offset() && columnChunk.isSetOffset_index_length()) { - return new IndexReference(columnChunk.getOffset_index_offset(), columnChunk.getOffset_index_length()); - } - return null; - } - - private static void validateFileMetadata(ParquetDataSourceId dataSourceId, org.apache.parquet.hadoop.metadata.FileMetaData fileMetaData, Optional parquetWriteValidation) + private static void validateFileMetadata(ParquetDataSourceId dataSourceId, FileMetadata fileMetaData, Optional parquetWriteValidation) throws ParquetCorruptionException { if (parquetWriteValidation.isEmpty()) { diff --git a/lib/trino-parquet/src/main/java/io/trino/parquet/reader/PageReader.java b/lib/trino-parquet/src/main/java/io/trino/parquet/reader/PageReader.java index 1c19115cffdd..d8ec35c52fbe 100644 --- a/lib/trino-parquet/src/main/java/io/trino/parquet/reader/PageReader.java +++ b/lib/trino-parquet/src/main/java/io/trino/parquet/reader/PageReader.java @@ -21,11 +21,12 @@ import io.trino.parquet.DataPageV2; import io.trino.parquet.DictionaryPage; import io.trino.parquet.Page; +import io.trino.parquet.ParquetDataSourceId; +import io.trino.parquet.metadata.ColumnChunkMetadata; import jakarta.annotation.Nullable; import org.apache.parquet.column.ColumnDescriptor; import org.apache.parquet.column.statistics.Statistics; import org.apache.parquet.format.CompressionCodec; -import org.apache.parquet.hadoop.metadata.ColumnChunkMetaData; import org.apache.parquet.internal.column.columnindex.OffsetIndex; import java.io.IOException; @@ -36,9 +37,11 @@ import static com.google.common.base.Preconditions.checkState; import static io.trino.parquet.ParquetCompressionUtils.decompress; import static io.trino.parquet.ParquetReaderUtils.isOnlyDictionaryEncodingPages; +import static java.util.Objects.requireNonNull; public final class PageReader { + private final ParquetDataSourceId dataSourceId; private final CompressionCodec codec; private final boolean hasOnlyDictionaryEncodedPages; private final boolean hasNoNulls; @@ -48,8 +51,9 @@ public final class PageReader private int dataPageReadCount; public static PageReader createPageReader( + ParquetDataSourceId dataSourceId, ChunkedInputStream columnChunk, - ColumnChunkMetaData metadata, + ColumnChunkMetadata metadata, ColumnDescriptor columnDescriptor, @Nullable OffsetIndex offsetIndex, Optional fileCreatedBy) @@ -61,21 +65,30 @@ public static PageReader createPageReader( boolean hasNoNulls = columnStatistics != null && columnStatistics.getNumNulls() == 0; boolean hasOnlyDictionaryEncodedPages = isOnlyDictionaryEncodingPages(metadata); ParquetColumnChunkIterator compressedPages = new ParquetColumnChunkIterator( + dataSourceId, fileCreatedBy, columnDescriptor, metadata, columnChunk, offsetIndex); - return new PageReader(metadata.getCodec().getParquetCompressionCodec(), compressedPages, hasOnlyDictionaryEncodedPages, hasNoNulls); + + return new PageReader( + dataSourceId, + metadata.getCodec().getParquetCompressionCodec(), + compressedPages, + hasOnlyDictionaryEncodedPages, + hasNoNulls); } @VisibleForTesting public PageReader( + ParquetDataSourceId dataSourceId, CompressionCodec codec, Iterator compressedPages, boolean hasOnlyDictionaryEncodedPages, boolean hasNoNulls) { + this.dataSourceId = requireNonNull(dataSourceId, "dataSourceId is null"); this.codec = codec; this.compressedPages = Iterators.peekingIterator(compressedPages); this.hasOnlyDictionaryEncodedPages = hasOnlyDictionaryEncodedPages; @@ -106,7 +119,7 @@ public DataPage readPage() return dataPageV1; } return new DataPageV1( - decompress(codec, dataPageV1.getSlice(), dataPageV1.getUncompressedSize()), + decompress(dataSourceId, codec, dataPageV1.getSlice(), dataPageV1.getUncompressedSize()), dataPageV1.getValueCount(), dataPageV1.getUncompressedSize(), dataPageV1.getFirstRowIndex(), @@ -128,7 +141,7 @@ public DataPage readPage() dataPageV2.getRepetitionLevels(), dataPageV2.getDefinitionLevels(), dataPageV2.getDataEncoding(), - decompress(codec, dataPageV2.getSlice(), uncompressedSize), + decompress(dataSourceId, codec, dataPageV2.getSlice(), uncompressedSize), dataPageV2.getUncompressedSize(), dataPageV2.getFirstRowIndex(), dataPageV2.getStatistics(), @@ -150,7 +163,7 @@ public DictionaryPage readDictionaryPage() try { DictionaryPage compressedDictionaryPage = (DictionaryPage) compressedPages.next(); return new DictionaryPage( - decompress(codec, compressedDictionaryPage.getSlice(), compressedDictionaryPage.getUncompressedSize()), + decompress(dataSourceId, codec, compressedDictionaryPage.getSlice(), compressedDictionaryPage.getUncompressedSize()), compressedDictionaryPage.getDictionarySize(), compressedDictionaryPage.getEncoding()); } diff --git a/lib/trino-parquet/src/main/java/io/trino/parquet/reader/ParquetColumnChunkIterator.java b/lib/trino-parquet/src/main/java/io/trino/parquet/reader/ParquetColumnChunkIterator.java index 1a14d579bc59..235c1b2d3d76 100644 --- a/lib/trino-parquet/src/main/java/io/trino/parquet/reader/ParquetColumnChunkIterator.java +++ b/lib/trino-parquet/src/main/java/io/trino/parquet/reader/ParquetColumnChunkIterator.java @@ -18,6 +18,8 @@ import io.trino.parquet.DictionaryPage; import io.trino.parquet.Page; import io.trino.parquet.ParquetCorruptionException; +import io.trino.parquet.ParquetDataSourceId; +import io.trino.parquet.metadata.ColumnChunkMetadata; import jakarta.annotation.Nullable; import org.apache.parquet.column.ColumnDescriptor; import org.apache.parquet.column.Encoding; @@ -26,7 +28,6 @@ import org.apache.parquet.format.DictionaryPageHeader; import org.apache.parquet.format.PageHeader; import org.apache.parquet.format.Util; -import org.apache.parquet.hadoop.metadata.ColumnChunkMetaData; import org.apache.parquet.internal.column.columnindex.OffsetIndex; import java.io.IOException; @@ -41,9 +42,10 @@ public final class ParquetColumnChunkIterator implements Iterator { + private final ParquetDataSourceId dataSourceId; private final Optional fileCreatedBy; private final ColumnDescriptor descriptor; - private final ColumnChunkMetaData metadata; + private final ColumnChunkMetadata metadata; private final ChunkedInputStream input; private final OffsetIndex offsetIndex; @@ -51,12 +53,14 @@ public final class ParquetColumnChunkIterator private int dataPageCount; public ParquetColumnChunkIterator( + ParquetDataSourceId dataSourceId, Optional fileCreatedBy, ColumnDescriptor descriptor, - ColumnChunkMetaData metadata, + ColumnChunkMetadata metadata, ChunkedInputStream input, @Nullable OffsetIndex offsetIndex) { + this.dataSourceId = requireNonNull(dataSourceId, "dataSourceId is null"); this.fileCreatedBy = requireNonNull(fileCreatedBy, "fileCreatedBy is null"); this.descriptor = requireNonNull(descriptor, "descriptor is null"); this.metadata = requireNonNull(metadata, "metadata is null"); @@ -83,7 +87,7 @@ public Page next() switch (pageHeader.type) { case DICTIONARY_PAGE: if (dataPageCount != 0) { - throw new ParquetCorruptionException("Column (%s) has a dictionary page after the first position in column chunk", descriptor); + throw new ParquetCorruptionException(dataSourceId, "Column (%s) has a dictionary page after the first position in column chunk", descriptor); } result = readDictionaryPage(pageHeader, pageHeader.getUncompressed_page_size(), pageHeader.getCompressed_page_size()); break; diff --git a/lib/trino-parquet/src/main/java/io/trino/parquet/reader/ParquetReader.java b/lib/trino-parquet/src/main/java/io/trino/parquet/reader/ParquetReader.java index d8c25025e0e3..6ab03775d17d 100644 --- a/lib/trino-parquet/src/main/java/io/trino/parquet/reader/ParquetReader.java +++ b/lib/trino-parquet/src/main/java/io/trino/parquet/reader/ParquetReader.java @@ -17,9 +17,13 @@ import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.common.collect.ListMultimap; +import com.google.errorprone.annotations.CheckReturnValue; +import com.google.errorprone.annotations.FormatMethod; import io.airlift.log.Logger; +import io.airlift.slice.Slice; import io.trino.memory.context.AggregatedMemoryContext; import io.trino.parquet.ChunkKey; +import io.trino.parquet.Column; import io.trino.parquet.DiskRange; import io.trino.parquet.Field; import io.trino.parquet.GroupField; @@ -28,12 +32,19 @@ import io.trino.parquet.ParquetReaderOptions; import io.trino.parquet.ParquetWriteValidation; import io.trino.parquet.PrimitiveField; +import io.trino.parquet.VariantField; +import io.trino.parquet.metadata.ColumnChunkMetadata; +import io.trino.parquet.metadata.PrunedBlockMetadata; import io.trino.parquet.predicate.TupleDomainParquetPredicate; import io.trino.parquet.reader.FilteredOffsetIndex.OffsetRange; +import io.trino.parquet.spark.Variant; import io.trino.plugin.base.metrics.LongCount; import io.trino.spi.Page; import io.trino.spi.block.ArrayBlock; import io.trino.spi.block.Block; +import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.DictionaryBlock; +import io.trino.spi.block.LongArrayBlock; import io.trino.spi.block.RowBlock; import io.trino.spi.block.RunLengthEncodedBlock; import io.trino.spi.metrics.Metric; @@ -42,13 +53,10 @@ import io.trino.spi.type.MapType; import io.trino.spi.type.RowType; import io.trino.spi.type.Type; -import io.trino.spi.type.TypeSignatureParameter; import jakarta.annotation.Nullable; import org.apache.parquet.column.ColumnDescriptor; import org.apache.parquet.filter2.compat.FilterCompat; import org.apache.parquet.filter2.predicate.FilterPredicate; -import org.apache.parquet.hadoop.metadata.BlockMetaData; -import org.apache.parquet.hadoop.metadata.ColumnChunkMetaData; import org.apache.parquet.hadoop.metadata.ColumnPath; import org.apache.parquet.internal.column.columnindex.OffsetIndex; import org.apache.parquet.internal.filter2.columnindex.ColumnIndexFilter; @@ -57,6 +65,8 @@ import java.io.Closeable; import java.io.IOException; +import java.time.ZoneId; +import java.util.Arrays; import java.util.HashMap; import java.util.List; import java.util.Map; @@ -65,7 +75,9 @@ import java.util.function.Function; import static com.google.common.base.Preconditions.checkArgument; +import static com.google.common.collect.ImmutableList.toImmutableList; import static com.google.common.collect.ImmutableSet.toImmutableSet; +import static io.airlift.slice.Slices.utf8Slice; import static io.trino.parquet.ParquetValidationUtils.validateParquet; import static io.trino.parquet.ParquetWriteValidation.StatisticsValidation; import static io.trino.parquet.ParquetWriteValidation.StatisticsValidation.createStatisticsValidationBuilder; @@ -73,11 +85,13 @@ import static io.trino.parquet.ParquetWriteValidation.WriteChecksumBuilder.createWriteChecksumBuilder; import static io.trino.parquet.reader.ListColumnReader.calculateCollectionOffsets; import static io.trino.parquet.reader.PageReader.createPageReader; +import static io.trino.spi.type.VarbinaryType.VARBINARY; +import static io.trino.spi.type.VarcharType.VARCHAR; import static java.lang.Math.max; import static java.lang.Math.min; import static java.lang.Math.toIntExact; import static java.lang.String.format; -import static java.util.Collections.nCopies; +import static java.util.Objects.checkIndex; import static java.util.Objects.requireNonNull; public class ParquetReader @@ -91,16 +105,17 @@ public class ParquetReader public static final String COLUMN_INDEX_ROWS_FILTERED = "ParquetColumnIndexRowsFiltered"; private final Optional fileCreatedBy; - private final List blocks; - private final List firstRowsOfBlocks; - private final List columnFields; + private final List rowGroups; + private final List columnFields; + private final boolean appendRowNumberColumn; private final List primitiveFields; private final ParquetDataSource dataSource; + private final ZoneId zoneId; private final ColumnReaderFactory columnReaderFactory; private final AggregatedMemoryContext memoryContext; private int currentRowGroup = -1; - private BlockMetaData currentBlockMetadata; + private PrunedBlockMetadata currentBlockMetadata; private long currentGroupRowCount; /** * Index in the Parquet file of the first row of the current group @@ -120,11 +135,11 @@ public class ParquetReader private AggregatedMemoryContext currentRowGroupMemoryContext; private final Map chunkReaders; - private final List> columnIndexStore; private final Optional writeValidation; private final Optional writeChecksumBuilder; private final Optional rowGroupStatisticsValidation; private final FilteredRowRanges[] blockRowRanges; + private final Function exceptionTransform; private final ParquetBlockFactory blockFactory; private final Map> codecMetrics; @@ -132,41 +147,26 @@ public class ParquetReader public ParquetReader( Optional fileCreatedBy, - List columnFields, - List blocks, - List firstRowsOfBlocks, - ParquetDataSource dataSource, - DateTimeZone timeZone, - AggregatedMemoryContext memoryContext, - ParquetReaderOptions options, - Function exceptionTransform) - throws IOException - { - this(fileCreatedBy, columnFields, blocks, firstRowsOfBlocks, dataSource, timeZone, memoryContext, options, exceptionTransform, Optional.empty(), nCopies(blocks.size(), Optional.empty()), Optional.empty()); - } - - public ParquetReader( - Optional fileCreatedBy, - List columnFields, - List blocks, - List firstRowsOfBlocks, + List columnFields, + boolean appendRowNumberColumn, + List rowGroups, ParquetDataSource dataSource, DateTimeZone timeZone, AggregatedMemoryContext memoryContext, ParquetReaderOptions options, Function exceptionTransform, Optional parquetPredicate, - List> columnIndexStore, Optional writeValidation) throws IOException { this.fileCreatedBy = requireNonNull(fileCreatedBy, "fileCreatedBy is null"); requireNonNull(columnFields, "columnFields is null"); this.columnFields = ImmutableList.copyOf(columnFields); - this.primitiveFields = getPrimitiveFields(columnFields); - this.blocks = requireNonNull(blocks, "blocks is null"); - this.firstRowsOfBlocks = requireNonNull(firstRowsOfBlocks, "firstRowsOfBlocks is null"); + this.appendRowNumberColumn = appendRowNumberColumn; + this.primitiveFields = getPrimitiveFields(columnFields.stream().map(Column::field).collect(toImmutableList())); + this.rowGroups = requireNonNull(rowGroups, "rowGroups is null"); this.dataSource = requireNonNull(dataSource, "dataSource is null"); + this.zoneId = requireNonNull(timeZone, "timeZone is null").toTimeZone().toZoneId(); this.columnReaderFactory = new ColumnReaderFactory(timeZone, options); this.memoryContext = requireNonNull(memoryContext, "memoryContext is null"); this.currentRowGroupMemoryContext = memoryContext.newAggregatedMemoryContext(); @@ -175,36 +175,35 @@ public ParquetReader( this.columnReaders = new HashMap<>(); this.maxBytesPerCell = new HashMap<>(); - checkArgument(blocks.size() == firstRowsOfBlocks.size(), "elements of firstRowsOfBlocks must correspond to blocks"); - this.writeValidation = requireNonNull(writeValidation, "writeValidation is null"); validateWrite( validation -> fileCreatedBy.equals(Optional.of(validation.getCreatedBy())), "Expected created by %s, found %s", writeValidation.map(ParquetWriteValidation::getCreatedBy), fileCreatedBy); - validateBlockMetadata(blocks); + validateBlockMetadata(rowGroups); this.writeChecksumBuilder = writeValidation.map(validation -> createWriteChecksumBuilder(validation.getTypes())); this.rowGroupStatisticsValidation = writeValidation.map(validation -> createStatisticsValidationBuilder(validation.getTypes())); requireNonNull(parquetPredicate, "parquetPredicate is null"); - this.columnIndexStore = requireNonNull(columnIndexStore, "columnIndexStore is null"); Optional filter = Optional.empty(); if (parquetPredicate.isPresent() && options.isUseColumnIndex()) { filter = parquetPredicate.get().toParquetFilter(timeZone); } - this.blockRowRanges = calculateFilteredRowRanges(blocks, filter, columnIndexStore, primitiveFields); + this.blockRowRanges = calculateFilteredRowRanges(rowGroups, filter, primitiveFields); + this.exceptionTransform = exceptionTransform; this.blockFactory = new ParquetBlockFactory(exceptionTransform); ListMultimap ranges = ArrayListMultimap.create(); Map codecMetrics = new HashMap<>(); - for (int rowGroup = 0; rowGroup < blocks.size(); rowGroup++) { - BlockMetaData metadata = blocks.get(rowGroup); + for (int rowGroup = 0; rowGroup < rowGroups.size(); rowGroup++) { + PrunedBlockMetadata blockMetadata = rowGroups.get(rowGroup).prunedBlockMetadata(); + long rowGroupRowCount = blockMetadata.getRowCount(); for (PrimitiveField field : primitiveFields) { int columnId = field.getId(); - ColumnChunkMetaData chunkMetadata = getColumnChunkMetaData(metadata, field.getDescriptor()); + ColumnChunkMetadata chunkMetadata = blockMetadata.getColumnChunkMetaData(field.getDescriptor()); ColumnPath columnPath = chunkMetadata.getPath(); - long rowGroupRowCount = metadata.getRowCount(); + long startingPosition = chunkMetadata.getStartingPos(); long totalLength = chunkMetadata.getTotalSize(); long totalDataSize = 0; @@ -266,12 +265,27 @@ public Page nextPage() } // create a lazy page blockFactory.nextPage(); - Block[] blocks = new Block[columnFields.size()]; - for (int channel = 0; channel < columnFields.size(); channel++) { - Field field = columnFields.get(channel); - blocks[channel] = blockFactory.createBlock(batchSize, () -> readBlock(field)); + Block[] blocks = new Block[columnFields.size() + (appendRowNumberColumn ? 1 : 0)]; + int rowNumberColumnIndex = appendRowNumberColumn ? columnFields.size() : -1; + SelectedPositions selectedPositions = new SelectedPositions(batchSize, null); + for (int channel = 0; channel < blocks.length; channel++) { + Block block; + if (channel == rowNumberColumnIndex) { + block = selectedPositions.createRowNumberBlock(lastBatchStartRow()); + } + else { + try { + // todo use selected positions to improve read performance + block = readBlock(columnFields.get(channel).field()); + } + catch (IOException e) { + throw exceptionTransform.apply(e); + } + block = selectedPositions.apply(block); + } + blocks[channel] = block; } - Page page = new Page(batchSize, blocks); + Page page = new Page(selectedPositions.positionCount(), blocks); validateWritePageChecksum(page); return page; } @@ -314,11 +328,12 @@ private boolean advanceToNextRowGroup() } currentRowGroup++; - if (currentRowGroup == blocks.size()) { + if (currentRowGroup == rowGroups.size()) { return false; } - currentBlockMetadata = blocks.get(currentRowGroup); - firstRowIndexInGroup = firstRowsOfBlocks.get(currentRowGroup); + RowGroupInfo rowGroupInfo = rowGroups.get(currentRowGroup); + currentBlockMetadata = rowGroupInfo.prunedBlockMetadata(); + firstRowIndexInGroup = rowGroupInfo.fileRowOffset(); currentGroupRowCount = currentBlockMetadata.getRowCount(); FilteredRowRanges currentGroupRowRanges = blockRowRanges[currentRowGroup]; log.debug("advanceToNextRowGroup dataSource %s, currentRowGroup %d, rowRanges %s, currentBlockMetadata %s", dataSource.getId(), currentRowGroup, currentGroupRowRanges, currentBlockMetadata); @@ -326,7 +341,9 @@ private boolean advanceToNextRowGroup() long rowCount = currentGroupRowRanges.getRowCount(); columnIndexRowsFiltered += currentGroupRowCount - rowCount; if (rowCount == 0) { - return false; + // Filters on multiple columns with page indexes may yield non-overlapping row ranges and eliminate the entire row group. + // Advance to next row group to ensure that we don't return a null Page and close the page source before all row groups are processed + return advanceToNextRowGroup(); } currentGroupRowCount = rowCount; } @@ -341,20 +358,45 @@ private void freeCurrentRowGroupBuffers() return; } - for (int column = 0; column < primitiveFields.size(); column++) { - ChunkedInputStream chunkedStream = chunkReaders.get(new ChunkKey(column, currentRowGroup)); + for (PrimitiveField field : primitiveFields) { + ChunkedInputStream chunkedStream = chunkReaders.get(new ChunkKey(field.getId(), currentRowGroup)); if (chunkedStream != null) { chunkedStream.close(); } } } + private ColumnChunk readVariant(VariantField field) + throws IOException + { + ColumnChunk valueChunk = readColumnChunk(field.getValue()); + + BlockBuilder variantBlock = VARCHAR.createBlockBuilder(null, 1); + if (valueChunk.getBlock().getPositionCount() == 0) { + variantBlock.appendNull(); + } + else { + ColumnChunk metadataChunk = readColumnChunk(field.getMetadata()); + Slice value = VARBINARY.getSlice(valueChunk.getBlock(), 0); + Slice metadata = VARBINARY.getSlice(metadataChunk.getBlock(), 0); + VARCHAR.writeSlice(variantBlock, value); + VARCHAR.writeSlice(variantBlock, metadata); + Variant variant = new Variant(value.byteArray(), metadata.byteArray()); + VARCHAR.writeSlice(variantBlock, utf8Slice(variant.toJson(zoneId))); + } + return new ColumnChunk(variantBlock.build(), valueChunk.getDefinitionLevels(), valueChunk.getRepetitionLevels()); + } + private ColumnChunk readArray(GroupField field) throws IOException { List parameters = field.getType().getTypeParameters(); checkArgument(parameters.size() == 1, "Arrays must have a single type parameter, found %s", parameters.size()); - Field elementField = field.getChildren().get(0).get(); + Optional children = field.getChildren().get(0); + if (children.isEmpty()) { + return new ColumnChunk(field.getType().createNullBlock(), new int[] {}, new int[] {}); + } + Field elementField = children.get(); ColumnChunk columnChunk = readColumnChunk(elementField); ListColumnReader.BlockPositions collectionPositions = calculateCollectionOffsets(field, columnChunk.getDefinitionLevels(), columnChunk.getRepetitionLevels()); @@ -372,7 +414,8 @@ private ColumnChunk readMap(GroupField field) ColumnChunk columnChunk = readColumnChunk(field.getChildren().get(0).get()); blocks[0] = columnChunk.getBlock(); - blocks[1] = readColumnChunk(field.getChildren().get(1).get()).getBlock(); + Optional valueField = field.getChildren().get(1); + blocks[1] = valueField.isPresent() ? readColumnChunk(valueField.get()).getBlock() : parameters.get(1).createNullBlock(); ListColumnReader.BlockPositions collectionPositions = calculateCollectionOffsets(field, columnChunk.getDefinitionLevels(), columnChunk.getRepetitionLevels()); Block mapBlock = ((MapType) field.getType()).createBlockFromKeyValue(collectionPositions.isNull(), collectionPositions.offsets(), blocks[0], blocks[1]); return new ColumnChunk(mapBlock, columnChunk.getDefinitionLevels(), columnChunk.getRepetitionLevels()); @@ -381,31 +424,73 @@ private ColumnChunk readMap(GroupField field) private ColumnChunk readStruct(GroupField field) throws IOException { - List fields = field.getType().getTypeSignature().getParameters(); - Block[] blocks = new Block[fields.size()]; + Block[] blocks = new Block[field.getType().getTypeParameters().size()]; ColumnChunk columnChunk = null; List> parameters = field.getChildren(); - for (int i = 0; i < fields.size(); i++) { + for (int i = 0; i < blocks.length; i++) { Optional parameter = parameters.get(i); if (parameter.isPresent()) { columnChunk = readColumnChunk(parameter.get()); blocks[i] = columnChunk.getBlock(); } } - for (int i = 0; i < fields.size(); i++) { + + if (columnChunk == null) { + throw new ParquetCorruptionException(dataSource.getId(), "Struct field does not have any children: %s", field); + } + + StructColumnReader.RowBlockPositions structIsNull = StructColumnReader.calculateStructOffsets(field, columnChunk.getDefinitionLevels(), columnChunk.getRepetitionLevels()); + Optional isNull = structIsNull.isNull(); + for (int i = 0; i < blocks.length; i++) { if (blocks[i] == null) { - blocks[i] = RunLengthEncodedBlock.create(field.getType().getTypeParameters().get(i), null, columnChunk.getBlock().getPositionCount()); + blocks[i] = RunLengthEncodedBlock.create(field.getType().getTypeParameters().get(i), null, structIsNull.positionsCount()); + } + else if (isNull.isPresent()) { + blocks[i] = toNotNullSupressedBlock(structIsNull.positionsCount(), isNull.get(), blocks[i]); } } - StructColumnReader.RowBlockPositions structIsNull = StructColumnReader.calculateStructOffsets(field, columnChunk.getDefinitionLevels(), columnChunk.getRepetitionLevels()); - Block rowBlock = RowBlock.fromFieldBlocks(structIsNull.positionsCount(), structIsNull.isNull(), blocks); + Block rowBlock = RowBlock.fromNotNullSuppressedFieldBlocks(structIsNull.positionsCount(), structIsNull.isNull(), blocks); return new ColumnChunk(rowBlock, columnChunk.getDefinitionLevels(), columnChunk.getRepetitionLevels()); } + private static Block toNotNullSupressedBlock(int positionCount, boolean[] rowIsNull, Block fieldBlock) + { + // find a existing position in the block that is null + int nullIndex = -1; + if (fieldBlock.mayHaveNull()) { + for (int position = 0; position < fieldBlock.getPositionCount(); position++) { + if (fieldBlock.isNull(position)) { + nullIndex = position; + break; + } + } + } + // if there are no null positions, append a null to the end of the block + if (nullIndex == -1) { + fieldBlock = fieldBlock.getLoadedBlock(); + nullIndex = fieldBlock.getPositionCount(); + fieldBlock = fieldBlock.copyWithAppendedNull(); + } + + // create a dictionary that maps null positions to the null index + int[] dictionaryIds = new int[positionCount]; + int nullSuppressedPosition = 0; + for (int position = 0; position < positionCount; position++) { + if (rowIsNull[position]) { + dictionaryIds[position] = nullIndex; + } + else { + dictionaryIds[position] = nullSuppressedPosition; + nullSuppressedPosition++; + } + } + return DictionaryBlock.create(positionCount, fieldBlock, dictionaryIds); + } + @Nullable private FilteredOffsetIndex getFilteredOffsetIndex(FilteredRowRanges rowRanges, int rowGroup, long rowGroupRowCount, ColumnPath columnPath) { - Optional rowGroupColumnIndexStore = this.columnIndexStore.get(rowGroup); + Optional rowGroupColumnIndexStore = this.rowGroups.get(rowGroup).columnIndexStore(); if (rowGroupColumnIndexStore.isEmpty()) { return null; } @@ -424,8 +509,8 @@ private ColumnChunk readPrimitive(PrimitiveField field) int fieldId = field.getId(); ColumnReader columnReader = columnReaders.get(fieldId); if (!columnReader.hasPageReader()) { - validateParquet(currentBlockMetadata.getRowCount() > 0, "Row group has 0 rows"); - ColumnChunkMetaData metadata = getColumnChunkMetaData(currentBlockMetadata, columnDescriptor); + validateParquet(currentBlockMetadata.getRowCount() > 0, dataSource.getId(), "Row group has 0 rows"); + ColumnChunkMetadata metadata = currentBlockMetadata.getColumnChunkMetaData(columnDescriptor); FilteredRowRanges rowRanges = blockRowRanges[currentRowGroup]; OffsetIndex offsetIndex = null; if (rowRanges != null) { @@ -433,7 +518,7 @@ private ColumnChunk readPrimitive(PrimitiveField field) } ChunkedInputStream columnChunkInputStream = chunkReaders.get(new ChunkKey(fieldId, currentRowGroup)); columnReader.setPageReader( - createPageReader(columnChunkInputStream, metadata, columnDescriptor, offsetIndex, fileCreatedBy), + createPageReader(dataSource.getId(), columnChunkInputStream, metadata, columnDescriptor, offsetIndex, fileCreatedBy), Optional.ofNullable(rowRanges)); } ColumnChunk columnChunk = columnReader.readPrimitive(); @@ -450,6 +535,11 @@ private ColumnChunk readPrimitive(PrimitiveField field) return columnChunk; } + public List getColumnFields() + { + return columnFields; + } + public Metrics getMetrics() { ImmutableMap.Builder> metrics = ImmutableMap.>builder() @@ -461,17 +551,6 @@ public Metrics getMetrics() return new Metrics(metrics.buildOrThrow()); } - private ColumnChunkMetaData getColumnChunkMetaData(BlockMetaData blockMetaData, ColumnDescriptor columnDescriptor) - throws IOException - { - for (ColumnChunkMetaData metadata : blockMetaData.getColumns()) { - if (metadata.getPath().equals(ColumnPath.get(columnDescriptor.getPath()))) { - return metadata; - } - } - throw new ParquetCorruptionException("Metadata is missing for column: %s", columnDescriptor); - } - private void initializeColumnReaders() { for (PrimitiveField field : primitiveFields) { @@ -499,6 +578,10 @@ else if (field instanceof GroupField groupField) { .flatMap(Optional::stream) .forEach(child -> parseField(child, primitiveFields)); } + else if (field instanceof VariantField variantField) { + parseField(variantField.getValue(), primitiveFields); + parseField(variantField.getMetadata(), primitiveFields); + } } public Block readBlock(Field field) @@ -511,7 +594,10 @@ private ColumnChunk readColumnChunk(Field field) throws IOException { ColumnChunk columnChunk; - if (field.getType() instanceof RowType) { + if (field instanceof VariantField variantField) { + columnChunk = readVariant(variantField); + } + else if (field.getType() instanceof RowType) { columnChunk = readStruct((GroupField) field); } else if (field.getType() instanceof MapType) { @@ -537,25 +623,24 @@ public AggregatedMemoryContext getMemoryContext() } private static FilteredRowRanges[] calculateFilteredRowRanges( - List blocks, + List rowGroups, Optional filter, - List> columnIndexStore, List primitiveFields) { - FilteredRowRanges[] blockRowRanges = new FilteredRowRanges[blocks.size()]; + FilteredRowRanges[] blockRowRanges = new FilteredRowRanges[rowGroups.size()]; if (filter.isEmpty()) { return blockRowRanges; } Set paths = primitiveFields.stream() .map(field -> ColumnPath.get(field.getDescriptor().getPath())) .collect(toImmutableSet()); - for (int rowGroup = 0; rowGroup < blocks.size(); rowGroup++) { - Optional rowGroupColumnIndexStore = columnIndexStore.get(rowGroup); + for (int rowGroup = 0; rowGroup < rowGroups.size(); rowGroup++) { + RowGroupInfo rowGroupInfo = rowGroups.get(rowGroup); + Optional rowGroupColumnIndexStore = rowGroupInfo.columnIndexStore(); if (rowGroupColumnIndexStore.isEmpty()) { continue; } - BlockMetaData metadata = blocks.get(rowGroup); - long rowGroupRowCount = metadata.getRowCount(); + long rowGroupRowCount = rowGroupInfo.prunedBlockMetadata().getRowCount(); FilteredRowRanges rowRanges = new FilteredRowRanges(ColumnIndexFilter.calculateRowRanges( FilterCompat.get(filter.get()), rowGroupColumnIndexStore.get(), @@ -577,14 +662,16 @@ private void validateWritePageChecksum(Page page) } } - private void validateBlockMetadata(List blockMetaData) + private void validateBlockMetadata(List rowGroups) throws ParquetCorruptionException { if (writeValidation.isPresent()) { - writeValidation.get().validateBlocksMetadata(dataSource.getId(), blockMetaData); + writeValidation.get().validateBlocksMetadata(dataSource.getId(), rowGroups); } } + @SuppressWarnings("FormatStringAnnotation") + @FormatMethod private void validateWrite(java.util.function.Predicate test, String messageFormat, Object... args) throws ParquetCorruptionException { @@ -592,4 +679,43 @@ private void validateWrite(java.util.function.Predicate throw new ParquetCorruptionException(dataSource.getId(), "Write validation failed: " + messageFormat, args); } } + + private record SelectedPositions(int positionCount, @Nullable int[] positions) + { + @CheckReturnValue + public Block apply(Block block) + { + if (positions == null) { + return block; + } + return block.getPositions(positions, 0, positionCount); + } + + public Block createRowNumberBlock(long startRowNumber) + { + long[] rowNumbers = new long[positionCount]; + for (int i = 0; i < positionCount; i++) { + int position = positions == null ? i : positions[i]; + rowNumbers[i] = startRowNumber + position; + } + return new LongArrayBlock(positionCount, Optional.empty(), rowNumbers); + } + + @CheckReturnValue + public SelectedPositions selectPositions(int[] positions, int offset, int size) + { + if (this.positions == null) { + for (int i = 0; i < size; i++) { + checkIndex(offset + i, positionCount); + } + return new SelectedPositions(size, Arrays.copyOfRange(positions, offset, offset + size)); + } + + int[] newPositions = new int[size]; + for (int i = 0; i < size; i++) { + newPositions[i] = this.positions[positions[offset + i]]; + } + return new SelectedPositions(size, newPositions); + } + } } diff --git a/lib/trino-parquet/src/main/java/io/trino/parquet/reader/RowGroupInfo.java b/lib/trino-parquet/src/main/java/io/trino/parquet/reader/RowGroupInfo.java new file mode 100644 index 000000000000..a35ecf379394 --- /dev/null +++ b/lib/trino-parquet/src/main/java/io/trino/parquet/reader/RowGroupInfo.java @@ -0,0 +1,21 @@ +/* + * 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.parquet.reader; + +import io.trino.parquet.metadata.PrunedBlockMetadata; +import org.apache.parquet.internal.filter2.columnindex.ColumnIndexStore; + +import java.util.Optional; + +public record RowGroupInfo(PrunedBlockMetadata prunedBlockMetadata, long fileRowOffset, Optional columnIndexStore) {} diff --git a/lib/trino-parquet/src/main/java/io/trino/parquet/reader/TrinoColumnIndexStore.java b/lib/trino-parquet/src/main/java/io/trino/parquet/reader/TrinoColumnIndexStore.java index 041c4e1250ae..c7f722e85322 100644 --- a/lib/trino-parquet/src/main/java/io/trino/parquet/reader/TrinoColumnIndexStore.java +++ b/lib/trino-parquet/src/main/java/io/trino/parquet/reader/TrinoColumnIndexStore.java @@ -18,26 +18,33 @@ import com.google.common.collect.ListMultimap; import io.trino.parquet.DiskRange; import io.trino.parquet.ParquetDataSource; +import io.trino.parquet.ParquetReaderOptions; +import io.trino.parquet.metadata.BlockMetadata; +import io.trino.parquet.metadata.ColumnChunkMetadata; +import io.trino.parquet.metadata.IndexReference; +import io.trino.spi.predicate.Domain; +import io.trino.spi.predicate.TupleDomain; import jakarta.annotation.Nullable; +import org.apache.parquet.column.ColumnDescriptor; import org.apache.parquet.format.Util; import org.apache.parquet.format.converter.ParquetMetadataConverter; -import org.apache.parquet.hadoop.metadata.BlockMetaData; -import org.apache.parquet.hadoop.metadata.ColumnChunkMetaData; import org.apache.parquet.hadoop.metadata.ColumnPath; import org.apache.parquet.internal.column.columnindex.ColumnIndex; import org.apache.parquet.internal.column.columnindex.OffsetIndex; import org.apache.parquet.internal.filter2.columnindex.ColumnIndexStore; -import org.apache.parquet.internal.hadoop.metadata.IndexReference; import org.apache.parquet.schema.PrimitiveType; import java.io.IOException; import java.io.InputStream; +import java.util.HashSet; import java.util.List; import java.util.Map; +import java.util.Optional; import java.util.Set; import java.util.function.BiFunction; import static com.google.common.collect.ImmutableMap.toImmutableMap; +import static com.google.common.collect.ImmutableSet.toImmutableSet; import static io.trino.memory.context.AggregatedMemoryContext.newSimpleAggregatedMemoryContext; import static java.util.Objects.requireNonNull; @@ -65,7 +72,7 @@ public class TrinoColumnIndexStore */ public TrinoColumnIndexStore( ParquetDataSource dataSource, - BlockMetaData block, + BlockMetadata block, Set columnsRead, Set columnsFiltered) { @@ -76,7 +83,7 @@ public TrinoColumnIndexStore( ImmutableList.Builder columnIndexBuilder = ImmutableList.builderWithExpectedSize(columnsFiltered.size()); ImmutableList.Builder offsetIndexBuilder = ImmutableList.builderWithExpectedSize(columnsRead.size()); - for (ColumnChunkMetaData column : block.getColumns()) { + for (ColumnChunkMetadata column : block.columns()) { ColumnPath path = column.getPath(); if (column.getColumnIndexReference() != null && columnsFiltered.contains(path)) { columnIndexBuilder.add(new ColumnIndexMetadata( @@ -129,6 +136,43 @@ public OffsetIndex getOffsetIndex(ColumnPath column) return offsetIndexStore.get(column); } + public static Optional getColumnIndexStore( + ParquetDataSource dataSource, + BlockMetadata blockMetadata, + Map, ColumnDescriptor> descriptorsByPath, + TupleDomain parquetTupleDomain, + ParquetReaderOptions options) + { + if (!options.isUseColumnIndex() || parquetTupleDomain.isAll() || parquetTupleDomain.isNone()) { + return Optional.empty(); + } + + boolean hasColumnIndex = false; + for (ColumnChunkMetadata column : blockMetadata.columns()) { + if (column.getColumnIndexReference() != null && column.getOffsetIndexReference() != null) { + hasColumnIndex = true; + break; + } + } + + if (!hasColumnIndex) { + return Optional.empty(); + } + + Set columnsReadPaths = new HashSet<>(descriptorsByPath.size()); + for (List path : descriptorsByPath.keySet()) { + columnsReadPaths.add(ColumnPath.get(path.toArray(new String[0]))); + } + + Map parquetDomains = parquetTupleDomain.getDomains() + .orElseThrow(() -> new IllegalStateException("Predicate other than none should have domains")); + Set columnsFilteredPaths = parquetDomains.keySet().stream() + .map(column -> ColumnPath.get(column.getPath())) + .collect(toImmutableSet()); + + return Optional.of(new TrinoColumnIndexStore(dataSource, blockMetadata, columnsReadPaths, columnsFilteredPaths)); + } + private static Map loadIndexes( ParquetDataSource dataSource, List indexMetadata, diff --git a/lib/trino-parquet/src/main/java/io/trino/parquet/reader/decoders/PlainValueDecoders.java b/lib/trino-parquet/src/main/java/io/trino/parquet/reader/decoders/PlainValueDecoders.java index 36f68aea3acc..70072c2b2c86 100644 --- a/lib/trino-parquet/src/main/java/io/trino/parquet/reader/decoders/PlainValueDecoders.java +++ b/lib/trino-parquet/src/main/java/io/trino/parquet/reader/decoders/PlainValueDecoders.java @@ -408,4 +408,39 @@ public void skip(int n) input.skip(n * typeLength); } } + + private abstract static class AbstractBooleanPlainValueDecoder + implements ValueDecoder + { + protected SimpleSliceInputStream input; + // Number of unread bits in the current byte + protected int alreadyReadBits; + // Partly read byte + protected byte partiallyReadByte; + + @Override + public void init(SimpleSliceInputStream input) + { + this.input = requireNonNull(input, "input is null"); + alreadyReadBits = 0; + } + + @Override + public void skip(int n) + { + if (alreadyReadBits != 0) { // Skip the partially read byte + int chunkSize = min(Byte.SIZE - alreadyReadBits, n); + n -= chunkSize; + alreadyReadBits = (alreadyReadBits + chunkSize) % Byte.SIZE; // Set to 0 when full byte reached + } + + // Skip full bytes + input.skip(n / Byte.SIZE); + + if (n % Byte.SIZE != 0) { // Partially skip the last byte + alreadyReadBits = n % Byte.SIZE; + partiallyReadByte = input.readByte(); + } + } + } } diff --git a/lib/trino-parquet/src/main/java/io/trino/parquet/reader/decoders/ValueDecoders.java b/lib/trino-parquet/src/main/java/io/trino/parquet/reader/decoders/ValueDecoders.java index 73d8b3c14502..c20b2d53b0aa 100644 --- a/lib/trino-parquet/src/main/java/io/trino/parquet/reader/decoders/ValueDecoders.java +++ b/lib/trino-parquet/src/main/java/io/trino/parquet/reader/decoders/ValueDecoders.java @@ -452,6 +452,32 @@ public void skip(int n) }; } + public ValueDecoder getInt96ToLongTimestampWithTimeZoneDecoder(ParquetEncoding encoding) + { + checkArgument( + field.getType() instanceof TimestampWithTimeZoneType timestampType && !timestampType.isShort(), + "Trino type %s is not a long timestamp", + field.getType()); + int precision = ((TimestampWithTimeZoneType) field.getType()).getPrecision(); + return new InlineTransformDecoder<>( + getInt96TimestampDecoder(encoding), + (values, offset, length) -> { + for (int i = offset; i < offset + length; i++) { + long epochSeconds = decodeFixed12First(values, i); + int nanosOfSecond = decodeFixed12Second(values, i); + if (precision < 9) { + nanosOfSecond = (int) round(nanosOfSecond, 9 - precision); + } + long utcMillis = epochSeconds * MILLISECONDS_PER_SECOND + (nanosOfSecond / NANOSECONDS_PER_MILLISECOND); + encodeFixed12( + packDateTimeWithZone(utcMillis, UTC_KEY), + (nanosOfSecond % NANOSECONDS_PER_MILLISECOND) * PICOSECONDS_PER_NANOSECOND, + values, + i); + } + }); + } + public ValueDecoder getInt64TimestampMillsToShortTimestampDecoder(ParquetEncoding encoding) { checkArgument( @@ -508,7 +534,46 @@ public ValueDecoder getInt64TimestampMillsToShortTimestampWithTimeZoneDe }); } - public ValueDecoder getInt64TimestampMicrosToShortTimestampDecoder(ParquetEncoding encoding) + public ValueDecoder getInt64TimestampMillisToShortTimestampDecoder(ParquetEncoding encoding, DateTimeZone timeZone) + { + checkArgument( + field.getType() instanceof TimestampType timestampType && timestampType.isShort(), + "Trino type %s is not a short timestamp", + field.getType()); + int precision = ((TimestampType) field.getType()).getPrecision(); + ValueDecoder valueDecoder = getLongDecoder(encoding); + if (precision < 3) { + return new InlineTransformDecoder<>( + valueDecoder, + (values, offset, length) -> { + // decoded values are epochMillis, round to lower precision and convert to epochMicros + for (int i = offset; i < offset + length; i++) { + long epochMillis = round(values[i], 3 - precision); + if (timeZone == DateTimeZone.UTC) { + values[i] = epochMillis * MICROSECONDS_PER_MILLISECOND; + } + else { + values[i] = timeZone.convertUTCToLocal(epochMillis) * MICROSECONDS_PER_MILLISECOND; + } + } + }); + } + return new InlineTransformDecoder<>( + valueDecoder, + (values, offset, length) -> { + // decoded values are epochMillis, convert to epochMicros + for (int i = offset; i < offset + length; i++) { + if (timeZone == DateTimeZone.UTC) { + values[i] = values[i] * MICROSECONDS_PER_MILLISECOND; + } + else { + values[i] = timeZone.convertUTCToLocal(values[i]) * MICROSECONDS_PER_MILLISECOND; + } + } + }); + } + + public ValueDecoder getInt64TimestampMicrosToShortTimestampDecoder(ParquetEncoding encoding, DateTimeZone timeZone) { checkArgument( field.getType() instanceof TimestampType timestampType && timestampType.isShort(), @@ -517,14 +582,32 @@ public ValueDecoder getInt64TimestampMicrosToShortTimestampDecoder(Parqu int precision = ((TimestampType) field.getType()).getPrecision(); ValueDecoder valueDecoder = getLongDecoder(encoding); if (precision == 6) { - return valueDecoder; + if (timeZone == DateTimeZone.UTC) { + return valueDecoder; + } + new InlineTransformDecoder<>( + valueDecoder, + (values, offset, length) -> { + for (int i = offset; i < offset + length; i++) { + long epochMicros = values[i]; + long localMillis = timeZone.convertUTCToLocal(floorDiv(epochMicros, MICROSECONDS_PER_MILLISECOND)); + values[i] = (localMillis * MICROSECONDS_PER_MILLISECOND) + floorMod(epochMicros, MICROSECONDS_PER_MILLISECOND); + } + }); } return new InlineTransformDecoder<>( valueDecoder, (values, offset, length) -> { // decoded values are epochMicros, round to lower precision for (int i = offset; i < offset + length; i++) { - values[i] = round(values[i], 6 - precision); + long epochMicros = round(values[i], 6 - precision); + if (timeZone == DateTimeZone.UTC) { + values[i] = epochMicros; + } + else { + long localMillis = timeZone.convertUTCToLocal(floorDiv(epochMicros, MICROSECONDS_PER_MILLISECOND)); + values[i] = (localMillis * MICROSECONDS_PER_MILLISECOND) + floorMod(epochMicros, MICROSECONDS_PER_MILLISECOND); + } } }); } @@ -546,7 +629,7 @@ public ValueDecoder getInt64TimestampMicrosToShortTimestampWithTimeZoneD }); } - public ValueDecoder getInt64TimestampNanosToShortTimestampDecoder(ParquetEncoding encoding) + public ValueDecoder getInt64TimestampNanosToShortTimestampDecoder(ParquetEncoding encoding, DateTimeZone timeZone) { checkArgument( field.getType() instanceof TimestampType timestampType && timestampType.isShort(), @@ -558,12 +641,19 @@ public ValueDecoder getInt64TimestampNanosToShortTimestampDecoder(Parque (values, offset, length) -> { // decoded values are epochNanos, round to lower precision and convert to epochMicros for (int i = offset; i < offset + length; i++) { - values[i] = round(values[i], 9 - precision) / NANOSECONDS_PER_MICROSECOND; + long epochNanos = round(values[i], 9 - precision); + if (timeZone == DateTimeZone.UTC) { + values[i] = epochNanos / NANOSECONDS_PER_MICROSECOND; + } + else { + long localMillis = timeZone.convertUTCToLocal(floorDiv(epochNanos, NANOSECONDS_PER_MILLISECOND)); + values[i] = (localMillis * MICROSECONDS_PER_MILLISECOND) + floorDiv(floorMod(epochNanos, NANOSECONDS_PER_MILLISECOND), NANOSECONDS_PER_MICROSECOND); + } } }); } - public ValueDecoder getInt64TimestampMillisToLongTimestampDecoder(ParquetEncoding encoding) + public ValueDecoder getInt64TimestampMillisToLongTimestampDecoder(ParquetEncoding encoding, DateTimeZone timeZone) { ValueDecoder delegate = getLongDecoder(encoding); return new ValueDecoder<>() @@ -581,7 +671,12 @@ public void read(int[] values, int offset, int length) delegate.read(buffer, 0, length); // decoded values are epochMillis, convert to epochMicros for (int i = 0; i < length; i++) { - encodeFixed12(buffer[i] * MICROSECONDS_PER_MILLISECOND, 0, values, i + offset); + if (timeZone == DateTimeZone.UTC) { + encodeFixed12(buffer[i] * MICROSECONDS_PER_MILLISECOND, 0, values, i + offset); + } + else { + encodeFixed12(timeZone.convertUTCToLocal(buffer[i]) * MICROSECONDS_PER_MILLISECOND, 0, values, i + offset); + } } } @@ -593,7 +688,7 @@ public void skip(int n) }; } - public ValueDecoder getInt64TimestampMicrosToLongTimestampDecoder(ParquetEncoding encoding) + public ValueDecoder getInt64TimestampMicrosToLongTimestampDecoder(ParquetEncoding encoding, DateTimeZone timeZone) { ValueDecoder delegate = getLongDecoder(encoding); return new ValueDecoder<>() @@ -611,7 +706,17 @@ public void read(int[] values, int offset, int length) delegate.read(buffer, 0, length); // decoded values are epochMicros for (int i = 0; i < length; i++) { - encodeFixed12(buffer[i], 0, values, i + offset); + long epochMicros = buffer[i]; + if (timeZone == DateTimeZone.UTC) { + encodeFixed12(epochMicros, 0, values, i + offset); + } + else { + long localMillis = timeZone.convertUTCToLocal(floorDiv(epochMicros, MICROSECONDS_PER_MILLISECOND)); + encodeFixed12((localMillis * MICROSECONDS_PER_MILLISECOND) + floorMod(epochMicros, MICROSECONDS_PER_MILLISECOND), + 0, + values, + i + offset); + } } } @@ -658,7 +763,7 @@ public void skip(int n) }; } - public ValueDecoder getInt64TimestampNanosToLongTimestampDecoder(ParquetEncoding encoding) + public ValueDecoder getInt64TimestampNanosToLongTimestampDecoder(ParquetEncoding encoding, DateTimeZone timeZone) { ValueDecoder delegate = getLongDecoder(encoding); return new ValueDecoder<>() @@ -674,14 +779,26 @@ public void read(int[] values, int offset, int length) { long[] buffer = new long[length]; delegate.read(buffer, 0, length); - // decoded values are epochNanos, convert to (epochMicros, picosOfMicro) + // decoded values are epochNanos, convert to (epochMicros, picosOfNanos) for (int i = 0; i < length; i++) { long epochNanos = buffer[i]; - encodeFixed12( - floorDiv(epochNanos, NANOSECONDS_PER_MICROSECOND), - floorMod(epochNanos, NANOSECONDS_PER_MICROSECOND) * PICOSECONDS_PER_NANOSECOND, - values, - i + offset); + int picosOfNanos = floorMod(epochNanos, NANOSECONDS_PER_MICROSECOND) * PICOSECONDS_PER_NANOSECOND; + if (timeZone == DateTimeZone.UTC) { + encodeFixed12( + floorDiv(epochNanos, NANOSECONDS_PER_MICROSECOND), + picosOfNanos, + values, + i + offset); + } + else { + long localMillis = timeZone.convertUTCToLocal(floorDiv(epochNanos, NANOSECONDS_PER_MILLISECOND)); + long microsFromNanos = floorMod(epochNanos, NANOSECONDS_PER_MILLISECOND) / NANOSECONDS_PER_MICROSECOND; + encodeFixed12( + (localMillis * MICROSECONDS_PER_MILLISECOND) + microsFromNanos, + picosOfNanos, + values, + i + offset); + } } } diff --git a/lib/trino-parquet/src/main/java/io/trino/parquet/spark/Variant.java b/lib/trino-parquet/src/main/java/io/trino/parquet/spark/Variant.java new file mode 100644 index 000000000000..12b7d6a69817 --- /dev/null +++ b/lib/trino-parquet/src/main/java/io/trino/parquet/spark/Variant.java @@ -0,0 +1,172 @@ +/* + * 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.parquet.spark; + +import com.fasterxml.jackson.core.JsonGenerator; + +import java.io.CharArrayWriter; +import java.io.IOException; +import java.time.Instant; +import java.time.LocalDate; +import java.time.ZoneId; +import java.time.ZoneOffset; +import java.time.format.DateTimeFormatter; +import java.time.format.DateTimeFormatterBuilder; +import java.util.Base64; +import java.util.Locale; + +import static com.google.common.base.Preconditions.checkArgument; +import static io.trino.parquet.spark.VariantUtil.SIZE_LIMIT; +import static io.trino.parquet.spark.VariantUtil.VERSION; +import static io.trino.parquet.spark.VariantUtil.VERSION_MASK; +import static io.trino.parquet.spark.VariantUtil.getBinary; +import static io.trino.parquet.spark.VariantUtil.getBoolean; +import static io.trino.parquet.spark.VariantUtil.getDecimal; +import static io.trino.parquet.spark.VariantUtil.getDouble; +import static io.trino.parquet.spark.VariantUtil.getFloat; +import static io.trino.parquet.spark.VariantUtil.getLong; +import static io.trino.parquet.spark.VariantUtil.getMetadataKey; +import static io.trino.parquet.spark.VariantUtil.getString; +import static io.trino.parquet.spark.VariantUtil.getType; +import static io.trino.parquet.spark.VariantUtil.handleArray; +import static io.trino.parquet.spark.VariantUtil.handleObject; +import static io.trino.parquet.spark.VariantUtil.readUnsigned; +import static io.trino.plugin.base.util.JsonUtils.jsonFactory; +import static java.time.format.DateTimeFormatter.ISO_LOCAL_DATE; +import static java.time.format.DateTimeFormatter.ISO_LOCAL_TIME; +import static java.time.temporal.ChronoUnit.MICROS; + +/** + * Copied from https://github.com/apache/spark/blob/53d65fd12dd9231139188227ef9040d40d759021/common/variant/src/main/java/org/apache/spark/types/variant/Variant.java + * and adjusted the code style. + */ +public final class Variant +{ + private static final DateTimeFormatter TIMESTAMP_NTZ_FORMATTER = new DateTimeFormatterBuilder() + .append(ISO_LOCAL_DATE) + .appendLiteral(' ') + .append(ISO_LOCAL_TIME) + .toFormatter(Locale.US); + + private static final DateTimeFormatter TIMESTAMP_FORMATTER = new DateTimeFormatterBuilder() + .append(TIMESTAMP_NTZ_FORMATTER) + .appendOffset("+HH:MM", "+00:00") + .toFormatter(Locale.US); + + private final byte[] value; + private final byte[] metadata; + // The variant value doesn't use the whole `value` binary, but starts from its `pos` index and + // spans a size of `valueSize(value, pos)`. This design avoids frequent copies of the value binary + // when reading a sub-variant in the array/object element. + private final int position; + + public Variant(byte[] value, byte[] metadata) + { + this(value, metadata, 0); + } + + private Variant(byte[] value, byte[] metadata, int position) + { + this.value = value; + this.metadata = metadata; + this.position = position; + checkArgument(metadata.length >= 1, "metadata must be present"); + checkArgument((metadata[0] & VERSION_MASK) == VERSION, "metadata version must be %s", VERSION); + // Don't attempt to use a Variant larger than 16 MiB. We'll never produce one, and it risks memory instability. + checkArgument(metadata.length <= SIZE_LIMIT, "max metadata size is %s: %s", SIZE_LIMIT, metadata.length); + checkArgument(value.length <= SIZE_LIMIT, "max value size is %s: %s", SIZE_LIMIT, value.length); + } + + // Stringify the variant in JSON format. + public String toJson(ZoneId zoneId) + { + StringBuilder json = new StringBuilder(); + toJsonImpl(value, metadata, position, json, zoneId); + return json.toString(); + } + + private static void toJsonImpl(byte[] value, byte[] metadata, int position, StringBuilder json, ZoneId zoneId) + { + switch (getType(value, position)) { + case NULL -> json.append("null"); + case BOOLEAN -> json.append(getBoolean(value, position)); + case LONG -> json.append(getLong(value, position)); + case FLOAT -> json.append(getFloat(value, position)); + case DOUBLE -> json.append(getDouble(value, position)); + case DECIMAL -> json.append(getDecimal(value, position).toPlainString()); + case STRING -> json.append(escapeJson(getString(value, position))); + case BINARY -> appendQuoted(json, Base64.getEncoder().encodeToString(getBinary(value, position))); + case DATE -> appendQuoted(json, LocalDate.ofEpochDay(getLong(value, position)).toString()); + case TIMESTAMP -> appendQuoted(json, TIMESTAMP_FORMATTER.format(microsToInstant(getLong(value, position)).atZone(zoneId))); + case TIMESTAMP_NTZ -> appendQuoted(json, TIMESTAMP_NTZ_FORMATTER.format(microsToInstant(getLong(value, position)).atZone(ZoneOffset.UTC))); + case ARRAY -> handleArray(value, position, (size, offsetSize, offsetStart, dataStart) -> { + json.append('['); + for (int i = 0; i < size; ++i) { + int offset = readUnsigned(value, offsetStart + offsetSize * i, offsetSize); + int elementPos = dataStart + offset; + if (i != 0) { + json.append(','); + } + toJsonImpl(value, metadata, elementPos, json, zoneId); + } + json.append(']'); + return null; + }); + case OBJECT -> handleObject(value, position, (size, idSize, offsetSize, idStart, offsetStart, dataStart) -> { + json.append('{'); + for (int i = 0; i < size; ++i) { + int id = readUnsigned(value, idStart + idSize * i, idSize); + int offset = readUnsigned(value, offsetStart + offsetSize * i, offsetSize); + int elementPosition = dataStart + offset; + if (i != 0) { + json.append(','); + } + json.append(escapeJson(getMetadataKey(metadata, id))); + json.append(':'); + toJsonImpl(value, metadata, elementPosition, json, zoneId); + } + json.append('}'); + return null; + }); + } + } + + private static Instant microsToInstant(long timestamp) + { + return Instant.EPOCH.plus(timestamp, MICROS); + } + + // A simplified and more performant version of `sb.append(escapeJson(value))`. It is used when we + // know `value` doesn't contain any special character that needs escaping. + private static void appendQuoted(StringBuilder json, String value) + { + json.append('"').append(value).append('"'); + } + + // Escape a string so that it can be pasted into JSON structure. + // For example, if `str` only contains a new-line character, then the result content is "\n" + // (4 characters). + private static String escapeJson(String value) + { + try (CharArrayWriter writer = new CharArrayWriter(); + JsonGenerator generator = jsonFactory().createGenerator(writer)) { + generator.writeString(value); + generator.flush(); + return writer.toString(); + } + catch (IOException e) { + throw new RuntimeException(e); + } + } +} diff --git a/lib/trino-parquet/src/main/java/io/trino/parquet/spark/VariantUtil.java b/lib/trino-parquet/src/main/java/io/trino/parquet/spark/VariantUtil.java new file mode 100644 index 000000000000..c17c33f647f0 --- /dev/null +++ b/lib/trino-parquet/src/main/java/io/trino/parquet/spark/VariantUtil.java @@ -0,0 +1,480 @@ +/* + * 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.parquet.spark; + +import java.math.BigDecimal; +import java.math.BigInteger; +import java.util.Arrays; + +import static java.nio.charset.StandardCharsets.UTF_8; + +/** + * Copied from https://github.com/apache/spark/blob/53d65fd12dd9231139188227ef9040d40d759021/common/variant/src/main/java/org/apache/spark/types/variant/VariantUtil.java + * + * This class defines constants related to the variant format and provides functions for + * manipulating variant binaries. + + * A variant is made up of 2 binaries: value and metadata. A variant value consists of a one-byte + * header and a number of content bytes (can be zero). The header byte is divided into upper 6 bits + * (called "type info") and lower 2 bits (called "basic type"). The content format is explained in + * the below constants for all possible basic type and type info values. + + * The variant metadata includes a version id and a dictionary of distinct strings (case-sensitive). + * Its binary format is: + * - Version: 1-byte unsigned integer. The only acceptable value is 1 currently. + * - Dictionary size: 4-byte little-endian unsigned integer. The number of keys in the + * dictionary. + * - Offsets: (size + 1) * 4-byte little-endian unsigned integers. `offsets[i]` represents the + * starting position of string i, counting starting from the address of `offsets[0]`. Strings + * must be stored contiguously, so we don’t need to store the string size, instead, we compute it + * with `offset[i + 1] - offset[i]`. + * - UTF-8 string data. + */ +public final class VariantUtil +{ + public static final int BASIC_TYPE_BITS = 2; + public static final int BASIC_TYPE_MASK = 0x3; + public static final int TYPE_INFO_MASK = 0x3F; + + // Below is all possible basic type values. + // Primitive value. The type info value must be one of the values in the below section. + public static final int PRIMITIVE = 0; + // Short string value. The type info value is the string size, which must be in `[0, + // kMaxShortStrSize]`. + // The string content bytes directly follow the header byte. + public static final int SHORT_STR = 1; + // Object value. The content contains a size, a list of field ids, a list of field offsets, and + // the actual field data. The length of the id list is `size`, while the length of the offset + // list is `size + 1`, where the last offset represent the total size of the field data. The + // fields in an object must be sorted by the field name in alphabetical order. Duplicate field + // names in one object are not allowed. + // We use 5 bits in the type info to specify the integer type of the object header: it should + // be 0_b4_b3b2_b1b0 (MSB is 0), where: + // - b4 specifies the type of size. When it is 0/1, `size` is a little-endian 1/4-byte + // unsigned integer. + // - b3b2/b1b0 specifies the integer type of id and offset. When the 2 bits are 0/1/2, the + // list contains 1/2/3-byte little-endian unsigned integers. + public static final int OBJECT = 2; + // Array value. The content contains a size, a list of field offsets, and the actual element + // data. It is similar to an object without the id list. The length of the offset list + // is `size + 1`, where the last offset represent the total size of the element data. + // Its type info should be: 000_b2_b1b0: + // - b2 specifies the type of size. + // - b1b0 specifies the integer type of offset. + public static final int ARRAY = 3; + + // Below is all possible type info values for `PRIMITIVE`. + // JSON Null value. Empty content. + public static final int NULL = 0; + // True value. Empty content. + public static final int TRUE = 1; + // False value. Empty content. + public static final int FALSE = 2; + // 1-byte little-endian signed integer. + public static final int INT1 = 3; + // 2-byte little-endian signed integer. + public static final int INT2 = 4; + // 4-byte little-endian signed integer. + public static final int INT4 = 5; + // 4-byte little-endian signed integer. + public static final int INT8 = 6; + // 8-byte IEEE double. + public static final int DOUBLE = 7; + // 4-byte decimal. Content is 1-byte scale + 4-byte little-endian signed integer. + public static final int DECIMAL4 = 8; + // 8-byte decimal. Content is 1-byte scale + 8-byte little-endian signed integer. + public static final int DECIMAL8 = 9; + // 16-byte decimal. Content is 1-byte scale + 16-byte little-endian signed integer. + public static final int DECIMAL16 = 10; + // Date value. Content is 4-byte little-endian signed integer that represents the number of days + // from the Unix epoch. + public static final int DATE = 11; + // Timestamp value. Content is 8-byte little-endian signed integer that represents the number of + // microseconds elapsed since the Unix epoch, 1970-01-01 00:00:00 UTC. It is displayed to users in + // their local time zones and may be displayed differently depending on the execution environment. + public static final int TIMESTAMP = 12; + // Timestamp_ntz value. It has the same content as `TIMESTAMP` but should always be interpreted + // as if the local time zone is UTC. + public static final int TIMESTAMP_NTZ = 13; + // 4-byte IEEE float. + public static final int FLOAT = 14; + // Binary value. The content is (4-byte little-endian unsigned integer representing the binary + // size) + (size bytes of binary content). + public static final int BINARY = 15; + // Long string value. The content is (4-byte little-endian unsigned integer representing the + // string size) + (size bytes of string content). + public static final int LONG_STR = 16; + + public static final byte VERSION = 1; + // The lower 4 bits of the first metadata byte contain the version. + public static final byte VERSION_MASK = 0x0F; + + public static final int U24_MAX = 0xFFFFFF; + public static final int U32_SIZE = 4; + + // Both variant value and variant metadata need to be no longer than 16MiB. + public static final int SIZE_LIMIT = U24_MAX + 1; + + public static final int MAX_DECIMAL4_PRECISION = 9; + public static final int MAX_DECIMAL8_PRECISION = 18; + public static final int MAX_DECIMAL16_PRECISION = 38; + + private VariantUtil() {} + + // Check the validity of an array index `position`. Throw `MALFORMED_VARIANT` if it is out of bound, + // meaning that the variant is malformed. + static void checkIndex(int position, int length) + { + if (position < 0 || position >= length) { + throw new IllegalArgumentException("Index out of bound: %s (length: %s)".formatted(position, length)); + } + } + + // Read a little-endian signed long value from `bytes[position, position + numBytes)`. + static long readLong(byte[] bytes, int position, int numBytes) + { + checkIndex(position, bytes.length); + checkIndex(position + numBytes - 1, bytes.length); + long result = 0; + // All bytes except the most significant byte should be unsign-extended and shifted (so we need + // `& 0xFF`). The most significant byte should be sign-extended and is handled after the loop. + for (int i = 0; i < numBytes - 1; ++i) { + long unsignedByteValue = bytes[position + i] & 0xFF; + result |= unsignedByteValue << (8 * i); + } + long signedByteValue = bytes[position + numBytes - 1]; + result |= signedByteValue << (8 * (numBytes - 1)); + return result; + } + + // Read a little-endian unsigned int value from `bytes[position, position + numBytes)`. The value must fit + // into a non-negative int (`[0, Integer.MAX_VALUE]`). + static int readUnsigned(byte[] bytes, int position, int numBytes) + { + checkIndex(position, bytes.length); + checkIndex(position + numBytes - 1, bytes.length); + int result = 0; + // Similar to the `readLong` loop, but all bytes should be unsign-extended. + for (int i = 0; i < numBytes; ++i) { + int unsignedByteValue = bytes[position + i] & 0xFF; + result |= unsignedByteValue << (8 * i); + } + if (result < 0) { + throw new IllegalArgumentException("Value out of bound: %s".formatted(result)); + } + return result; + } + + // The value type of variant value. It is determined by the header byte but not a 1:1 mapping + // (for example, INT1/2/4/8 all maps to `Type.LONG`). + public enum Type + { + NULL, + BOOLEAN, + LONG, + FLOAT, + DOUBLE, + DECIMAL, + STRING, + BINARY, + DATE, + TIMESTAMP, + TIMESTAMP_NTZ, + ARRAY, + OBJECT, + } + + // Get the value type of variant value `value[position...]`. It is only legal to call `get*` if + // `getType` returns this type (for example, it is only legal to call `getLong` if `getType` + // returns `Type.Long`). + // Throw `MALFORMED_VARIANT` if the variant is malformed. + public static Type getType(byte[] value, int position) + { + checkIndex(position, value.length); + int basicType = value[position] & BASIC_TYPE_MASK; + int typeInfo = (value[position] >> BASIC_TYPE_BITS) & TYPE_INFO_MASK; + return switch (basicType) { + case SHORT_STR -> Type.STRING; + case OBJECT -> Type.OBJECT; + case ARRAY -> Type.ARRAY; + default -> switch (typeInfo) { + case NULL -> Type.NULL; + case TRUE, FALSE -> Type.BOOLEAN; + case INT1, INT2, INT4, INT8 -> Type.LONG; + case DOUBLE -> Type.DOUBLE; + case DECIMAL4, DECIMAL8, DECIMAL16 -> Type.DECIMAL; + case DATE -> Type.DATE; + case TIMESTAMP -> Type.TIMESTAMP; + case TIMESTAMP_NTZ -> Type.TIMESTAMP_NTZ; + case FLOAT -> Type.FLOAT; + case BINARY -> Type.BINARY; + case LONG_STR -> Type.STRING; + default -> throw new IllegalArgumentException("Unexpected type: " + typeInfo); + }; + }; + } + + private static IllegalStateException unexpectedType(Type type) + { + return new IllegalStateException("Expect type to be " + type); + } + + // Get a boolean value from variant value `value[position...]`. + // Throw `MALFORMED_VARIANT` if the variant is malformed. + public static boolean getBoolean(byte[] value, int position) + { + checkIndex(position, value.length); + int basicType = value[position] & BASIC_TYPE_MASK; + int typeInfo = (value[position] >> BASIC_TYPE_BITS) & TYPE_INFO_MASK; + if (basicType != PRIMITIVE || (typeInfo != TRUE && typeInfo != FALSE)) { + throw unexpectedType(Type.BOOLEAN); + } + return typeInfo == TRUE; + } + + // Get a long value from variant value `value[position...]`. + // It is only legal to call it if `getType` returns one of `Type.LONG/DATE/TIMESTAMP/ + // TIMESTAMP_NTZ`. If the type is `DATE`, the return value is guaranteed to fit into an int and + // represents the number of days from the Unix epoch. If the type is `TIMESTAMP/TIMESTAMP_NTZ`, + // the return value represents the number of microseconds from the Unix epoch. + // Throw `MALFORMED_VARIANT` if the variant is malformed. + public static long getLong(byte[] value, int position) + { + checkIndex(position, value.length); + int basicType = value[position] & BASIC_TYPE_MASK; + int typeInfo = (value[position] >> BASIC_TYPE_BITS) & TYPE_INFO_MASK; + String exceptionMessage = "Expect type to be LONG/DATE/TIMESTAMP/TIMESTAMP_NTZ"; + if (basicType != PRIMITIVE) { + throw new IllegalStateException(exceptionMessage); + } + return switch (typeInfo) { + case INT1 -> readLong(value, position + 1, 1); + case INT2 -> readLong(value, position + 1, 2); + case INT4, DATE -> readLong(value, position + 1, 4); + case INT8, TIMESTAMP, TIMESTAMP_NTZ -> readLong(value, position + 1, 8); + default -> throw new IllegalStateException(exceptionMessage); + }; + } + + // Get a double value from variant value `value[position...]`. + // Throw `MALFORMED_VARIANT` if the variant is malformed. + public static double getDouble(byte[] value, int position) + { + checkIndex(position, value.length); + int basicType = value[position] & BASIC_TYPE_MASK; + int typeInfo = (value[position] >> BASIC_TYPE_BITS) & TYPE_INFO_MASK; + if (basicType != PRIMITIVE || typeInfo != DOUBLE) { + throw unexpectedType(Type.DOUBLE); + } + return Double.longBitsToDouble(readLong(value, position + 1, 8)); + } + + // Check whether the precision and scale of the decimal are within the limit. + private static void checkDecimal(BigDecimal decimal, int maxPrecision) + { + if (decimal.precision() > maxPrecision || decimal.scale() > maxPrecision) { + throw new IllegalArgumentException("Decimal out of bound: " + decimal); + } + } + + // Get a decimal value from variant value `value[position...]`. + // Throw `MALFORMED_VARIANT` if the variant is malformed. + public static BigDecimal getDecimal(byte[] value, int position) + { + checkIndex(position, value.length); + int basicType = value[position] & BASIC_TYPE_MASK; + int typeInfo = (value[position] >> BASIC_TYPE_BITS) & TYPE_INFO_MASK; + if (basicType != PRIMITIVE) { + throw unexpectedType(Type.DECIMAL); + } + // Interpret the scale byte as unsigned. If it is a negative byte, the unsigned value must be + // greater than `MAX_DECIMAL16_PRECISION` and will trigger an error in `checkDecimal`. + int scale = value[position + 1] & 0xFF; + BigDecimal result; + switch (typeInfo) { + case DECIMAL4: + result = BigDecimal.valueOf(readLong(value, position + 2, 4), scale); + checkDecimal(result, MAX_DECIMAL4_PRECISION); + break; + case DECIMAL8: + result = BigDecimal.valueOf(readLong(value, position + 2, 8), scale); + checkDecimal(result, MAX_DECIMAL8_PRECISION); + break; + case DECIMAL16: + checkIndex(position + 17, value.length); + byte[] bytes = new byte[16]; + // Copy the bytes reversely because the `BigInteger` constructor expects a big-endian + // representation. + for (int i = 0; i < 16; ++i) { + bytes[i] = value[position + 17 - i]; + } + result = new BigDecimal(new BigInteger(bytes), scale); + checkDecimal(result, MAX_DECIMAL16_PRECISION); + break; + default: + throw unexpectedType(Type.DECIMAL); + } + return result.stripTrailingZeros(); + } + + // Get a float value from variant value `value[position...]`. + // Throw `MALFORMED_VARIANT` if the variant is malformed. + public static float getFloat(byte[] value, int position) + { + checkIndex(position, value.length); + int basicType = value[position] & BASIC_TYPE_MASK; + int typeInfo = (value[position] >> BASIC_TYPE_BITS) & TYPE_INFO_MASK; + if (basicType != PRIMITIVE || typeInfo != FLOAT) { + throw unexpectedType(Type.FLOAT); + } + return Float.intBitsToFloat((int) readLong(value, position + 1, 4)); + } + + // Get a binary value from variant value `value[position...]`. + // Throw `MALFORMED_VARIANT` if the variant is malformed. + public static byte[] getBinary(byte[] value, int position) + { + checkIndex(position, value.length); + int basicType = value[position] & BASIC_TYPE_MASK; + int typeInfo = (value[position] >> BASIC_TYPE_BITS) & TYPE_INFO_MASK; + if (basicType != PRIMITIVE || typeInfo != BINARY) { + throw unexpectedType(Type.BINARY); + } + int start = position + 1 + U32_SIZE; + int length = readUnsigned(value, position + 1, U32_SIZE); + checkIndex(start + length - 1, value.length); + return Arrays.copyOfRange(value, start, start + length); + } + + // Get a string value from variant value `value[position...]`. + // Throw `MALFORMED_VARIANT` if the variant is malformed. + public static String getString(byte[] value, int position) + { + checkIndex(position, value.length); + int basicType = value[position] & BASIC_TYPE_MASK; + int typeInfo = (value[position] >> BASIC_TYPE_BITS) & TYPE_INFO_MASK; + if (basicType == SHORT_STR || (basicType == PRIMITIVE && typeInfo == LONG_STR)) { + int start; + int length; + if (basicType == SHORT_STR) { + start = position + 1; + length = typeInfo; + } + else { + start = position + 1 + U32_SIZE; + length = readUnsigned(value, position + 1, U32_SIZE); + } + checkIndex(start + length - 1, value.length); + return new String(value, start, length, UTF_8); + } + throw unexpectedType(Type.STRING); + } + + public interface ObjectHandler + { + /** + * @param size Number of object fields. + * @param idSize The integer size of the field id list. + * @param offsetSize The integer size of the offset list. + * @param idStart The starting index of the field id list in the variant value array. + * @param offsetStart The starting index of the offset list in the variant value array. + * @param dataStart The starting index of field data in the variant value array. + */ + T apply(int size, int idSize, int offsetSize, int idStart, int offsetStart, int dataStart); + } + + // A helper function to access a variant object. It provides `handler` with its required + // parameters and returns what it returns. + public static T handleObject(byte[] value, int position, ObjectHandler handler) + { + checkIndex(position, value.length); + int basicType = value[position] & BASIC_TYPE_MASK; + int typeInfo = (value[position] >> BASIC_TYPE_BITS) & TYPE_INFO_MASK; + if (basicType != OBJECT) { + throw unexpectedType(Type.OBJECT); + } + // Refer to the comment of the `OBJECT` constant for the details of the object header encoding. + // Suppose `typeInfo` has a bit representation of 0_b4_b3b2_b1b0, the following line extracts + // b4 to determine whether the object uses a 1/4-byte size. + boolean largeSize = ((typeInfo >> 4) & 0x1) != 0; + int sizeBytes = (largeSize ? U32_SIZE : 1); + int size = readUnsigned(value, position + 1, sizeBytes); + // Extracts b3b2 to determine the integer size of the field id list. + int idSize = ((typeInfo >> 2) & 0x3) + 1; + // Extracts b1b0 to determine the integer size of the offset list. + int offsetSize = (typeInfo & 0x3) + 1; + int idStart = position + 1 + sizeBytes; + int offsetStart = idStart + size * idSize; + int dataStart = offsetStart + (size + 1) * offsetSize; + return handler.apply(size, idSize, offsetSize, idStart, offsetStart, dataStart); + } + + public interface ArrayHandler + { + /** + * @param size Number of array elements. + * @param offsetSize The integer size of the offset list. + * @param offsetStart The starting index of the offset list in the variant value array. + * @param dataStart The starting index of element data in the variant value array. + */ + T apply(int size, int offsetSize, int offsetStart, int dataStart); + } + + // A helper function to access a variant array. + public static T handleArray(byte[] value, int position, ArrayHandler handler) + { + checkIndex(position, value.length); + int basicType = value[position] & BASIC_TYPE_MASK; + int typeInfo = (value[position] >> BASIC_TYPE_BITS) & TYPE_INFO_MASK; + if (basicType != ARRAY) { + throw unexpectedType(Type.ARRAY); + } + // Refer to the comment of the `ARRAY` constant for the details of the object header encoding. + // Suppose `typeInfo` has a bit representation of 000_b2_b1b0, the following line extracts + // b2 to determine whether the object uses a 1/4-byte size. + boolean largeSize = ((typeInfo >> 2) & 0x1) != 0; + int sizeBytes = (largeSize ? U32_SIZE : 1); + int size = readUnsigned(value, position + 1, sizeBytes); + // Extracts b1b0 to determine the integer size of the offset list. + int offsetSize = (typeInfo & 0x3) + 1; + int offsetStart = position + 1 + sizeBytes; + int dataStart = offsetStart + (size + 1) * offsetSize; + return handler.apply(size, offsetSize, offsetStart, dataStart); + } + + // Get a key at `id` in the variant metadata. + // Throw `MALFORMED_VARIANT` if the variant is malformed. An out-of-bound `id` is also considered + // a malformed variant because it is read from the corresponding variant value. + public static String getMetadataKey(byte[] metadata, int id) + { + checkIndex(0, metadata.length); + // Extracts the highest 2 bits in the metadata header to determine the integer size of the + // offset list. + int offsetSize = ((metadata[0] >> 6) & 0x3) + 1; + int dictSize = readUnsigned(metadata, 1, offsetSize); + if (id >= dictSize) { + throw new IllegalArgumentException("Index out of bound: %s (size: %s)".formatted(id, dictSize)); + } + // There are a header byte, a `dictSize` with `offsetSize` bytes, and `(dictSize + 1)` offsets + // before the string data. + int stringStart = 1 + (dictSize + 2) * offsetSize; + int offset = readUnsigned(metadata, 1 + (id + 1) * offsetSize, offsetSize); + int nextOffset = readUnsigned(metadata, 1 + (id + 2) * offsetSize, offsetSize); + if (offset > nextOffset) { + throw new IllegalArgumentException("Invalid offset: %s > %s".formatted(offset, nextOffset)); + } + checkIndex(stringStart + nextOffset - 1, metadata.length); + return new String(metadata, stringStart + offset, nextOffset - offset, UTF_8); + } +} diff --git a/lib/trino-parquet/src/main/java/io/trino/parquet/writer/ColumnWriter.java b/lib/trino-parquet/src/main/java/io/trino/parquet/writer/ColumnWriter.java index b82d4a4942a1..1e25b0da2631 100644 --- a/lib/trino-parquet/src/main/java/io/trino/parquet/writer/ColumnWriter.java +++ b/lib/trino-parquet/src/main/java/io/trino/parquet/writer/ColumnWriter.java @@ -13,10 +13,13 @@ */ package io.trino.parquet.writer; +import org.apache.parquet.column.values.bloomfilter.BloomFilter; import org.apache.parquet.format.ColumnMetaData; import java.io.IOException; import java.util.List; +import java.util.Optional; +import java.util.OptionalInt; import static java.util.Objects.requireNonNull; @@ -38,10 +41,14 @@ class BufferData { private final ColumnMetaData metaData; private final List data; + private final OptionalInt dictionaryPageSize; + private final Optional bloomFilter; - public BufferData(List data, ColumnMetaData metaData) + public BufferData(List data, OptionalInt dictionaryPageSize, Optional bloomFilter, ColumnMetaData metaData) { this.data = requireNonNull(data, "data is null"); + this.dictionaryPageSize = requireNonNull(dictionaryPageSize, "dictionaryPageSize is null"); + this.bloomFilter = requireNonNull(bloomFilter, "bloomFilter is null"); this.metaData = requireNonNull(metaData, "metaData is null"); } @@ -54,5 +61,15 @@ public List getData() { return data; } + + public OptionalInt getDictionaryPageSize() + { + return dictionaryPageSize; + } + + public Optional getBloomFilter() + { + return bloomFilter; + } } } diff --git a/lib/trino-parquet/src/main/java/io/trino/parquet/writer/MessageTypeConverter.java b/lib/trino-parquet/src/main/java/io/trino/parquet/writer/MessageTypeConverter.java index 3b64ad6dfc7c..6fc67f4329c9 100644 --- a/lib/trino-parquet/src/main/java/io/trino/parquet/writer/MessageTypeConverter.java +++ b/lib/trino-parquet/src/main/java/io/trino/parquet/writer/MessageTypeConverter.java @@ -88,7 +88,7 @@ public void visit(GroupType groupType) element.setConverted_type(getConvertedType(groupType.getOriginalType())); } if (groupType.getLogicalTypeAnnotation() != null) { - element.setLogicalType(convertToLogicalType(new ParquetMetadataConverter(), groupType.getLogicalTypeAnnotation())); + element.setLogicalType(io.trino.parquet.ParquetMetadataConverter.convertToLogicalType(groupType.getLogicalTypeAnnotation())); } if (groupType.getId() != null) { element.setField_id(groupType.getId().intValue()); diff --git a/lib/trino-parquet/src/main/java/io/trino/parquet/writer/ParquetWriter.java b/lib/trino-parquet/src/main/java/io/trino/parquet/writer/ParquetWriter.java index 38d5170e65c6..fc702b13a2d6 100644 --- a/lib/trino-parquet/src/main/java/io/trino/parquet/writer/ParquetWriter.java +++ b/lib/trino-parquet/src/main/java/io/trino/parquet/writer/ParquetWriter.java @@ -19,25 +19,36 @@ import io.airlift.slice.OutputStreamSliceOutput; import io.airlift.slice.Slice; import io.airlift.slice.Slices; -import io.trino.parquet.Field; +import io.trino.parquet.Column; import io.trino.parquet.ParquetCorruptionException; import io.trino.parquet.ParquetDataSource; import io.trino.parquet.ParquetReaderOptions; import io.trino.parquet.ParquetWriteValidation; +import io.trino.parquet.metadata.BlockMetadata; +import io.trino.parquet.metadata.FileMetadata; +import io.trino.parquet.metadata.ParquetMetadata; import io.trino.parquet.reader.MetadataReader; import io.trino.parquet.reader.ParquetReader; +import io.trino.parquet.reader.RowGroupInfo; import io.trino.parquet.writer.ColumnWriter.BufferData; import io.trino.spi.Page; import io.trino.spi.type.Type; -import org.apache.parquet.column.ParquetProperties; +import jakarta.annotation.Nullable; +import org.apache.parquet.column.ColumnDescriptor; +import org.apache.parquet.column.values.bloomfilter.BloomFilter; +import org.apache.parquet.format.BloomFilterAlgorithm; +import org.apache.parquet.format.BloomFilterCompression; +import org.apache.parquet.format.BloomFilterHash; +import org.apache.parquet.format.BloomFilterHeader; import org.apache.parquet.format.ColumnMetaData; import org.apache.parquet.format.CompressionCodec; import org.apache.parquet.format.FileMetaData; import org.apache.parquet.format.KeyValue; import org.apache.parquet.format.RowGroup; +import org.apache.parquet.format.SplitBlockAlgorithm; +import org.apache.parquet.format.Uncompressed; import org.apache.parquet.format.Util; -import org.apache.parquet.hadoop.metadata.BlockMetaData; -import org.apache.parquet.hadoop.metadata.ParquetMetadata; +import org.apache.parquet.format.XxHash; import org.apache.parquet.io.MessageColumnIO; import org.apache.parquet.schema.MessageType; import org.joda.time.DateTimeZone; @@ -45,9 +56,11 @@ import java.io.Closeable; import java.io.IOException; import java.io.OutputStream; +import java.io.UncheckedIOException; import java.util.List; import java.util.Map; import java.util.Optional; +import java.util.OptionalInt; import java.util.function.Consumer; import static com.google.common.base.Preconditions.checkArgument; @@ -61,32 +74,38 @@ import static io.trino.memory.context.AggregatedMemoryContext.newSimpleAggregatedMemoryContext; import static io.trino.parquet.ParquetTypeUtils.constructField; import static io.trino.parquet.ParquetTypeUtils.getColumnIO; +import static io.trino.parquet.ParquetTypeUtils.getDescriptors; import static io.trino.parquet.ParquetTypeUtils.lookupColumnByName; import static io.trino.parquet.ParquetWriteValidation.ParquetWriteValidationBuilder; +import static io.trino.parquet.metadata.PrunedBlockMetadata.createPrunedColumnsMetadata; import static io.trino.parquet.writer.ParquetDataOutput.createDataOutput; +import static io.trino.spi.type.BigintType.BIGINT; +import static io.trino.spi.type.DoubleType.DOUBLE; +import static io.trino.spi.type.IntegerType.INTEGER; +import static io.trino.spi.type.RealType.REAL; +import static io.trino.spi.type.UuidType.UUID; +import static io.trino.spi.type.VarbinaryType.VARBINARY; +import static io.trino.spi.type.VarcharType.VARCHAR; import static java.lang.Math.max; import static java.lang.Math.min; import static java.nio.charset.StandardCharsets.US_ASCII; -import static java.util.Collections.nCopies; import static java.util.Objects.requireNonNull; -import static org.apache.parquet.column.ParquetProperties.WriterVersion.PARQUET_1_0; public class ParquetWriter implements Closeable { private static final int INSTANCE_SIZE = instanceSize(ParquetWriter.class); + public static final List SUPPORTED_BLOOM_FILTER_TYPES = ImmutableList.of(BIGINT, DOUBLE, INTEGER, REAL, UUID, VARBINARY, VARCHAR); private final OutputStreamSliceOutput outputStream; private final ParquetWriterOptions writerOption; private final MessageType messageType; - private final String createdBy; - private final int chunkMaxLogicalBytes; + private final int chunkMaxBytes; private final Map, Type> primitiveTypes; private final CompressionCodec compressionCodec; - private final boolean useBatchColumnReadersForVerification; private final Optional parquetTimeZone; - - private final ImmutableList.Builder rowGroupBuilder = ImmutableList.builder(); + private final FileFooter fileFooter; + private final ImmutableList.Builder>> bloomFilterGroups = ImmutableList.builder(); private final Optional validationBuilder; private List columnWriters; @@ -94,6 +113,8 @@ public class ParquetWriter private long bufferedBytes; private boolean closed; private boolean writeHeader; + @Nullable + private FileMetaData fileMetaData; public static final Slice MAGIC = wrappedBuffer("PAR1".getBytes(US_ASCII)); @@ -104,7 +125,6 @@ public ParquetWriter( ParquetWriterOptions writerOption, CompressionCodec compressionCodec, String trinoVersion, - boolean useBatchColumnReadersForVerification, Optional parquetTimeZone, Optional validationBuilder) { @@ -114,15 +134,15 @@ public ParquetWriter( this.primitiveTypes = requireNonNull(primitiveTypes, "primitiveTypes is null"); this.writerOption = requireNonNull(writerOption, "writerOption is null"); this.compressionCodec = requireNonNull(compressionCodec, "compressionCodec is null"); - this.useBatchColumnReadersForVerification = useBatchColumnReadersForVerification; this.parquetTimeZone = requireNonNull(parquetTimeZone, "parquetTimeZone is null"); - this.createdBy = formatCreatedBy(requireNonNull(trinoVersion, "trinoVersion is null")); + String createdBy = formatCreatedBy(requireNonNull(trinoVersion, "trinoVersion is null")); + this.fileFooter = new FileFooter(messageType, createdBy, parquetTimeZone); recordValidation(validation -> validation.setTimeZone(parquetTimeZone.map(DateTimeZone::getID))); recordValidation(validation -> validation.setColumns(messageType.getColumns())); recordValidation(validation -> validation.setCreatedBy(createdBy)); initColumnWriters(); - this.chunkMaxLogicalBytes = max(1, writerOption.getMaxRowGroupSize() / 2); + this.chunkMaxBytes = max(1, writerOption.getMaxRowGroupSize() / 2); } public long getWrittenBytes() @@ -162,7 +182,7 @@ public void write(Page page) Page chunk = page.getRegion(writeOffset, min(page.getPositionCount() - writeOffset, writerOption.getBatchSize())); // avoid chunk with huge logical size - while (chunk.getPositionCount() > 1 && chunk.getLogicalSizeInBytes() > chunkMaxLogicalBytes) { + while (chunk.getPositionCount() > 1 && chunk.getLogicalSizeInBytes() > chunkMaxBytes) { chunk = page.getRegion(writeOffset, chunk.getPositionCount() / 2); } @@ -204,6 +224,8 @@ public void close() columnWriters.forEach(ColumnWriter::close); flush(); columnWriters = ImmutableList.of(); + fileMetaData = fileFooter.createFileMetadata(); + writeBloomFilters(fileMetaData.getRow_groups(), bloomFilterGroups.build()); writeFooter(); } bufferedBytes = 0; @@ -231,40 +253,47 @@ public void validate(ParquetDataSource input) } } + public FileMetaData getFileMetaData() + { + checkState(closed, "fileMetaData is available only after writer is closed"); + return requireNonNull(fileMetaData, "fileMetaData is null"); + } + private ParquetReader createParquetReader(ParquetDataSource input, ParquetMetadata parquetMetadata, ParquetWriteValidation writeValidation) throws IOException { - org.apache.parquet.hadoop.metadata.FileMetaData fileMetaData = parquetMetadata.getFileMetaData(); + FileMetadata fileMetaData = parquetMetadata.getFileMetaData(); MessageColumnIO messageColumnIO = getColumnIO(fileMetaData.getSchema(), fileMetaData.getSchema()); - ImmutableList.Builder columnFields = ImmutableList.builder(); + ImmutableList.Builder columnFields = ImmutableList.builder(); for (int i = 0; i < writeValidation.getTypes().size(); i++) { - columnFields.add(constructField( - writeValidation.getTypes().get(i), - lookupColumnByName(messageColumnIO, writeValidation.getColumnNames().get(i))) - .orElseThrow()); + columnFields.add(new Column( + messageColumnIO.getName(), + constructField( + writeValidation.getTypes().get(i), + lookupColumnByName(messageColumnIO, writeValidation.getColumnNames().get(i))) + .orElseThrow())); } + Map, ColumnDescriptor> descriptorsByPath = getDescriptors(fileMetaData.getSchema(), fileMetaData.getSchema()); long nextStart = 0; - ImmutableList.Builder blockStartsBuilder = ImmutableList.builder(); - for (BlockMetaData block : parquetMetadata.getBlocks()) { - blockStartsBuilder.add(nextStart); - nextStart += block.getRowCount(); + ImmutableList.Builder rowGroupInfoBuilder = ImmutableList.builder(); + for (BlockMetadata block : parquetMetadata.getBlocks()) { + rowGroupInfoBuilder.add(new RowGroupInfo(createPrunedColumnsMetadata(block, input.getId(), descriptorsByPath), nextStart, Optional.empty())); + nextStart += block.rowCount(); } - List blockStarts = blockStartsBuilder.build(); return new ParquetReader( Optional.ofNullable(fileMetaData.getCreatedBy()), columnFields.build(), - parquetMetadata.getBlocks(), - blockStarts, + false, + rowGroupInfoBuilder.build(), input, parquetTimeZone.orElseThrow(), newSimpleAggregatedMemoryContext(), - new ParquetReaderOptions().withBatchColumnReaders(useBatchColumnReadersForVerification), + new ParquetReaderOptions(), exception -> { throwIfUnchecked(exception); return new RuntimeException(exception); }, Optional.empty(), - nCopies(blockStarts.size(), Optional.empty()), Optional.of(writeValidation)); } @@ -307,32 +336,35 @@ private void flush() } // update stats - long stripeStartOffset = outputStream.longSize(); - List columns = bufferDataList.stream() - .map(BufferData::getMetaData) - .collect(toImmutableList()); - - long currentOffset = stripeStartOffset; - for (ColumnMetaData columnMetaData : columns) { - columnMetaData.setData_page_offset(currentOffset); + long currentOffset = outputStream.longSize(); + ImmutableList.Builder columnMetaDataBuilder = ImmutableList.builder(); + for (BufferData bufferData : bufferDataList) { + ColumnMetaData columnMetaData = bufferData.getMetaData(); + OptionalInt dictionaryPageSize = bufferData.getDictionaryPageSize(); + if (dictionaryPageSize.isPresent()) { + columnMetaData.setDictionary_page_offset(currentOffset); + } + columnMetaData.setData_page_offset(currentOffset + dictionaryPageSize.orElse(0)); + columnMetaDataBuilder.add(columnMetaData); currentOffset += columnMetaData.getTotal_compressed_size(); } - updateRowGroups(columns); + updateRowGroups(columnMetaDataBuilder.build(), outputStream.longSize()); // flush pages for (BufferData bufferData : bufferDataList) { bufferData.getData() .forEach(data -> data.writeData(outputStream)); } + + bloomFilterGroups.add(bufferDataList.stream().map(BufferData::getBloomFilter).collect(toImmutableList())); } private void writeFooter() throws IOException { checkState(closed); - List rowGroups = rowGroupBuilder.build(); - Slice footer = getFooter(rowGroups, messageType); - recordValidation(validation -> validation.setRowGroups(rowGroups)); + Slice footer = serializeFooter(fileMetaData); + recordValidation(validation -> validation.setRowGroups(fileMetaData.getRow_groups())); createDataOutput(footer).writeData(outputStream); Slice footerSize = Slices.allocate(SIZE_OF_INT); @@ -342,31 +374,55 @@ private void writeFooter() createDataOutput(MAGIC).writeData(outputStream); } - Slice getFooter(List rowGroups, MessageType messageType) - throws IOException + private void writeBloomFilters(List rowGroups, List>> rowGroupBloomFilters) { - FileMetaData fileMetaData = new FileMetaData(); - fileMetaData.setVersion(1); - fileMetaData.setCreated_by(createdBy); - fileMetaData.setSchema(MessageTypeConverter.toParquetSchema(messageType)); - // Added based on org.apache.hadoop.hive.ql.io.parquet.write.DataWritableWriteSupport - parquetTimeZone.ifPresent(dateTimeZone -> fileMetaData.setKey_value_metadata( - ImmutableList.of(new KeyValue("writer.time.zone").setValue(dateTimeZone.getID())))); - long totalRows = rowGroups.stream().mapToLong(RowGroup::getNum_rows).sum(); - fileMetaData.setNum_rows(totalRows); - fileMetaData.setRow_groups(ImmutableList.copyOf(rowGroups)); + checkArgument(rowGroups.size() == rowGroupBloomFilters.size(), "Row groups size %s should match row group Bloom filter size %s", rowGroups.size(), rowGroupBloomFilters.size()); + for (int group = 0; group < rowGroups.size(); group++) { + List columns = rowGroups.get(group).getColumns(); + List> bloomFilters = rowGroupBloomFilters.get(group); + for (int i = 0; i < columns.size(); i++) { + if (bloomFilters.get(i).isEmpty()) { + continue; + } - DynamicSliceOutput dynamicSliceOutput = new DynamicSliceOutput(40); - Util.writeFileMetaData(fileMetaData, dynamicSliceOutput); - return dynamicSliceOutput.slice(); + BloomFilter bloomFilter = bloomFilters.get(i).orElseThrow(); + long bloomFilterOffset = outputStream.longSize(); + try { + Util.writeBloomFilterHeader( + new BloomFilterHeader( + bloomFilter.getBitsetSize(), + BloomFilterAlgorithm.BLOCK(new SplitBlockAlgorithm()), + BloomFilterHash.XXHASH(new XxHash()), + BloomFilterCompression.UNCOMPRESSED(new Uncompressed())), + outputStream, + null, + null); + bloomFilter.writeTo(outputStream); + columns.get(i).getMeta_data().setBloom_filter_offset(bloomFilterOffset); + } + catch (IOException e) { + throw new UncheckedIOException(e); + } + } + } } - private void updateRowGroups(List columnMetaData) + private void updateRowGroups(List columnMetaData, long fileOffset) { long totalCompressedBytes = columnMetaData.stream().mapToLong(ColumnMetaData::getTotal_compressed_size).sum(); long totalBytes = columnMetaData.stream().mapToLong(ColumnMetaData::getTotal_uncompressed_size).sum(); ImmutableList columnChunks = columnMetaData.stream().map(ParquetWriter::toColumnChunk).collect(toImmutableList()); - rowGroupBuilder.add(new RowGroup(columnChunks, totalBytes, rows).setTotal_compressed_size(totalCompressedBytes)); + fileFooter.addRowGroup(new RowGroup(columnChunks, totalBytes, rows) + .setTotal_compressed_size(totalCompressedBytes) + .setFile_offset(fileOffset)); + } + + private static Slice serializeFooter(FileMetaData fileMetaData) + throws IOException + { + DynamicSliceOutput dynamicSliceOutput = new DynamicSliceOutput(40); + Util.writeFileMetaData(fileMetaData, dynamicSliceOutput); + return dynamicSliceOutput.slice(); } private static org.apache.parquet.format.ColumnChunk toColumnChunk(ColumnMetaData metaData) @@ -388,11 +444,52 @@ static String formatCreatedBy(String trinoVersion) private void initColumnWriters() { - ParquetProperties parquetProperties = ParquetProperties.builder() - .withWriterVersion(PARQUET_1_0) - .withPageSize(writerOption.getMaxPageSize()) - .build(); + this.columnWriters = ParquetWriters.getColumnWriters( + messageType, + primitiveTypes, + compressionCodec, + writerOption, + parquetTimeZone); + } + + private static class FileFooter + { + private final MessageType messageType; + private final String createdBy; + private final Optional parquetTimeZone; + + @Nullable + private ImmutableList.Builder rowGroupBuilder = ImmutableList.builder(); + + private FileFooter(MessageType messageType, String createdBy, Optional parquetTimeZone) + { + this.messageType = messageType; + this.createdBy = createdBy; + this.parquetTimeZone = parquetTimeZone; + } - this.columnWriters = ParquetWriters.getColumnWriters(messageType, primitiveTypes, parquetProperties, compressionCodec, parquetTimeZone); + public void addRowGroup(RowGroup rowGroup) + { + checkState(rowGroupBuilder != null, "rowGroupBuilder is null"); + rowGroupBuilder.add(rowGroup); + } + + public FileMetaData createFileMetadata() + { + checkState(rowGroupBuilder != null, "rowGroupBuilder is null"); + List rowGroups = rowGroupBuilder.build(); + rowGroupBuilder = null; + long totalRows = rowGroups.stream().mapToLong(RowGroup::getNum_rows).sum(); + FileMetaData fileMetaData = new FileMetaData( + 1, + MessageTypeConverter.toParquetSchema(messageType), + totalRows, + ImmutableList.copyOf(rowGroups)); + fileMetaData.setCreated_by(createdBy); + // Added based on org.apache.hadoop.hive.ql.io.parquet.write.DataWritableWriteSupport + parquetTimeZone.ifPresent(dateTimeZone -> fileMetaData.setKey_value_metadata( + ImmutableList.of(new KeyValue("writer.time.zone").setValue(dateTimeZone.getID())))); + return fileMetaData; + } } } diff --git a/lib/trino-parquet/src/main/java/io/trino/parquet/writer/ParquetWriterOptions.java b/lib/trino-parquet/src/main/java/io/trino/parquet/writer/ParquetWriterOptions.java index 7b7af7591ac3..c41cabf5a006 100644 --- a/lib/trino-parquet/src/main/java/io/trino/parquet/writer/ParquetWriterOptions.java +++ b/lib/trino-parquet/src/main/java/io/trino/parquet/writer/ParquetWriterOptions.java @@ -13,15 +13,26 @@ */ package io.trino.parquet.writer; +import com.google.common.collect.ImmutableSet; import com.google.common.primitives.Ints; import io.airlift.units.DataSize; -import org.apache.parquet.hadoop.ParquetWriter; +import org.apache.parquet.column.ParquetProperties; + +import java.util.Set; + +import static com.google.common.base.Preconditions.checkArgument; +import static io.airlift.units.DataSize.Unit.MEGABYTE; public class ParquetWriterOptions { - private static final DataSize DEFAULT_MAX_ROW_GROUP_SIZE = DataSize.ofBytes(ParquetWriter.DEFAULT_BLOCK_SIZE); - private static final DataSize DEFAULT_MAX_PAGE_SIZE = DataSize.ofBytes(ParquetWriter.DEFAULT_PAGE_SIZE); + private static final DataSize DEFAULT_MAX_ROW_GROUP_SIZE = DataSize.of(128, MEGABYTE); + private static final DataSize DEFAULT_MAX_PAGE_SIZE = DataSize.ofBytes(ParquetProperties.DEFAULT_PAGE_SIZE); + // org.apache.parquet.column.DEFAULT_PAGE_ROW_COUNT_LIMIT is 20_000 to improve selectivity of page indexes + // This value should be revisited when TODO https://github.com/trinodb/trino/issues/9359 is implemented + public static final int DEFAULT_MAX_PAGE_VALUE_COUNT = 60_000; public static final int DEFAULT_BATCH_SIZE = 10_000; + public static final DataSize DEFAULT_MAX_BLOOM_FILTER_SIZE = DataSize.of(1, MEGABYTE); + public static final double DEFAULT_BLOOM_FILTER_FPP = 0.05; public static ParquetWriterOptions.Builder builder() { @@ -30,13 +41,30 @@ public static ParquetWriterOptions.Builder builder() private final int maxRowGroupSize; private final int maxPageSize; + private final int maxPageValueCount; private final int batchSize; + private final int maxBloomFilterSize; + private final double bloomFilterFpp; + // Set of column dot paths to columns with bloom filters + private final Set bloomFilterColumns; - private ParquetWriterOptions(DataSize maxBlockSize, DataSize maxPageSize, int batchSize) + private ParquetWriterOptions( + DataSize maxBlockSize, + DataSize maxPageSize, + int maxPageValueCount, + int batchSize, + DataSize maxBloomFilterSize, + double bloomFilterFpp, + Set bloomFilterColumns) { this.maxRowGroupSize = Ints.saturatedCast(maxBlockSize.toBytes()); this.maxPageSize = Ints.saturatedCast(maxPageSize.toBytes()); + this.maxPageValueCount = maxPageValueCount; this.batchSize = batchSize; + this.maxBloomFilterSize = Ints.saturatedCast(maxBloomFilterSize.toBytes()); + this.bloomFilterFpp = bloomFilterFpp; + this.bloomFilterColumns = ImmutableSet.copyOf(bloomFilterColumns); + checkArgument(this.bloomFilterFpp > 0.0 && this.bloomFilterFpp < 1.0, "bloomFilterFpp should be > 0.0 & < 1.0"); } public int getMaxRowGroupSize() @@ -49,16 +77,40 @@ public int getMaxPageSize() return maxPageSize; } + public int getMaxPageValueCount() + { + return maxPageValueCount; + } + public int getBatchSize() { return batchSize; } + public int getMaxBloomFilterSize() + { + return maxBloomFilterSize; + } + + public Set getBloomFilterColumns() + { + return bloomFilterColumns; + } + + public double getBLoomFilterFpp() + { + return bloomFilterFpp; + } + public static class Builder { private DataSize maxBlockSize = DEFAULT_MAX_ROW_GROUP_SIZE; private DataSize maxPageSize = DEFAULT_MAX_PAGE_SIZE; + private int maxPageValueCount = DEFAULT_MAX_PAGE_VALUE_COUNT; private int batchSize = DEFAULT_BATCH_SIZE; + private DataSize maxBloomFilterSize = DEFAULT_MAX_BLOOM_FILTER_SIZE; + private Set bloomFilterColumns = ImmutableSet.of(); + private double bloomFilterFpp = DEFAULT_BLOOM_FILTER_FPP; public Builder setMaxBlockSize(DataSize maxBlockSize) { @@ -72,15 +124,34 @@ public Builder setMaxPageSize(DataSize maxPageSize) return this; } + public Builder setMaxPageValueCount(int maxPageValueCount) + { + this.maxPageValueCount = maxPageValueCount; + return this; + } + public Builder setBatchSize(int batchSize) { this.batchSize = batchSize; return this; } + public Builder setBloomFilterColumns(Set columns) + { + this.bloomFilterColumns = columns; + return this; + } + public ParquetWriterOptions build() { - return new ParquetWriterOptions(maxBlockSize, maxPageSize, batchSize); + return new ParquetWriterOptions( + maxBlockSize, + maxPageSize, + maxPageValueCount, + batchSize, + maxBloomFilterSize, + bloomFilterFpp, + bloomFilterColumns); } } } diff --git a/lib/trino-parquet/src/main/java/io/trino/parquet/writer/ParquetWriters.java b/lib/trino-parquet/src/main/java/io/trino/parquet/writer/ParquetWriters.java index 0697feec1ee3..2afb33822f48 100644 --- a/lib/trino-parquet/src/main/java/io/trino/parquet/writer/ParquetWriters.java +++ b/lib/trino-parquet/src/main/java/io/trino/parquet/writer/ParquetWriters.java @@ -13,6 +13,7 @@ */ package io.trino.parquet.writer; +import com.google.common.base.Joiner; import com.google.common.collect.ImmutableList; import io.trino.parquet.writer.valuewriter.BigintValueWriter; import io.trino.parquet.writer.valuewriter.BinaryValueWriter; @@ -32,6 +33,7 @@ import io.trino.parquet.writer.valuewriter.TimestampNanosValueWriter; import io.trino.parquet.writer.valuewriter.TimestampTzMicrosValueWriter; import io.trino.parquet.writer.valuewriter.TimestampTzMillisValueWriter; +import io.trino.parquet.writer.valuewriter.TrinoValuesWriterFactory; import io.trino.parquet.writer.valuewriter.UuidValueWriter; import io.trino.spi.TrinoException; import io.trino.spi.type.CharType; @@ -42,8 +44,9 @@ import io.trino.spi.type.VarbinaryType; import io.trino.spi.type.VarcharType; import org.apache.parquet.column.ColumnDescriptor; -import org.apache.parquet.column.ParquetProperties; import org.apache.parquet.column.values.ValuesWriter; +import org.apache.parquet.column.values.bloomfilter.BlockSplitBloomFilter; +import org.apache.parquet.column.values.bloomfilter.BloomFilter; import org.apache.parquet.format.CompressionCodec; import org.apache.parquet.schema.GroupType; import org.apache.parquet.schema.LogicalTypeAnnotation; @@ -57,9 +60,13 @@ import java.util.List; import java.util.Map; import java.util.Optional; +import java.util.Set; import java.util.function.Predicate; import static com.google.common.base.Preconditions.checkArgument; +import static io.trino.parquet.writer.ParquetWriter.SUPPORTED_BLOOM_FILTER_TYPES; +import static io.trino.parquet.writer.valuewriter.ColumnDescriptorValuesWriter.newDefinitionLevelWriter; +import static io.trino.parquet.writer.valuewriter.ColumnDescriptorValuesWriter.newRepetitionLevelWriter; import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED; import static io.trino.spi.type.BigintType.BIGINT; import static io.trino.spi.type.BooleanType.BOOLEAN; @@ -82,16 +89,26 @@ final class ParquetWriters { + private static final int DEFAULT_DICTIONARY_PAGE_SIZE = 1024 * 1024; + static final int BLOOM_FILTER_EXPECTED_ENTRIES = 100_000; + private ParquetWriters() {} static List getColumnWriters( MessageType messageType, Map, Type> trinoTypes, - ParquetProperties parquetProperties, CompressionCodec compressionCodec, + ParquetWriterOptions writerOptions, Optional parquetTimeZone) { - WriteBuilder writeBuilder = new WriteBuilder(messageType, trinoTypes, parquetProperties, compressionCodec, parquetTimeZone); + TrinoValuesWriterFactory valuesWriterFactory = new TrinoValuesWriterFactory(writerOptions.getMaxPageSize(), DEFAULT_DICTIONARY_PAGE_SIZE); + WriteBuilder writeBuilder = new WriteBuilder( + messageType, + trinoTypes, + valuesWriterFactory, + compressionCodec, + writerOptions, + parquetTimeZone); ParquetTypeVisitor.visit(messageType, writeBuilder); return writeBuilder.build(); } @@ -101,22 +118,33 @@ private static class WriteBuilder { private final MessageType type; private final Map, Type> trinoTypes; - private final ParquetProperties parquetProperties; + private final TrinoValuesWriterFactory valuesWriterFactory; private final CompressionCodec compressionCodec; + private final int maxPageSize; + private final int pageValueCountLimit; + private final Set bloomFilterColumns; private final Optional parquetTimeZone; private final ImmutableList.Builder builder = ImmutableList.builder(); + private final int maxBloomFilterSize; + private final double bloomFilterFpp; WriteBuilder( MessageType messageType, Map, Type> trinoTypes, - ParquetProperties parquetProperties, + TrinoValuesWriterFactory valuesWriterFactory, CompressionCodec compressionCodec, + ParquetWriterOptions writerOptions, Optional parquetTimeZone) { this.type = requireNonNull(messageType, "messageType is null"); this.trinoTypes = requireNonNull(trinoTypes, "trinoTypes is null"); - this.parquetProperties = requireNonNull(parquetProperties, "parquetProperties is null"); + this.valuesWriterFactory = requireNonNull(valuesWriterFactory, "valuesWriterFactory is null"); this.compressionCodec = requireNonNull(compressionCodec, "compressionCodec is null"); + this.maxPageSize = writerOptions.getMaxPageSize(); + this.pageValueCountLimit = writerOptions.getMaxPageValueCount(); + this.maxBloomFilterSize = writerOptions.getMaxBloomFilterSize(); + this.bloomFilterColumns = requireNonNull(writerOptions.getBloomFilterColumns(), "bloomFilterColumns is null"); + this.bloomFilterFpp = writerOptions.getBLoomFilterFpp(); this.parquetTimeZone = requireNonNull(parquetTimeZone, "parquetTimeZone is null"); } @@ -166,13 +194,16 @@ public ColumnWriter primitive(PrimitiveType primitive) int fieldRepetitionLevel = type.getMaxRepetitionLevel(path); ColumnDescriptor columnDescriptor = new ColumnDescriptor(path, primitive, fieldRepetitionLevel, fieldDefinitionLevel); Type trinoType = requireNonNull(trinoTypes.get(ImmutableList.copyOf(path)), "Trino type is null"); + Optional bloomFilter = createBloomFilter(bloomFilterColumns, maxBloomFilterSize, bloomFilterFpp, columnDescriptor, trinoType); return new PrimitiveColumnWriter( columnDescriptor, - getValueWriter(parquetProperties.newValuesWriter(columnDescriptor), trinoType, columnDescriptor.getPrimitiveType(), parquetTimeZone), - parquetProperties.newDefinitionLevelWriter(columnDescriptor), - parquetProperties.newRepetitionLevelWriter(columnDescriptor), + getValueWriter(valuesWriterFactory.newValuesWriter(columnDescriptor, bloomFilter), trinoType, columnDescriptor.getPrimitiveType(), parquetTimeZone), + newDefinitionLevelWriter(columnDescriptor, maxPageSize), + newRepetitionLevelWriter(columnDescriptor, maxPageSize), compressionCodec, - parquetProperties.getPageSizeThreshold()); + maxPageSize, + pageValueCountLimit, + bloomFilter); } private String[] currentPath() @@ -186,6 +217,20 @@ private String[] currentPath() } return path; } + + private static Optional createBloomFilter(Set bloomFilterColumns, int maxBloomFilterSize, double bloomFilterFpp, ColumnDescriptor columnDescriptor, Type colummType) + { + if (!SUPPORTED_BLOOM_FILTER_TYPES.contains(colummType)) { + return Optional.empty(); + } + // TODO: Enable use of AdaptiveBlockSplitBloomFilter once parquet-mr 1.14.0 is released + String dotPath = Joiner.on('.').join(columnDescriptor.getPath()); + if (bloomFilterColumns.contains(dotPath)) { + int optimalNumOfBits = BlockSplitBloomFilter.optimalNumOfBits(BLOOM_FILTER_EXPECTED_ENTRIES, bloomFilterFpp); + return Optional.of(new BlockSplitBloomFilter(optimalNumOfBits / 8, maxBloomFilterSize)); + } + return Optional.empty(); + } } private static PrimitiveValueWriter getValueWriter(ValuesWriter valuesWriter, Type type, PrimitiveType parquetType, Optional parquetTimeZone) diff --git a/lib/trino-parquet/src/main/java/io/trino/parquet/writer/PrimitiveColumnWriter.java b/lib/trino-parquet/src/main/java/io/trino/parquet/writer/PrimitiveColumnWriter.java index 24a482b34205..6715be46ecec 100644 --- a/lib/trino-parquet/src/main/java/io/trino/parquet/writer/PrimitiveColumnWriter.java +++ b/lib/trino-parquet/src/main/java/io/trino/parquet/writer/PrimitiveColumnWriter.java @@ -15,24 +15,29 @@ import com.google.common.collect.ImmutableList; import io.airlift.slice.Slices; +import io.trino.parquet.ParquetMetadataConverter; import io.trino.parquet.writer.repdef.DefLevelWriterProvider; import io.trino.parquet.writer.repdef.DefLevelWriterProviders; import io.trino.parquet.writer.repdef.RepLevelWriterProvider; import io.trino.parquet.writer.repdef.RepLevelWriterProviders; +import io.trino.parquet.writer.valuewriter.ColumnDescriptorValuesWriter; import io.trino.parquet.writer.valuewriter.PrimitiveValueWriter; import io.trino.plugin.base.io.ChunkedSliceOutput; import jakarta.annotation.Nullable; import org.apache.parquet.bytes.BytesInput; import org.apache.parquet.column.ColumnDescriptor; import org.apache.parquet.column.Encoding; +import org.apache.parquet.column.EncodingStats; import org.apache.parquet.column.page.DictionaryPage; import org.apache.parquet.column.statistics.Statistics; -import org.apache.parquet.column.values.ValuesWriter; +import org.apache.parquet.column.values.bloomfilter.BloomFilter; import org.apache.parquet.format.ColumnMetaData; import org.apache.parquet.format.CompressionCodec; +import org.apache.parquet.format.DataPageHeader; +import org.apache.parquet.format.DictionaryPageHeader; import org.apache.parquet.format.PageEncodingStats; +import org.apache.parquet.format.PageHeader; import org.apache.parquet.format.PageType; -import org.apache.parquet.format.converter.ParquetMetadataConverter; import java.io.ByteArrayOutputStream; import java.io.IOException; @@ -40,11 +45,15 @@ import java.util.HashSet; import java.util.List; import java.util.Map; +import java.util.Optional; +import java.util.OptionalInt; import java.util.Set; import static com.google.common.base.Preconditions.checkState; import static com.google.common.collect.ImmutableList.toImmutableList; import static io.airlift.slice.SizeOf.instanceSize; +import static io.trino.parquet.ParquetMetadataConverter.convertEncodingStats; +import static io.trino.parquet.ParquetMetadataConverter.getEncoding; import static io.trino.parquet.writer.ParquetCompressor.getCompressor; import static io.trino.parquet.writer.ParquetDataOutput.createDataOutput; import static io.trino.parquet.writer.repdef.DefLevelWriterProvider.DefinitionLevelWriter; @@ -52,6 +61,7 @@ import static io.trino.parquet.writer.repdef.RepLevelWriterProvider.RepetitionLevelWriter; import static io.trino.parquet.writer.repdef.RepLevelWriterProvider.getRootRepetitionLevelWriter; import static java.util.Objects.requireNonNull; +import static org.apache.parquet.format.Util.writePageHeader; public class PrimitiveColumnWriter implements ColumnWriter @@ -59,15 +69,16 @@ public class PrimitiveColumnWriter private static final int INSTANCE_SIZE = instanceSize(PrimitiveColumnWriter.class); private static final int MINIMUM_OUTPUT_BUFFER_CHUNK_SIZE = 8 * 1024; private static final int MAXIMUM_OUTPUT_BUFFER_CHUNK_SIZE = 2 * 1024 * 1024; + // ParquetMetadataConverter.MAX_STATS_SIZE is 4096, we need a value which would guarantee that min and max + // don't add up to 4096 (so less than 2048). Using 1K as that is big enough for most use cases. + private static final int MAX_STATISTICS_LENGTH_IN_BYTES = 1024; private final ColumnDescriptor columnDescriptor; private final CompressionCodec compressionCodec; private final PrimitiveValueWriter primitiveValueWriter; - private final ValuesWriter definitionLevelWriter; - private final ValuesWriter repetitionLevelWriter; - - private final ParquetMetadataConverter parquetMetadataConverter = new ParquetMetadataConverter(); + private final ColumnDescriptorValuesWriter definitionLevelWriter; + private final ColumnDescriptorValuesWriter repetitionLevelWriter; private boolean closed; private boolean getDataStreamsCalled; @@ -80,10 +91,11 @@ public class PrimitiveColumnWriter private final Set encodings = new HashSet<>(); private final Map dataPagesWithEncoding = new HashMap<>(); private final Map dictionaryPagesWithEncoding = new HashMap<>(); + private final Statistics columnStatistics; + private final Optional bloomFilter; private long totalCompressedSize; private long totalUnCompressedSize; private long totalValues; - private Statistics columnStatistics; private final int maxDefinitionLevel; @@ -93,13 +105,22 @@ public class PrimitiveColumnWriter private final ParquetCompressor compressor; private final int pageSizeThreshold; + private final int pageValueCountLimit; // Total size of compressed parquet pages and the current uncompressed page buffered in memory // Used by ParquetWriter to decide when a row group is big enough to flush private long bufferedBytes; private long pageBufferedBytes; - public PrimitiveColumnWriter(ColumnDescriptor columnDescriptor, PrimitiveValueWriter primitiveValueWriter, ValuesWriter definitionLevelWriter, ValuesWriter repetitionLevelWriter, CompressionCodec compressionCodec, int pageSizeThreshold) + public PrimitiveColumnWriter( + ColumnDescriptor columnDescriptor, + PrimitiveValueWriter primitiveValueWriter, + ColumnDescriptorValuesWriter definitionLevelWriter, + ColumnDescriptorValuesWriter repetitionLevelWriter, + CompressionCodec compressionCodec, + int pageSizeThreshold, + int pageValueCountLimit, + Optional bloomFilter) { this.columnDescriptor = requireNonNull(columnDescriptor, "columnDescriptor is null"); this.maxDefinitionLevel = columnDescriptor.getMaxDefinitionLevel(); @@ -109,8 +130,10 @@ public PrimitiveColumnWriter(ColumnDescriptor columnDescriptor, PrimitiveValueWr this.compressionCodec = requireNonNull(compressionCodec, "compressionCodec is null"); this.compressor = getCompressor(compressionCodec); this.pageSizeThreshold = pageSizeThreshold; + this.pageValueCountLimit = pageValueCountLimit; this.columnStatistics = Statistics.createStats(columnDescriptor.getPrimitiveType()); this.compressedOutputStream = new ChunkedSliceOutput(MINIMUM_OUTPUT_BUFFER_CHUNK_SIZE, MAXIMUM_OUTPUT_BUFFER_CHUNK_SIZE); + this.bloomFilter = requireNonNull(bloomFilter, "bloomFilter is null"); } @Override @@ -142,7 +165,7 @@ public void writeBlock(ColumnChunk columnChunk) } long currentPageBufferedBytes = getCurrentPageBufferedBytes(); - if (currentPageBufferedBytes >= pageSizeThreshold) { + if (valueCount >= pageValueCountLimit || currentPageBufferedBytes >= pageSizeThreshold) { flushCurrentPageToBuffer(); } else { @@ -161,7 +184,17 @@ public List getBuffer() throws IOException { checkState(closed); - return ImmutableList.of(new BufferData(getDataStreams(), getColumnMetaData())); + DataStreams dataStreams = getDataStreams(); + ColumnMetaData columnMetaData = getColumnMetaData(); + + EncodingStats stats = convertEncodingStats(columnMetaData.getEncoding_stats()); + boolean isOnlyDictionaryEncodingPages = stats.hasDictionaryPages() && !stats.hasNonDictionaryEncodedPages(); + + return ImmutableList.of(new BufferData( + dataStreams.data(), + dataStreams.dictionaryPageSize(), + isOnlyDictionaryEncodingPages ? Optional.empty() : dataStreams.bloomFilter(), + columnMetaData)); } // Returns ColumnMetaData that offset is invalid @@ -171,14 +204,14 @@ private ColumnMetaData getColumnMetaData() ColumnMetaData columnMetaData = new ColumnMetaData( ParquetTypeConverter.getType(columnDescriptor.getPrimitiveType().getPrimitiveTypeName()), - encodings.stream().map(parquetMetadataConverter::getEncoding).collect(toImmutableList()), + encodings.stream().map(ParquetMetadataConverter::getEncoding).collect(toImmutableList()), ImmutableList.copyOf(columnDescriptor.getPath()), compressionCodec, totalValues, totalUnCompressedSize, totalCompressedSize, -1); - columnMetaData.setStatistics(ParquetMetadataConverter.toParquetStatistics(columnStatistics)); + columnMetaData.setStatistics(ParquetMetadataConverter.toParquetStatistics(columnStatistics, MAX_STATISTICS_LENGTH_IN_BYTES)); ImmutableList.Builder pageEncodingStats = ImmutableList.builder(); dataPagesWithEncoding.entrySet().stream() .map(encodingAndCount -> new PageEncodingStats(PageType.DATA_PAGE, encodingAndCount.getKey(), encodingAndCount.getValue())) @@ -213,16 +246,17 @@ private void flushCurrentPageToBuffer() columnStatistics.mergeStatistics(statistics); int writtenBytesSoFar = compressedOutputStream.size(); - parquetMetadataConverter.writeDataPageV1Header(uncompressedSize, + PageHeader header = dataPageV1Header( + uncompressedSize, compressedSize, valueCount, repetitionLevelWriter.getEncoding(), definitionLevelWriter.getEncoding(), - primitiveValueWriter.getEncoding(), - compressedOutputStream); + primitiveValueWriter.getEncoding()); + writePageHeader(header, compressedOutputStream); int pageHeaderSize = compressedOutputStream.size() - writtenBytesSoFar; - dataPagesWithEncoding.merge(parquetMetadataConverter.getEncoding(primitiveValueWriter.getEncoding()), 1, Integer::sum); + dataPagesWithEncoding.merge(getEncoding(primitiveValueWriter.getEncoding()), 1, Integer::sum); // update total stats totalUnCompressedSize += pageHeaderSize + uncompressedSize; @@ -248,7 +282,7 @@ private void flushCurrentPageToBuffer() updateBufferedBytes(getCurrentPageBufferedBytes()); } - private List getDataStreams() + private DataStreams getDataStreams() throws IOException { ImmutableList.Builder outputs = ImmutableList.builder(); @@ -257,6 +291,7 @@ private List getDataStreams() } // write dict page if possible DictionaryPage dictionaryPage = primitiveValueWriter.toDictPageAndClose(); + OptionalInt dictionaryPageSize = OptionalInt.empty(); if (dictionaryPage != null) { int uncompressedSize = dictionaryPage.getUncompressedSize(); byte[] pageBytes = dictionaryPage.getBytes().toByteArray(); @@ -266,25 +301,26 @@ private List getDataStreams() int compressedSize = pageData.size(); ByteArrayOutputStream dictStream = new ByteArrayOutputStream(); - parquetMetadataConverter.writeDictionaryPageHeader( + PageHeader header = dictionaryPageHeader( uncompressedSize, compressedSize, dictionaryPage.getDictionarySize(), - dictionaryPage.getEncoding(), - dictStream); + dictionaryPage.getEncoding()); + writePageHeader(header, dictStream); ParquetDataOutput pageHeader = createDataOutput(dictStream); outputs.add(pageHeader); outputs.add(pageData); totalCompressedSize += pageHeader.size() + compressedSize; totalUnCompressedSize += pageHeader.size() + uncompressedSize; - dictionaryPagesWithEncoding.merge(new ParquetMetadataConverter().getEncoding(dictionaryPage.getEncoding()), 1, Integer::sum); + dictionaryPagesWithEncoding.merge(getEncoding(dictionaryPage.getEncoding()), 1, Integer::sum); + dictionaryPageSize = OptionalInt.of(pageHeader.size() + compressedSize); primitiveValueWriter.resetDictionary(); } getDataStreamsCalled = true; outputs.add(createDataOutput(compressedOutputStream)); - return outputs.build(); + return new DataStreams(outputs.build(), dictionaryPageSize, bloomFilter); } @Override @@ -314,4 +350,34 @@ private long getCurrentPageBufferedBytes() repetitionLevelWriter.getBufferedSize() + primitiveValueWriter.getBufferedSize(); } + + private static PageHeader dataPageV1Header( + int uncompressedSize, + int compressedSize, + int valueCount, + org.apache.parquet.column.Encoding rlEncoding, + org.apache.parquet.column.Encoding dlEncoding, + org.apache.parquet.column.Encoding valuesEncoding) + { + PageHeader header = new PageHeader(PageType.DATA_PAGE, uncompressedSize, compressedSize); + header.setData_page_header(new DataPageHeader( + valueCount, + getEncoding(valuesEncoding), + getEncoding(dlEncoding), + getEncoding(rlEncoding))); + return header; + } + + private static PageHeader dictionaryPageHeader( + int uncompressedSize, + int compressedSize, + int valueCount, + org.apache.parquet.column.Encoding valuesEncoding) + { + PageHeader header = new PageHeader(PageType.DICTIONARY_PAGE, uncompressedSize, compressedSize); + header.setDictionary_page_header(new DictionaryPageHeader(valueCount, getEncoding(valuesEncoding))); + return header; + } + + private record DataStreams(List data, OptionalInt dictionaryPageSize, Optional bloomFilter) {} } diff --git a/lib/trino-parquet/src/main/java/io/trino/parquet/writer/StructColumnWriter.java b/lib/trino-parquet/src/main/java/io/trino/parquet/writer/StructColumnWriter.java index a03351ce13be..181c1942ea35 100644 --- a/lib/trino-parquet/src/main/java/io/trino/parquet/writer/StructColumnWriter.java +++ b/lib/trino-parquet/src/main/java/io/trino/parquet/writer/StructColumnWriter.java @@ -19,13 +19,12 @@ import io.trino.parquet.writer.repdef.RepLevelWriterProvider; import io.trino.parquet.writer.repdef.RepLevelWriterProviders; import io.trino.spi.block.Block; -import io.trino.spi.block.ColumnarRow; +import io.trino.spi.block.RowBlock; import java.io.IOException; import java.util.List; import static io.airlift.slice.SizeOf.instanceSize; -import static io.trino.spi.block.ColumnarRow.toColumnarRow; import static java.util.Objects.requireNonNull; import static org.apache.parquet.Preconditions.checkArgument; @@ -47,22 +46,23 @@ public StructColumnWriter(List columnWriters, int maxDefinitionLev public void writeBlock(ColumnChunk columnChunk) throws IOException { - ColumnarRow columnarRow = toColumnarRow(columnChunk.getBlock()); - checkArgument(columnarRow.getFieldCount() == columnWriters.size(), "ColumnarRow field size %s is not equal to columnWriters size %s", columnarRow.getFieldCount(), columnWriters.size()); + Block block = columnChunk.getBlock(); + List fields = RowBlock.getNullSuppressedRowFieldsFromBlock(block); + checkArgument(fields.size() == columnWriters.size(), "Row field size %s is not equal to columnWriters size %s", fields.size(), columnWriters.size()); List defLevelWriterProviders = ImmutableList.builder() .addAll(columnChunk.getDefLevelWriterProviders()) - .add(DefLevelWriterProviders.of(columnarRow, maxDefinitionLevel)) + .add(DefLevelWriterProviders.of(block, maxDefinitionLevel)) .build(); List repLevelWriterProviders = ImmutableList.builder() .addAll(columnChunk.getRepLevelWriterProviders()) - .add(RepLevelWriterProviders.of(columnarRow)) + .add(RepLevelWriterProviders.of(block)) .build(); for (int i = 0; i < columnWriters.size(); ++i) { ColumnWriter columnWriter = columnWriters.get(i); - Block block = columnarRow.getField(i); - columnWriter.writeBlock(new ColumnChunk(block, defLevelWriterProviders, repLevelWriterProviders)); + Block field = fields.get(i); + columnWriter.writeBlock(new ColumnChunk(field, defLevelWriterProviders, repLevelWriterProviders)); } } diff --git a/lib/trino-parquet/src/main/java/io/trino/parquet/writer/repdef/DefLevelWriterProvider.java b/lib/trino-parquet/src/main/java/io/trino/parquet/writer/repdef/DefLevelWriterProvider.java index 8b3bff8e8c97..3eb3fe979f01 100644 --- a/lib/trino-parquet/src/main/java/io/trino/parquet/writer/repdef/DefLevelWriterProvider.java +++ b/lib/trino-parquet/src/main/java/io/trino/parquet/writer/repdef/DefLevelWriterProvider.java @@ -14,14 +14,14 @@ package io.trino.parquet.writer.repdef; import com.google.common.collect.Iterables; -import org.apache.parquet.column.values.ValuesWriter; +import io.trino.parquet.writer.valuewriter.ColumnDescriptorValuesWriter; import java.util.List; import java.util.Optional; public interface DefLevelWriterProvider { - DefinitionLevelWriter getDefinitionLevelWriter(Optional nestedWriter, ValuesWriter encoder); + DefinitionLevelWriter getDefinitionLevelWriter(Optional nestedWriter, ColumnDescriptorValuesWriter encoder); interface DefinitionLevelWriter { @@ -34,7 +34,7 @@ record ValuesCount(int totalValuesCount, int maxDefinitionLevelValuesCount) { } - static DefinitionLevelWriter getRootDefinitionLevelWriter(List defLevelWriterProviders, ValuesWriter encoder) + static DefinitionLevelWriter getRootDefinitionLevelWriter(List defLevelWriterProviders, ColumnDescriptorValuesWriter encoder) { // Constructs hierarchy of DefinitionLevelWriter from leaf to root DefinitionLevelWriter rootDefinitionLevelWriter = Iterables.getLast(defLevelWriterProviders) diff --git a/lib/trino-parquet/src/main/java/io/trino/parquet/writer/repdef/DefLevelWriterProviders.java b/lib/trino-parquet/src/main/java/io/trino/parquet/writer/repdef/DefLevelWriterProviders.java index 7de391070c78..aad756eef270 100644 --- a/lib/trino-parquet/src/main/java/io/trino/parquet/writer/repdef/DefLevelWriterProviders.java +++ b/lib/trino-parquet/src/main/java/io/trino/parquet/writer/repdef/DefLevelWriterProviders.java @@ -13,15 +13,18 @@ */ package io.trino.parquet.writer.repdef; +import io.trino.parquet.writer.valuewriter.ColumnDescriptorValuesWriter; +import io.trino.spi.block.ArrayBlock; import io.trino.spi.block.Block; import io.trino.spi.block.ColumnarArray; import io.trino.spi.block.ColumnarMap; -import io.trino.spi.block.ColumnarRow; -import org.apache.parquet.column.values.ValuesWriter; +import io.trino.spi.block.MapBlock; +import io.trino.spi.block.RowBlock; import java.util.Optional; import static com.google.common.base.Preconditions.checkArgument; +import static io.trino.spi.PageBlockUtil.getUnderlyingValueBlock; import static java.lang.String.format; import static java.util.Objects.requireNonNull; @@ -31,14 +34,12 @@ private DefLevelWriterProviders() {} public static DefLevelWriterProvider of(Block block, int maxDefinitionLevel) { + if (getUnderlyingValueBlock(block) instanceof RowBlock) { + return new RowDefLevelWriterProvider(block, maxDefinitionLevel); + } return new PrimitiveDefLevelWriterProvider(block, maxDefinitionLevel); } - public static DefLevelWriterProvider of(ColumnarRow columnarRow, int maxDefinitionLevel) - { - return new ColumnRowDefLevelWriterProvider(columnarRow, maxDefinitionLevel); - } - public static DefLevelWriterProvider of(ColumnarArray columnarArray, int maxDefinitionLevel) { return new ColumnArrayDefLevelWriterProvider(columnarArray, maxDefinitionLevel); @@ -59,10 +60,13 @@ static class PrimitiveDefLevelWriterProvider { this.block = requireNonNull(block, "block is null"); this.maxDefinitionLevel = maxDefinitionLevel; + checkArgument(!(getUnderlyingValueBlock(block) instanceof RowBlock), "block is a row block"); + checkArgument(!(getUnderlyingValueBlock(block) instanceof ArrayBlock), "block is an array block"); + checkArgument(!(getUnderlyingValueBlock(block) instanceof MapBlock), "block is a map block"); } @Override - public DefinitionLevelWriter getDefinitionLevelWriter(Optional nestedWriter, ValuesWriter encoder) + public DefinitionLevelWriter getDefinitionLevelWriter(Optional nestedWriter, ColumnDescriptorValuesWriter encoder) { checkArgument(nestedWriter.isEmpty(), "nestedWriter should be empty for primitive definition level writer"); return new DefinitionLevelWriter() @@ -81,9 +85,7 @@ public ValuesCount writeDefinitionLevels(int positionsCount) checkValidPosition(offset, positionsCount, block.getPositionCount()); int nonNullsCount = 0; if (!block.mayHaveNull()) { - for (int position = offset; position < offset + positionsCount; position++) { - encoder.writeInteger(maxDefinitionLevel); - } + encoder.writeRepeatInteger(maxDefinitionLevel, positionsCount); nonNullsCount = positionsCount; } else { @@ -100,20 +102,21 @@ public ValuesCount writeDefinitionLevels(int positionsCount) } } - static class ColumnRowDefLevelWriterProvider + static class RowDefLevelWriterProvider implements DefLevelWriterProvider { - private final ColumnarRow columnarRow; + private final Block block; private final int maxDefinitionLevel; - ColumnRowDefLevelWriterProvider(ColumnarRow columnarRow, int maxDefinitionLevel) + RowDefLevelWriterProvider(Block block, int maxDefinitionLevel) { - this.columnarRow = requireNonNull(columnarRow, "columnarRow is null"); + this.block = requireNonNull(block, "block is null"); this.maxDefinitionLevel = maxDefinitionLevel; + checkArgument(getUnderlyingValueBlock(block) instanceof RowBlock, "block is not a row block"); } @Override - public DefinitionLevelWriter getDefinitionLevelWriter(Optional nestedWriterOptional, ValuesWriter encoder) + public DefinitionLevelWriter getDefinitionLevelWriter(Optional nestedWriterOptional, ColumnDescriptorValuesWriter encoder) { checkArgument(nestedWriterOptional.isPresent(), "nestedWriter should be present for column row definition level writer"); return new DefinitionLevelWriter() @@ -125,21 +128,21 @@ public DefinitionLevelWriter getDefinitionLevelWriter(Optional nestedWriterOptional, ValuesWriter encoder) + public DefinitionLevelWriter getDefinitionLevelWriter(Optional nestedWriterOptional, ColumnDescriptorValuesWriter encoder) { checkArgument(nestedWriterOptional.isPresent(), "nestedWriter should be present for column map definition level writer"); return new DefinitionLevelWriter() @@ -261,7 +264,7 @@ static class ColumnArrayDefLevelWriterProvider } @Override - public DefinitionLevelWriter getDefinitionLevelWriter(Optional nestedWriterOptional, ValuesWriter encoder) + public DefinitionLevelWriter getDefinitionLevelWriter(Optional nestedWriterOptional, ColumnDescriptorValuesWriter encoder) { checkArgument(nestedWriterOptional.isPresent(), "nestedWriter should be present for column map definition level writer"); return new DefinitionLevelWriter() diff --git a/lib/trino-parquet/src/main/java/io/trino/parquet/writer/repdef/RepLevelWriterProvider.java b/lib/trino-parquet/src/main/java/io/trino/parquet/writer/repdef/RepLevelWriterProvider.java index c835c9060c54..941219ef1878 100644 --- a/lib/trino-parquet/src/main/java/io/trino/parquet/writer/repdef/RepLevelWriterProvider.java +++ b/lib/trino-parquet/src/main/java/io/trino/parquet/writer/repdef/RepLevelWriterProvider.java @@ -14,14 +14,14 @@ package io.trino.parquet.writer.repdef; import com.google.common.collect.Iterables; -import org.apache.parquet.column.values.ValuesWriter; +import io.trino.parquet.writer.valuewriter.ColumnDescriptorValuesWriter; import java.util.List; import java.util.Optional; public interface RepLevelWriterProvider { - RepetitionLevelWriter getRepetitionLevelWriter(Optional nestedWriter, ValuesWriter encoder); + RepetitionLevelWriter getRepetitionLevelWriter(Optional nestedWriter, ColumnDescriptorValuesWriter encoder); /** * Parent repetition level marks at which level either: @@ -36,7 +36,7 @@ interface RepetitionLevelWriter void writeRepetitionLevels(int parentLevel); } - static RepetitionLevelWriter getRootRepetitionLevelWriter(List repLevelWriterProviders, ValuesWriter encoder) + static RepetitionLevelWriter getRootRepetitionLevelWriter(List repLevelWriterProviders, ColumnDescriptorValuesWriter encoder) { // Constructs hierarchy of RepetitionLevelWriter from leaf to root RepetitionLevelWriter rootRepetitionLevelWriter = Iterables.getLast(repLevelWriterProviders) diff --git a/lib/trino-parquet/src/main/java/io/trino/parquet/writer/repdef/RepLevelWriterProviders.java b/lib/trino-parquet/src/main/java/io/trino/parquet/writer/repdef/RepLevelWriterProviders.java index c40ccd0fa1cf..c46661ed6134 100644 --- a/lib/trino-parquet/src/main/java/io/trino/parquet/writer/repdef/RepLevelWriterProviders.java +++ b/lib/trino-parquet/src/main/java/io/trino/parquet/writer/repdef/RepLevelWriterProviders.java @@ -13,15 +13,18 @@ */ package io.trino.parquet.writer.repdef; +import io.trino.parquet.writer.valuewriter.ColumnDescriptorValuesWriter; +import io.trino.spi.block.ArrayBlock; import io.trino.spi.block.Block; import io.trino.spi.block.ColumnarArray; import io.trino.spi.block.ColumnarMap; -import io.trino.spi.block.ColumnarRow; -import org.apache.parquet.column.values.ValuesWriter; +import io.trino.spi.block.MapBlock; +import io.trino.spi.block.RowBlock; import java.util.Optional; import static com.google.common.base.Preconditions.checkArgument; +import static io.trino.spi.PageBlockUtil.getUnderlyingValueBlock; import static java.lang.String.format; import static java.util.Objects.requireNonNull; @@ -31,14 +34,12 @@ private RepLevelWriterProviders() {} public static RepLevelWriterProvider of(Block block) { + if (getUnderlyingValueBlock(block) instanceof RowBlock) { + return new RowRepLevelWriterProvider(block); + } return new PrimitiveRepLevelWriterProvider(block); } - public static RepLevelWriterProvider of(ColumnarRow columnarRow) - { - return new ColumnRowRepLevelWriterProvider(columnarRow); - } - public static RepLevelWriterProvider of(ColumnarArray columnarArray, int maxRepetitionLevel) { return new ColumnArrayRepLevelWriterProvider(columnarArray, maxRepetitionLevel); @@ -57,10 +58,13 @@ static class PrimitiveRepLevelWriterProvider PrimitiveRepLevelWriterProvider(Block block) { this.block = requireNonNull(block, "block is null"); + checkArgument(!(getUnderlyingValueBlock(block) instanceof RowBlock), "block is a row block"); + checkArgument(!(getUnderlyingValueBlock(block) instanceof ArrayBlock), "block is an array block"); + checkArgument(!(getUnderlyingValueBlock(block) instanceof MapBlock), "block is a map block"); } @Override - public RepetitionLevelWriter getRepetitionLevelWriter(Optional nestedWriter, ValuesWriter encoder) + public RepetitionLevelWriter getRepetitionLevelWriter(Optional nestedWriter, ColumnDescriptorValuesWriter encoder) { checkArgument(nestedWriter.isEmpty(), "nestedWriter should be empty for primitive repetition level writer"); return new RepetitionLevelWriter() @@ -77,27 +81,26 @@ public void writeRepetitionLevels(int parentLevel) public void writeRepetitionLevels(int parentLevel, int positionsCount) { checkValidPosition(offset, positionsCount, block.getPositionCount()); - for (int i = 0; i < positionsCount; i++) { - encoder.writeInteger(parentLevel); - } + encoder.writeRepeatInteger(parentLevel, positionsCount); offset += positionsCount; } }; } } - static class ColumnRowRepLevelWriterProvider + static class RowRepLevelWriterProvider implements RepLevelWriterProvider { - private final ColumnarRow columnarRow; + private final Block block; - ColumnRowRepLevelWriterProvider(ColumnarRow columnarRow) + RowRepLevelWriterProvider(Block block) { - this.columnarRow = requireNonNull(columnarRow, "columnarRow is null"); + this.block = requireNonNull(block, "block is null"); + checkArgument(getUnderlyingValueBlock(block) instanceof RowBlock, "block is not a row block"); } @Override - public RepetitionLevelWriter getRepetitionLevelWriter(Optional nestedWriterOptional, ValuesWriter encoder) + public RepetitionLevelWriter getRepetitionLevelWriter(Optional nestedWriterOptional, ColumnDescriptorValuesWriter encoder) { checkArgument(nestedWriterOptional.isPresent(), "nestedWriter should be present for column row repetition level writer"); return new RepetitionLevelWriter() @@ -109,28 +112,28 @@ public RepetitionLevelWriter getRepetitionLevelWriter(Optional nestedWriterOptional, ValuesWriter encoder) + public RepetitionLevelWriter getRepetitionLevelWriter(Optional nestedWriterOptional, ColumnDescriptorValuesWriter encoder) { checkArgument(nestedWriterOptional.isPresent(), "nestedWriter should be present for column map repetition level writer"); return new RepetitionLevelWriter() @@ -220,7 +223,7 @@ static class ColumnArrayRepLevelWriterProvider } @Override - public RepetitionLevelWriter getRepetitionLevelWriter(Optional nestedWriterOptional, ValuesWriter encoder) + public RepetitionLevelWriter getRepetitionLevelWriter(Optional nestedWriterOptional, ColumnDescriptorValuesWriter encoder) { checkArgument(nestedWriterOptional.isPresent(), "nestedWriter should be present for column map repetition level writer"); return new RepetitionLevelWriter() diff --git a/lib/trino-parquet/src/main/java/io/trino/parquet/writer/valuewriter/BloomFilterValuesWriter.java b/lib/trino-parquet/src/main/java/io/trino/parquet/writer/valuewriter/BloomFilterValuesWriter.java new file mode 100644 index 000000000000..2f1d44cafc2a --- /dev/null +++ b/lib/trino-parquet/src/main/java/io/trino/parquet/writer/valuewriter/BloomFilterValuesWriter.java @@ -0,0 +1,152 @@ +/* + * 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.parquet.writer.valuewriter; + +import com.google.common.annotations.VisibleForTesting; +import org.apache.parquet.bytes.BytesInput; +import org.apache.parquet.column.Encoding; +import org.apache.parquet.column.page.DictionaryPage; +import org.apache.parquet.column.values.ValuesWriter; +import org.apache.parquet.column.values.bloomfilter.BloomFilter; +import org.apache.parquet.io.api.Binary; + +import java.util.Optional; + +public class BloomFilterValuesWriter + extends ValuesWriter +{ + private final ValuesWriter writer; + private final BloomFilter bloomFilter; + + public static ValuesWriter createBloomFilterValuesWriter(ValuesWriter writer, Optional bloomFilter) + { + if (bloomFilter.isPresent()) { + return new BloomFilterValuesWriter(writer, bloomFilter.orElseThrow()); + } + return writer; + } + + private BloomFilterValuesWriter(ValuesWriter writer, BloomFilter bloomFilter) + { + this.writer = writer; + this.bloomFilter = bloomFilter; + } + + @VisibleForTesting + public ValuesWriter getWriter() + { + return writer; + } + + @Override + public long getBufferedSize() + { + return writer.getBufferedSize(); + } + + @Override + public BytesInput getBytes() + { + return writer.getBytes(); + } + + @Override + public Encoding getEncoding() + { + return writer.getEncoding(); + } + + @Override + public void reset() + { + writer.reset(); + } + + @Override + public void close() + { + writer.close(); + } + + @Override + public DictionaryPage toDictPageAndClose() + { + return writer.toDictPageAndClose(); + } + + @Override + public void resetDictionary() + { + writer.resetDictionary(); + } + + @Override + public long getAllocatedSize() + { + return writer.getAllocatedSize() + bloomFilter.getBitsetSize(); + } + + @Override + public void writeByte(int value) + { + throw new UnsupportedOperationException(); + } + + @Override + public void writeBoolean(boolean v) + { + throw new UnsupportedOperationException(); + } + + @Override + public void writeBytes(Binary v) + { + writer.writeBytes(v); + bloomFilter.insertHash(bloomFilter.hash(v)); + } + + @Override + public void writeInteger(int v) + { + writer.writeInteger(v); + bloomFilter.insertHash(bloomFilter.hash(v)); + } + + @Override + public void writeLong(long v) + { + writer.writeLong(v); + bloomFilter.insertHash(bloomFilter.hash(v)); + } + + @Override + public void writeDouble(double v) + { + writer.writeDouble(v); + bloomFilter.insertHash(bloomFilter.hash(v)); + } + + @Override + public void writeFloat(float v) + { + writer.writeFloat(v); + bloomFilter.insertHash(bloomFilter.hash(v)); + } + + @Override + public String memUsageString(String s) + { + return writer.memUsageString(s); + } +} diff --git a/lib/trino-parquet/src/main/java/io/trino/parquet/writer/valuewriter/ColumnDescriptorValuesWriter.java b/lib/trino-parquet/src/main/java/io/trino/parquet/writer/valuewriter/ColumnDescriptorValuesWriter.java new file mode 100644 index 000000000000..435a1e0ecfd2 --- /dev/null +++ b/lib/trino-parquet/src/main/java/io/trino/parquet/writer/valuewriter/ColumnDescriptorValuesWriter.java @@ -0,0 +1,82 @@ +/* + * 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.parquet.writer.valuewriter; + +import org.apache.parquet.bytes.BytesInput; +import org.apache.parquet.column.ColumnDescriptor; +import org.apache.parquet.column.Encoding; + +import static org.apache.parquet.bytes.BytesUtils.getWidthFromMaxInt; + +/** + * Used for writing repetition and definition levels + */ +public interface ColumnDescriptorValuesWriter +{ + /** + * @param value the value to encode + */ + void writeInteger(int value); + + /** + * @param value the value to encode + * @param valueRepetitions number of times the input value is repeated in the input stream + */ + void writeRepeatInteger(int value, int valueRepetitions); + + /** + * used to decide if we want to work to the next page + * + * @return the size of the currently buffered data (in bytes) + */ + long getBufferedSize(); + + /** + * @return the allocated size of the buffer + */ + long getAllocatedSize(); + + /** + * @return the bytes buffered so far to write to the current page + */ + BytesInput getBytes(); + + /** + * @return the encoding that was used to encode the bytes + */ + Encoding getEncoding(); + + /** + * called after getBytes() to reset the current buffer and start writing the next page + */ + void reset(); + + static ColumnDescriptorValuesWriter newRepetitionLevelWriter(ColumnDescriptor path, int pageSizeThreshold) + { + return newColumnDescriptorValuesWriter(path.getMaxRepetitionLevel(), pageSizeThreshold); + } + + static ColumnDescriptorValuesWriter newDefinitionLevelWriter(ColumnDescriptor path, int pageSizeThreshold) + { + return newColumnDescriptorValuesWriter(path.getMaxDefinitionLevel(), pageSizeThreshold); + } + + private static ColumnDescriptorValuesWriter newColumnDescriptorValuesWriter(int maxLevel, int pageSizeThreshold) + { + if (maxLevel == 0) { + return new DevNullValuesWriter(); + } + return new RunLengthBitPackingHybridValuesWriter(getWidthFromMaxInt(maxLevel), pageSizeThreshold); + } +} diff --git a/lib/trino-parquet/src/main/java/io/trino/parquet/writer/valuewriter/DevNullValuesWriter.java b/lib/trino-parquet/src/main/java/io/trino/parquet/writer/valuewriter/DevNullValuesWriter.java new file mode 100644 index 000000000000..ed54be167d94 --- /dev/null +++ b/lib/trino-parquet/src/main/java/io/trino/parquet/writer/valuewriter/DevNullValuesWriter.java @@ -0,0 +1,60 @@ +/* + * 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.parquet.writer.valuewriter; + +import org.apache.parquet.bytes.BytesInput; +import org.apache.parquet.column.Encoding; + +/** + * This is a special writer that doesn't write anything. The idea being that + * some columns will always be the same value, and this will capture that. An + * example is the set of repetition levels for a schema with no repeated fields. + */ +public class DevNullValuesWriter + implements ColumnDescriptorValuesWriter +{ + @Override + public long getBufferedSize() + { + return 0; + } + + @Override + public void reset() {} + + @Override + public void writeInteger(int v) {} + + @Override + public void writeRepeatInteger(int value, int valueRepetitions) {} + + @Override + public BytesInput getBytes() + { + return BytesInput.empty(); + } + + @Override + public long getAllocatedSize() + { + return 0; + } + + @Override + @SuppressWarnings("deprecation") + public Encoding getEncoding() + { + return Encoding.BIT_PACKED; + } +} diff --git a/lib/trino-parquet/src/main/java/io/trino/parquet/writer/valuewriter/DictionaryFallbackValuesWriter.java b/lib/trino-parquet/src/main/java/io/trino/parquet/writer/valuewriter/DictionaryFallbackValuesWriter.java new file mode 100644 index 000000000000..305a591f4fea --- /dev/null +++ b/lib/trino-parquet/src/main/java/io/trino/parquet/writer/valuewriter/DictionaryFallbackValuesWriter.java @@ -0,0 +1,234 @@ +/* + * 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.parquet.writer.valuewriter; + +import com.google.common.annotations.VisibleForTesting; +import jakarta.annotation.Nullable; +import org.apache.parquet.bytes.BytesInput; +import org.apache.parquet.column.Encoding; +import org.apache.parquet.column.page.DictionaryPage; +import org.apache.parquet.column.values.ValuesWriter; +import org.apache.parquet.column.values.dictionary.DictionaryValuesWriter; +import org.apache.parquet.io.api.Binary; + +import static com.google.common.base.Verify.verify; +import static java.util.Objects.requireNonNull; + +/** + * Based on org.apache.parquet.column.values.fallback.FallbackValuesWriter + */ +public class DictionaryFallbackValuesWriter + extends ValuesWriter +{ + private final ValuesWriter fallBackWriter; + + private boolean fellBackAlready; + private ValuesWriter currentWriter; + @Nullable + private DictionaryValuesWriter initialWriter; + private boolean initialUsedAndHadDictionary; + /* size of raw data, even if dictionary is used, it will not have effect on raw data size, it is used to decide + * if fall back to plain encoding is better by comparing rawDataByteSize with Encoded data size + * It's also used in getBufferedSize, so the page will be written based on raw data size + */ + private long rawDataByteSize; + // indicates if this is the first page being processed + private boolean firstPage = true; + + public DictionaryFallbackValuesWriter(DictionaryValuesWriter initialWriter, ValuesWriter fallBackWriter) + { + super(); + this.initialWriter = initialWriter; + this.fallBackWriter = fallBackWriter; + this.currentWriter = initialWriter; + } + + @Override + public long getBufferedSize() + { + // use raw data size to decide if we want to flush the page + // so the actual size of the page written could be much more smaller + // due to dictionary encoding. This prevents page being too big when fallback happens. + return rawDataByteSize; + } + + @Override + public BytesInput getBytes() + { + if (!fellBackAlready && firstPage) { + // we use the first page to decide if we're going to use this encoding + BytesInput bytes = initialWriter.getBytes(); + if (!initialWriter.isCompressionSatisfying(rawDataByteSize, bytes.size())) { + fallBack(); + // Since fallback happened on first page itself, we can drop the contents of initialWriter + initialWriter.close(); + initialWriter = null; + verify(!initialUsedAndHadDictionary, "initialUsedAndHadDictionary should be false when falling back to PLAIN in first page"); + } + else { + return bytes; + } + } + return currentWriter.getBytes(); + } + + @Override + public Encoding getEncoding() + { + Encoding encoding = currentWriter.getEncoding(); + if (!fellBackAlready && !initialUsedAndHadDictionary) { + initialUsedAndHadDictionary = encoding.usesDictionary(); + } + return encoding; + } + + @Override + public void reset() + { + rawDataByteSize = 0; + firstPage = false; + currentWriter.reset(); + } + + @Override + public void close() + { + if (initialWriter != null) { + initialWriter.close(); + } + fallBackWriter.close(); + } + + @Override + public DictionaryPage toDictPageAndClose() + { + if (initialUsedAndHadDictionary) { + return initialWriter.toDictPageAndClose(); + } + else { + return currentWriter.toDictPageAndClose(); + } + } + + @Override + public void resetDictionary() + { + if (initialUsedAndHadDictionary) { + initialWriter.resetDictionary(); + } + else { + currentWriter.resetDictionary(); + } + currentWriter = initialWriter; + fellBackAlready = false; + initialUsedAndHadDictionary = false; + firstPage = true; + } + + @Override + public long getAllocatedSize() + { + return fallBackWriter.getAllocatedSize() + (initialWriter != null ? initialWriter.getAllocatedSize() : 0); + } + + @Override + public String memUsageString(String prefix) + { + return String.format( + "%s FallbackValuesWriter{\n" + + "%s\n" + + "%s\n" + + "%s}\n", + prefix, + initialWriter != null ? initialWriter.memUsageString(prefix + " initial:") : "", + fallBackWriter.memUsageString(prefix + " fallback:"), + prefix); + } + + // passthrough writing the value + @Override + public void writeByte(int value) + { + rawDataByteSize += Byte.BYTES; + currentWriter.writeByte(value); + checkFallback(); + } + + @Override + public void writeBytes(Binary value) + { + // For raw data, length(4 bytes int) is stored, followed by the binary content itself + rawDataByteSize += value.length() + Integer.BYTES; + currentWriter.writeBytes(value); + checkFallback(); + } + + @Override + public void writeInteger(int value) + { + rawDataByteSize += Integer.BYTES; + currentWriter.writeInteger(value); + checkFallback(); + } + + @Override + public void writeLong(long value) + { + rawDataByteSize += Long.BYTES; + currentWriter.writeLong(value); + checkFallback(); + } + + @Override + public void writeFloat(float value) + { + rawDataByteSize += Float.BYTES; + currentWriter.writeFloat(value); + checkFallback(); + } + + @Override + public void writeDouble(double value) + { + rawDataByteSize += Double.BYTES; + currentWriter.writeDouble(value); + checkFallback(); + } + + @VisibleForTesting + public DictionaryValuesWriter getInitialWriter() + { + return requireNonNull(initialWriter, "initialWriter is null"); + } + + @VisibleForTesting + public ValuesWriter getFallBackWriter() + { + return fallBackWriter; + } + + private void checkFallback() + { + if (!fellBackAlready && initialWriter.shouldFallBack()) { + fallBack(); + } + } + + private void fallBack() + { + fellBackAlready = true; + initialWriter.fallBackAllValuesTo(fallBackWriter); + currentWriter = fallBackWriter; + } +} diff --git a/lib/trino-parquet/src/main/java/io/trino/parquet/writer/valuewriter/RunLengthBitPackingHybridEncoder.java b/lib/trino-parquet/src/main/java/io/trino/parquet/writer/valuewriter/RunLengthBitPackingHybridEncoder.java new file mode 100644 index 000000000000..d3482ff0eaeb --- /dev/null +++ b/lib/trino-parquet/src/main/java/io/trino/parquet/writer/valuewriter/RunLengthBitPackingHybridEncoder.java @@ -0,0 +1,318 @@ +/* + * 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.parquet.writer.valuewriter; + +import org.apache.parquet.bytes.BytesInput; +import org.apache.parquet.bytes.BytesUtils; +import org.apache.parquet.bytes.CapacityByteArrayOutputStream; +import org.apache.parquet.bytes.HeapByteBufferAllocator; +import org.apache.parquet.column.values.bitpacking.BytePacker; +import org.apache.parquet.column.values.bitpacking.Packer; +import org.apache.parquet.column.values.rle.RunLengthBitPackingHybridValuesWriter; + +import java.io.IOException; + +import static com.google.common.base.Preconditions.checkArgument; + +/** + * Encodes values using a combination of run length encoding and bit packing, + * according to the following grammar: + * + * 

+ * {@code
+ * rle-bit-packed-hybrid:  
+ * length := length of the  in bytes stored as 4 bytes little endian
+ * encoded-data := *
+ * run :=  | 
+ * bit-packed-run :=  
+ * bit-packed-header := varint-encode( << 1 | 1)
+ * // we always bit-pack a multiple of 8 values at a time, so we only store the number of values / 8
+ * bit-pack-count := (number of values in this run) / 8
+ * bit-packed-values :=  bit packed back to back, from LSB to MSB
+ * rle-run :=  
+ * rle-header := varint-encode( (number of times repeated) << 1)
+ * repeated-value := value that is repeated, using a fixed-width of round-up-to-next-byte(bit-width)
+ * }
+ *
+ * NOTE: this class is only responsible for creating and returning the {@code } + * portion of the above grammar. The {@code } portion is done by + * {@link RunLengthBitPackingHybridValuesWriter} + *

+ * Only supports positive values (including 0) + */ +public class RunLengthBitPackingHybridEncoder +{ + private static final int INITIAL_SLAB_SIZE = 64; + + private final BytePacker packer; + private final CapacityByteArrayOutputStream baos; + + /** + * The bit width used for bit-packing and for writing + * the repeated-value + */ + private final int bitWidth; + /** + * Values that are bit-packed 8 at a time are packed into this + * buffer, which is then written to baos + */ + private final byte[] packBuffer; + /** + * Previous value written, used to detect repeated values + */ + private int previousValue; + + /** + * We buffer 8 values at a time, and either bit pack them + * or discard them after writing a rle-run + */ + private final int[] bufferedValues; + private int numBufferedValues; + + /** + * How many times a value has been repeated + */ + private int repeatCount; + /** + * How many groups of 8 values have been written + * to the current bit-packed-run + */ + private int bitPackedGroupCount; + + /** + * A "pointer" to a single byte in baos, + * which we use as our bit-packed-header. It's really + * the logical index of the byte in baos. + *

+ * We are only using one byte for this header, + * which limits us to writing 504 values per bit-packed-run. + *

+ * MSB must be 0 for varint encoding, LSB must be 1 to signify + * that this is a bit-packed-header leaves 6 bits to write the + * number of 8-groups -> (2^6 - 1) * 8 = 504 + */ + private long bitPackedRunHeaderPointer; + private boolean toBytesCalled; + + public RunLengthBitPackingHybridEncoder(int bitWidth, int maxCapacityHint) + { + checkArgument(bitWidth >= 0 && bitWidth <= 32, "bitWidth must be >= 0 and <= 32"); + + this.bitWidth = bitWidth; + this.baos = new CapacityByteArrayOutputStream(INITIAL_SLAB_SIZE, maxCapacityHint, new HeapByteBufferAllocator()); + this.packBuffer = new byte[bitWidth]; + this.bufferedValues = new int[8]; + this.packer = Packer.LITTLE_ENDIAN.newBytePacker(bitWidth); + reset(false); + } + + private void reset(boolean resetBaos) + { + if (resetBaos) { + this.baos.reset(); + } + this.previousValue = 0; + this.numBufferedValues = 0; + this.repeatCount = 0; + this.bitPackedGroupCount = 0; + this.bitPackedRunHeaderPointer = -1; + this.toBytesCalled = false; + } + + public void writeInt(int value) + throws IOException + { + writeRepeatedInteger(value, 1); + } + + public void writeRepeatedInteger(int value, int valueRepetitions) + throws IOException + { + if (valueRepetitions == 0) { + return; + } + // Process 1st occurrence of new value + if (value != previousValue) { + // This is a new value, check if it signals the end of an rle-run + if (repeatCount >= 8) { + // it does! write an rle-run + writeRleRun(); + } + + // this is a new value so we've only seen it once + repeatCount = 1; + valueRepetitions--; + // start tracking this value for repeats + previousValue = value; + + bufferedValues[numBufferedValues++] = value; + if (numBufferedValues == 8) { + // we've encountered less than 8 repeated values, so + // either start a new bit-packed-run or append to the + // current bit-packed-run + writeOrAppendBitPackedRun(); + // we're going to see this value at least 8 times, so + // just count remaining repeats for an rle-run + if (valueRepetitions >= 8) { + repeatCount = valueRepetitions; + return; + } + } + } + + // Process remaining repetitions of value + while (valueRepetitions > 0) { + repeatCount++; + valueRepetitions--; + if (repeatCount >= 8) { + // we've seen this at least 8 times, we're + // certainly going to write an rle-run, + // so just keep on counting repeats for now + repeatCount += valueRepetitions; + return; + } + + bufferedValues[numBufferedValues++] = value; + if (numBufferedValues == 8) { + // we've encountered less than 8 repeated values, so + // either start a new bit-packed-run or append to the + // current bit-packed-run + writeOrAppendBitPackedRun(); + if (valueRepetitions >= 8) { + // we're going to see this value at least 8 times, so + // just count remaining repeats for an rle-run + repeatCount = valueRepetitions; + return; + } + } + } + } + + private void writeOrAppendBitPackedRun() + throws IOException + { + if (bitPackedGroupCount >= 63) { + // we've packed as many values as we can for this run, + // end it and start a new one + endPreviousBitPackedRun(); + } + + if (bitPackedRunHeaderPointer == -1) { + // this is a new bit-packed-run, allocate a byte for the header + // and keep a "pointer" to it so that it can be mutated later + baos.write(0); // write a sentinel value + bitPackedRunHeaderPointer = baos.getCurrentIndex(); + } + + packer.pack8Values(bufferedValues, 0, packBuffer, 0); + baos.write(packBuffer); + + // empty the buffer, they've all been written + numBufferedValues = 0; + + // clear the repeat count, as some repeated values + // may have just been bit packed into this run + repeatCount = 0; + + ++bitPackedGroupCount; + } + + /** + * If we are currently writing a bit-packed-run, update the + * bit-packed-header and consider this run to be over + *

+ * does nothing if we're not currently writing a bit-packed run + */ + private void endPreviousBitPackedRun() + { + if (bitPackedRunHeaderPointer == -1) { + // we're not currently in a bit-packed-run + return; + } + + // create bit-packed-header, which needs to fit in 1 byte + byte bitPackHeader = (byte) ((bitPackedGroupCount << 1) | 1); + + // update this byte + baos.setByte(bitPackedRunHeaderPointer, bitPackHeader); + + // mark that this run is over + bitPackedRunHeaderPointer = -1; + + // reset the number of groups + bitPackedGroupCount = 0; + } + + private void writeRleRun() + throws IOException + { + // we may have been working on a bit-packed-run + // so close that run if it exists before writing this + // rle-run + endPreviousBitPackedRun(); + + // write the rle-header (lsb of 0 signifies a rle run) + BytesUtils.writeUnsignedVarInt(repeatCount << 1, baos); + // write the repeated-value + BytesUtils.writeIntLittleEndianPaddedOnBitWidth(baos, previousValue, bitWidth); + + // reset the repeat count + repeatCount = 0; + + // throw away all the buffered values, they were just repeats and they've been written + numBufferedValues = 0; + } + + public BytesInput toBytes() + throws IOException + { + checkArgument(!toBytesCalled, "You cannot call toBytes() more than once without calling reset()"); + + // write anything that is buffered / queued up for an rle-run + if (repeatCount >= 8) { + writeRleRun(); + } + else if (numBufferedValues > 0) { + for (int i = numBufferedValues; i < 8; i++) { + bufferedValues[i] = 0; + } + writeOrAppendBitPackedRun(); + endPreviousBitPackedRun(); + } + else { + endPreviousBitPackedRun(); + } + + toBytesCalled = true; + return BytesInput.from(baos); + } + + /** + * Reset this encoder for re-use + */ + public void reset() + { + reset(true); + } + + public long getBufferedSize() + { + return baos.size(); + } + + public long getAllocatedSize() + { + return baos.getCapacity(); + } +} diff --git a/lib/trino-parquet/src/main/java/io/trino/parquet/writer/valuewriter/RunLengthBitPackingHybridValuesWriter.java b/lib/trino-parquet/src/main/java/io/trino/parquet/writer/valuewriter/RunLengthBitPackingHybridValuesWriter.java new file mode 100644 index 000000000000..4129b72ce481 --- /dev/null +++ b/lib/trino-parquet/src/main/java/io/trino/parquet/writer/valuewriter/RunLengthBitPackingHybridValuesWriter.java @@ -0,0 +1,93 @@ +/* + * 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.parquet.writer.valuewriter; + +import org.apache.parquet.bytes.BytesInput; +import org.apache.parquet.column.Encoding; +import org.apache.parquet.io.ParquetEncodingException; + +import java.io.IOException; + +import static java.lang.Math.toIntExact; +import static org.apache.parquet.column.Encoding.RLE; + +public class RunLengthBitPackingHybridValuesWriter + implements ColumnDescriptorValuesWriter +{ + private final RunLengthBitPackingHybridEncoder encoder; + + public RunLengthBitPackingHybridValuesWriter(int bitWidth, int maxCapacityHint) + { + this.encoder = new RunLengthBitPackingHybridEncoder(bitWidth, maxCapacityHint); + } + + @Override + public void writeInteger(int value) + { + try { + encoder.writeInt(value); + } + catch (IOException e) { + throw new ParquetEncodingException(e); + } + } + + @Override + public void writeRepeatInteger(int value, int valueRepetitions) + { + try { + encoder.writeRepeatedInteger(value, valueRepetitions); + } + catch (IOException e) { + throw new ParquetEncodingException(e); + } + } + + @Override + public long getBufferedSize() + { + return encoder.getBufferedSize(); + } + + @Override + public long getAllocatedSize() + { + return encoder.getAllocatedSize(); + } + + @Override + public BytesInput getBytes() + { + try { + // prepend the length of the column + BytesInput rle = encoder.toBytes(); + return BytesInput.concat(BytesInput.fromInt(toIntExact(rle.size())), rle); + } + catch (IOException e) { + throw new ParquetEncodingException(e); + } + } + + @Override + public Encoding getEncoding() + { + return RLE; + } + + @Override + public void reset() + { + encoder.reset(); + } +} diff --git a/lib/trino-parquet/src/main/java/io/trino/parquet/writer/valuewriter/TrinoValuesWriterFactory.java b/lib/trino-parquet/src/main/java/io/trino/parquet/writer/valuewriter/TrinoValuesWriterFactory.java new file mode 100644 index 000000000000..bbcc966d2763 --- /dev/null +++ b/lib/trino-parquet/src/main/java/io/trino/parquet/writer/valuewriter/TrinoValuesWriterFactory.java @@ -0,0 +1,140 @@ +/* + * 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.parquet.writer.valuewriter; + +import org.apache.parquet.bytes.HeapByteBufferAllocator; +import org.apache.parquet.column.ColumnDescriptor; +import org.apache.parquet.column.Encoding; +import org.apache.parquet.column.values.ValuesWriter; +import org.apache.parquet.column.values.bloomfilter.BloomFilter; +import org.apache.parquet.column.values.dictionary.DictionaryValuesWriter; +import org.apache.parquet.column.values.plain.BooleanPlainValuesWriter; +import org.apache.parquet.column.values.plain.FixedLenByteArrayPlainValuesWriter; +import org.apache.parquet.column.values.plain.PlainValuesWriter; + +import java.util.Optional; + +import static io.trino.parquet.writer.valuewriter.BloomFilterValuesWriter.createBloomFilterValuesWriter; +import static org.apache.parquet.column.Encoding.PLAIN_DICTIONARY; + +/** + * Based on org.apache.parquet.column.values.factory.DefaultV1ValuesWriterFactory + */ +public class TrinoValuesWriterFactory +{ + private static final int INITIAL_SLAB_SIZE = 64; + + private final int maxPageSize; + private final int maxDictionaryPageSize; + + public TrinoValuesWriterFactory(int maxPageSize, int maxDictionaryPageSize) + { + this.maxPageSize = maxPageSize; + this.maxDictionaryPageSize = maxDictionaryPageSize; + } + + public ValuesWriter newValuesWriter(ColumnDescriptor descriptor, Optional bloomFilter) + { + return switch (descriptor.getPrimitiveType().getPrimitiveTypeName()) { + case BOOLEAN -> new BooleanPlainValuesWriter(); // no dictionary encoding for boolean + case FIXED_LEN_BYTE_ARRAY -> getFixedLenByteArrayValuesWriter(descriptor, bloomFilter); + case BINARY -> getBinaryValuesWriter(descriptor, bloomFilter); + case INT32 -> getInt32ValuesWriter(descriptor, bloomFilter); + case INT64 -> getInt64ValuesWriter(descriptor, bloomFilter); + case INT96 -> getInt96ValuesWriter(descriptor, bloomFilter); + case DOUBLE -> getDoubleValuesWriter(descriptor, bloomFilter); + case FLOAT -> getFloatValuesWriter(descriptor, bloomFilter); + }; + } + + private ValuesWriter getFixedLenByteArrayValuesWriter(ColumnDescriptor path, Optional bloomFilter) + { + // dictionary encoding was not enabled in PARQUET 1.0 + return createBloomFilterValuesWriter(new FixedLenByteArrayPlainValuesWriter(path.getPrimitiveType().getTypeLength(), INITIAL_SLAB_SIZE, maxPageSize, new HeapByteBufferAllocator()), bloomFilter); + } + + private ValuesWriter getBinaryValuesWriter(ColumnDescriptor path, Optional bloomFilter) + { + ValuesWriter fallbackWriter = new PlainValuesWriter(INITIAL_SLAB_SIZE, maxPageSize, new HeapByteBufferAllocator()); + return createBloomFilterValuesWriter(dictWriterWithFallBack(path, getEncodingForDictionaryPage(), getEncodingForDataPage(), fallbackWriter), bloomFilter); + } + + private ValuesWriter getInt32ValuesWriter(ColumnDescriptor path, Optional bloomFilter) + { + ValuesWriter fallbackWriter = new PlainValuesWriter(INITIAL_SLAB_SIZE, maxPageSize, new HeapByteBufferAllocator()); + return createBloomFilterValuesWriter(dictWriterWithFallBack(path, getEncodingForDictionaryPage(), getEncodingForDataPage(), fallbackWriter), bloomFilter); + } + + private ValuesWriter getInt64ValuesWriter(ColumnDescriptor path, Optional bloomFilter) + { + ValuesWriter fallbackWriter = new PlainValuesWriter(INITIAL_SLAB_SIZE, maxPageSize, new HeapByteBufferAllocator()); + return createBloomFilterValuesWriter(dictWriterWithFallBack(path, getEncodingForDictionaryPage(), getEncodingForDataPage(), fallbackWriter), bloomFilter); + } + + private ValuesWriter getInt96ValuesWriter(ColumnDescriptor path, Optional bloomFilter) + { + ValuesWriter fallbackWriter = new FixedLenByteArrayPlainValuesWriter(12, INITIAL_SLAB_SIZE, maxPageSize, new HeapByteBufferAllocator()); + return createBloomFilterValuesWriter(dictWriterWithFallBack(path, getEncodingForDictionaryPage(), getEncodingForDataPage(), fallbackWriter), bloomFilter); + } + + private ValuesWriter getDoubleValuesWriter(ColumnDescriptor path, Optional bloomFilter) + { + ValuesWriter fallbackWriter = new PlainValuesWriter(INITIAL_SLAB_SIZE, maxPageSize, new HeapByteBufferAllocator()); + return createBloomFilterValuesWriter(dictWriterWithFallBack(path, getEncodingForDictionaryPage(), getEncodingForDataPage(), fallbackWriter), bloomFilter); + } + + private ValuesWriter getFloatValuesWriter(ColumnDescriptor path, Optional bloomFilter) + { + ValuesWriter fallbackWriter = new PlainValuesWriter(INITIAL_SLAB_SIZE, maxPageSize, new HeapByteBufferAllocator()); + return createBloomFilterValuesWriter(dictWriterWithFallBack(path, getEncodingForDictionaryPage(), getEncodingForDataPage(), fallbackWriter), bloomFilter); + } + + @SuppressWarnings("deprecation") + private static Encoding getEncodingForDataPage() + { + return PLAIN_DICTIONARY; + } + + @SuppressWarnings("deprecation") + private static Encoding getEncodingForDictionaryPage() + { + return PLAIN_DICTIONARY; + } + + private DictionaryValuesWriter dictionaryWriter(ColumnDescriptor path, Encoding dictPageEncoding, Encoding dataPageEncoding) + { + return switch (path.getPrimitiveType().getPrimitiveTypeName()) { + case BOOLEAN -> throw new IllegalArgumentException("no dictionary encoding for BOOLEAN"); + case BINARY -> + new DictionaryValuesWriter.PlainBinaryDictionaryValuesWriter(maxDictionaryPageSize, dataPageEncoding, dictPageEncoding, new HeapByteBufferAllocator()); + case INT32 -> + new DictionaryValuesWriter.PlainIntegerDictionaryValuesWriter(maxDictionaryPageSize, dataPageEncoding, dictPageEncoding, new HeapByteBufferAllocator()); + case INT64 -> + new DictionaryValuesWriter.PlainLongDictionaryValuesWriter(maxDictionaryPageSize, dataPageEncoding, dictPageEncoding, new HeapByteBufferAllocator()); + case INT96 -> + new DictionaryValuesWriter.PlainFixedLenArrayDictionaryValuesWriter(maxDictionaryPageSize, 12, dataPageEncoding, dictPageEncoding, new HeapByteBufferAllocator()); + case DOUBLE -> + new DictionaryValuesWriter.PlainDoubleDictionaryValuesWriter(maxDictionaryPageSize, dataPageEncoding, dictPageEncoding, new HeapByteBufferAllocator()); + case FLOAT -> + new DictionaryValuesWriter.PlainFloatDictionaryValuesWriter(maxDictionaryPageSize, dataPageEncoding, dictPageEncoding, new HeapByteBufferAllocator()); + case FIXED_LEN_BYTE_ARRAY -> + new DictionaryValuesWriter.PlainFixedLenArrayDictionaryValuesWriter(maxDictionaryPageSize, path.getPrimitiveType().getTypeLength(), dataPageEncoding, dictPageEncoding, new HeapByteBufferAllocator()); + }; + } + + private ValuesWriter dictWriterWithFallBack(ColumnDescriptor path, Encoding dictPageEncoding, Encoding dataPageEncoding, ValuesWriter writerToFallBackTo) + { + return new DictionaryFallbackValuesWriter(dictionaryWriter(path, dictPageEncoding, dataPageEncoding), writerToFallBackTo); + } +} diff --git a/lib/trino-parquet/src/test/java/io/trino/parquet/ParquetTestUtils.java b/lib/trino-parquet/src/test/java/io/trino/parquet/ParquetTestUtils.java index 06c9c79e5d30..18ee3e71561b 100644 --- a/lib/trino-parquet/src/test/java/io/trino/parquet/ParquetTestUtils.java +++ b/lib/trino-parquet/src/test/java/io/trino/parquet/ParquetTestUtils.java @@ -14,11 +14,14 @@ package io.trino.parquet; import com.google.common.collect.ImmutableList; -import com.google.common.primitives.Booleans; import io.airlift.slice.Slice; import io.airlift.slice.Slices; import io.trino.memory.context.AggregatedMemoryContext; +import io.trino.parquet.metadata.FileMetadata; +import io.trino.parquet.metadata.ParquetMetadata; +import io.trino.parquet.predicate.TupleDomainParquetPredicate; import io.trino.parquet.reader.ParquetReader; +import io.trino.parquet.reader.RowGroupInfo; import io.trino.parquet.writer.ParquetSchemaConverter; import io.trino.parquet.writer.ParquetWriter; import io.trino.parquet.writer.ParquetWriterOptions; @@ -26,13 +29,15 @@ import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; import io.trino.spi.block.LongArrayBlock; +import io.trino.spi.block.RowBlock; +import io.trino.spi.predicate.TupleDomain; import io.trino.spi.type.MapType; import io.trino.spi.type.Type; import io.trino.spi.type.TypeOperators; +import org.apache.parquet.column.ColumnDescriptor; import org.apache.parquet.format.CompressionCodec; -import org.apache.parquet.hadoop.metadata.BlockMetaData; -import org.apache.parquet.hadoop.metadata.ParquetMetadata; import org.apache.parquet.io.MessageColumnIO; +import org.apache.parquet.schema.MessageType; import org.joda.time.DateTimeZone; import java.io.ByteArrayOutputStream; @@ -40,6 +45,7 @@ import java.io.OutputStream; import java.util.Arrays; import java.util.List; +import java.util.Map; import java.util.Optional; import java.util.Random; @@ -48,14 +54,15 @@ import static com.google.common.collect.ImmutableList.toImmutableList; import static io.trino.parquet.ParquetTypeUtils.constructField; import static io.trino.parquet.ParquetTypeUtils.getColumnIO; +import static io.trino.parquet.ParquetTypeUtils.getDescriptors; import static io.trino.parquet.ParquetTypeUtils.lookupColumnByName; +import static io.trino.parquet.predicate.PredicateUtils.buildPredicate; +import static io.trino.parquet.predicate.PredicateUtils.getFilteredRowGroups; import static io.trino.spi.block.ArrayBlock.fromElementBlock; import static io.trino.spi.block.MapBlock.fromKeyValueBlock; -import static io.trino.spi.block.RowBlock.fromFieldBlocks; import static io.trino.spi.type.BigintType.BIGINT; import static io.trino.spi.type.TypeUtils.writeNativeValue; -import static java.lang.Math.toIntExact; -import static java.util.Collections.nCopies; +import static java.util.Locale.ENGLISH; import static org.joda.time.DateTimeZone.UTC; public class ParquetTestUtils @@ -69,7 +76,7 @@ public static Slice writeParquetFile(ParquetWriterOptions writerOptions, List types, List columnNames) + public static ParquetWriter createParquetWriter(OutputStream outputStream, ParquetWriterOptions writerOptions, List types, List columnNames, CompressionCodec compression) { checkArgument(types.size() == columnNames.size()); ParquetSchemaConverter schemaConverter = new ParquetSchemaConverter(types, columnNames, false, false); @@ -88,9 +95,8 @@ public static ParquetWriter createParquetWriter(OutputStream outputStream, Parqu schemaConverter.getMessageType(), schemaConverter.getPrimitiveTypes(), writerOptions, - CompressionCodec.SNAPPY, + compression, "test-version", - false, Optional.of(DateTimeZone.getDefault()), Optional.empty()); } @@ -103,37 +109,60 @@ public static ParquetReader createParquetReader( List columnNames) throws IOException { - org.apache.parquet.hadoop.metadata.FileMetaData fileMetaData = parquetMetadata.getFileMetaData(); - MessageColumnIO messageColumnIO = getColumnIO(fileMetaData.getSchema(), fileMetaData.getSchema()); - ImmutableList.Builder columnFields = ImmutableList.builder(); + return createParquetReader(input, parquetMetadata, new ParquetReaderOptions(), memoryContext, types, columnNames, TupleDomain.all()); + } + + public static ParquetReader createParquetReader( + ParquetDataSource input, + ParquetMetadata parquetMetadata, + ParquetReaderOptions options, + AggregatedMemoryContext memoryContext, + List types, + List columnNames, + TupleDomain predicate) + throws IOException + { + FileMetadata fileMetaData = parquetMetadata.getFileMetaData(); + MessageType fileSchema = fileMetaData.getSchema(); + MessageColumnIO messageColumnIO = getColumnIO(fileSchema, fileSchema); + ImmutableList.Builder columnFields = ImmutableList.builder(); for (int i = 0; i < types.size(); i++) { - columnFields.add(constructField( - types.get(i), - lookupColumnByName(messageColumnIO, columnNames.get(i))) - .orElseThrow()); + columnFields.add(new Column( + messageColumnIO.getName(), + constructField( + types.get(i), + lookupColumnByName(messageColumnIO, columnNames.get(i))) + .orElseThrow())); } - long nextStart = 0; - ImmutableList.Builder blockStartsBuilder = ImmutableList.builder(); - for (BlockMetaData block : parquetMetadata.getBlocks()) { - blockStartsBuilder.add(nextStart); - nextStart += block.getRowCount(); - } - List blockStarts = blockStartsBuilder.build(); + Map, ColumnDescriptor> descriptorsByPath = getDescriptors(fileSchema, fileSchema); + TupleDomain parquetTupleDomain = predicate.transformKeys( + columnName -> descriptorsByPath.get(ImmutableList.of(columnName.toLowerCase(ENGLISH)))); + TupleDomainParquetPredicate parquetPredicate = buildPredicate(fileSchema, parquetTupleDomain, descriptorsByPath, UTC); + List rowGroups = getFilteredRowGroups( + 0, + input.getEstimatedSize(), + input, + parquetMetadata, + ImmutableList.of(parquetTupleDomain), + ImmutableList.of(parquetPredicate), + descriptorsByPath, + UTC, + 1000, + options); return new ParquetReader( Optional.ofNullable(fileMetaData.getCreatedBy()), columnFields.build(), - parquetMetadata.getBlocks(), - blockStarts, + false, + rowGroups, input, UTC, memoryContext, - new ParquetReaderOptions(), + options, exception -> { throwIfUnchecked(exception); return new RuntimeException(exception); }, - Optional.empty(), - nCopies(blockStarts.size(), Optional.empty()), + Optional.of(parquetPredicate), Optional.empty()); } @@ -163,24 +192,26 @@ public static List generateGroupSizes(int positionsCount) return groupsBuilder.build(); } - public static Block createRowBlock(Optional rowIsNull, int positionCount) + public static RowBlock createRowBlock(Optional rowIsNull, int positionCount) { - int fieldPositionCount = rowIsNull.map(nulls -> toIntExact(Booleans.asList(nulls).stream().filter(isNull -> !isNull).count())) - .orElse(positionCount); - int fieldCount = 4; - Block[] fieldBlocks = new Block[fieldCount]; + // TODO test with nested null fields and without nulls + Block[] fieldBlocks = new Block[4]; // no nulls block - fieldBlocks[0] = new LongArrayBlock(fieldPositionCount, Optional.empty(), new long[fieldPositionCount]); + fieldBlocks[0] = new LongArrayBlock(positionCount, rowIsNull, new long[positionCount]); // no nulls with mayHaveNull block - fieldBlocks[1] = new LongArrayBlock(fieldPositionCount, Optional.of(new boolean[fieldPositionCount]), new long[fieldPositionCount]); + fieldBlocks[1] = new LongArrayBlock(positionCount, rowIsNull.or(() -> Optional.of(new boolean[positionCount])), new long[positionCount]); // all nulls block - boolean[] allNulls = new boolean[fieldPositionCount]; - Arrays.fill(allNulls, false); - fieldBlocks[2] = new LongArrayBlock(fieldPositionCount, Optional.of(allNulls), new long[fieldPositionCount]); + boolean[] allNulls = new boolean[positionCount]; + Arrays.fill(allNulls, true); + fieldBlocks[2] = new LongArrayBlock(positionCount, Optional.of(allNulls), new long[positionCount]); // random nulls block - fieldBlocks[3] = createLongsBlockWithRandomNulls(fieldPositionCount); + boolean[] valueIsNull = rowIsNull.map(boolean[]::clone).orElseGet(() -> new boolean[positionCount]); + for (int i = 0; i < positionCount; i++) { + valueIsNull[i] |= RANDOM.nextBoolean(); + } + fieldBlocks[3] = new LongArrayBlock(positionCount, Optional.of(valueIsNull), new long[positionCount]); - return fromFieldBlocks(positionCount, rowIsNull, fieldBlocks); + return RowBlock.fromNotNullSuppressedFieldBlocks(positionCount, rowIsNull, fieldBlocks); } public static Block createArrayBlock(Optional valueIsNull, int positionCount) diff --git a/lib/trino-parquet/src/test/java/io/trino/parquet/reader/AbstractColumnReaderBenchmark.java b/lib/trino-parquet/src/test/java/io/trino/parquet/reader/AbstractColumnReaderBenchmark.java index 44325171784f..1d9e92314a48 100644 --- a/lib/trino-parquet/src/test/java/io/trino/parquet/reader/AbstractColumnReaderBenchmark.java +++ b/lib/trino-parquet/src/test/java/io/trino/parquet/reader/AbstractColumnReaderBenchmark.java @@ -17,6 +17,7 @@ import io.airlift.slice.Slices; import io.trino.parquet.DataPage; import io.trino.parquet.DataPageV1; +import io.trino.parquet.ParquetDataSourceId; import io.trino.parquet.ParquetReaderOptions; import io.trino.parquet.PrimitiveField; import org.apache.parquet.column.values.ValuesWriter; @@ -61,7 +62,7 @@ public abstract class AbstractColumnReaderBenchmark private static final int DATA_GENERATION_BATCH_SIZE = 16384; private static final int READ_BATCH_SIZE = 4096; - private final ColumnReaderFactory columnReaderFactory = new ColumnReaderFactory(UTC, new ParquetReaderOptions().withBatchColumnReaders(true)); + private final ColumnReaderFactory columnReaderFactory = new ColumnReaderFactory(UTC, new ParquetReaderOptions()); private final List dataPages = new ArrayList<>(); private int dataPositions; @@ -104,7 +105,8 @@ public int read() throws IOException { ColumnReader columnReader = columnReaderFactory.create(field, newSimpleAggregatedMemoryContext()); - columnReader.setPageReader(new PageReader(UNCOMPRESSED, dataPages.iterator(), false, false), Optional.empty()); + PageReader pageReader = new PageReader(new ParquetDataSourceId("test"), UNCOMPRESSED, dataPages.iterator(), false, false); + columnReader.setPageReader(pageReader, Optional.empty()); int rowsRead = 0; while (rowsRead < dataPositions) { int remaining = dataPositions - rowsRead; diff --git a/lib/trino-parquet/src/test/java/io/trino/parquet/reader/AbstractColumnReaderRowRangesTest.java b/lib/trino-parquet/src/test/java/io/trino/parquet/reader/AbstractColumnReaderRowRangesTest.java index 1322f241082e..36391e6d2603 100644 --- a/lib/trino-parquet/src/test/java/io/trino/parquet/reader/AbstractColumnReaderRowRangesTest.java +++ b/lib/trino-parquet/src/test/java/io/trino/parquet/reader/AbstractColumnReaderRowRangesTest.java @@ -18,6 +18,7 @@ import io.trino.parquet.DataPage; import io.trino.parquet.DataPageV2; import io.trino.parquet.Page; +import io.trino.parquet.ParquetDataSourceId; import io.trino.parquet.PrimitiveField; import io.trino.parquet.reader.decoders.ValueDecoder; import io.trino.parquet.reader.decoders.ValueDecoders; @@ -558,6 +559,7 @@ else if (dictionaryEncoding == DictionaryEncoding.MIXED) { inputPages = ImmutableList.builder().add(toTrinoDictionaryPage(encoder.toDictPageAndClose())).addAll(inputPages).build(); } return new PageReader( + new ParquetDataSourceId("test"), UNCOMPRESSED, inputPages.iterator(), dictionaryEncoding == DictionaryEncoding.ALL || (dictionaryEncoding == DictionaryEncoding.MIXED && testingPages.size() == 1), diff --git a/lib/trino-parquet/src/test/java/io/trino/parquet/reader/AbstractColumnReaderTest.java b/lib/trino-parquet/src/test/java/io/trino/parquet/reader/AbstractColumnReaderTest.java index 0a6a09c3bac6..371faf2acbf0 100644 --- a/lib/trino-parquet/src/test/java/io/trino/parquet/reader/AbstractColumnReaderTest.java +++ b/lib/trino-parquet/src/test/java/io/trino/parquet/reader/AbstractColumnReaderTest.java @@ -23,6 +23,7 @@ import io.trino.parquet.DataPageV2; import io.trino.parquet.DictionaryPage; import io.trino.parquet.Page; +import io.trino.parquet.ParquetDataSourceId; import io.trino.parquet.ParquetEncoding; import io.trino.parquet.ParquetReaderOptions; import io.trino.parquet.PrimitiveField; @@ -538,7 +539,7 @@ public void testMemoryUsage(DataPageVersion version, ColumnReaderFormat f // Create reader PrimitiveField field = createField(format, true); AggregatedMemoryContext memoryContext = newSimpleAggregatedMemoryContext(); - ColumnReaderFactory columnReaderFactory = new ColumnReaderFactory(UTC, new ParquetReaderOptions().withBatchColumnReaders(true)); + ColumnReaderFactory columnReaderFactory = new ColumnReaderFactory(UTC, new ParquetReaderOptions()); ColumnReader reader = columnReaderFactory.create(field, memoryContext); // Write data DictionaryValuesWriter dictionaryWriter = format.getDictionaryWriter(); @@ -687,6 +688,7 @@ protected static PageReader getPageReaderMock(List dataPages, @Nullabl pagesBuilder.add(dictionaryPage); } return new PageReader( + new ParquetDataSourceId("test"), UNCOMPRESSED, pagesBuilder.addAll(dataPages).build().iterator(), dataPages.stream() diff --git a/lib/trino-parquet/src/test/java/io/trino/parquet/reader/TestColumnReaderFactory.java b/lib/trino-parquet/src/test/java/io/trino/parquet/reader/TestColumnReaderFactory.java index b4749d526b4c..260e01c15f7d 100644 --- a/lib/trino-parquet/src/test/java/io/trino/parquet/reader/TestColumnReaderFactory.java +++ b/lib/trino-parquet/src/test/java/io/trino/parquet/reader/TestColumnReaderFactory.java @@ -30,41 +30,30 @@ public class TestColumnReaderFactory { @Test - public void testUseBatchedColumnReaders() + public void testTopLevelPrimitiveFields() { - PrimitiveField field = new PrimitiveField( + ColumnReaderFactory columnReaderFactory = new ColumnReaderFactory(UTC, new ParquetReaderOptions()); + PrimitiveType primitiveType = new PrimitiveType(OPTIONAL, INT32, "test"); + + PrimitiveField topLevelRepeatedPrimitiveField = new PrimitiveField( INTEGER, - false, - new ColumnDescriptor(new String[] {"test"}, new PrimitiveType(OPTIONAL, INT32, "test"), 0, 1), + true, + new ColumnDescriptor(new String[] {"topLevelRepeatedPrimitiveField test"}, primitiveType, 1, 1), 0); - ColumnReaderFactory columnReaderFactory = new ColumnReaderFactory(UTC, new ParquetReaderOptions().withBatchColumnReaders(false)); - assertThat(columnReaderFactory.create(field, newSimpleAggregatedMemoryContext())) - .isNotInstanceOf(AbstractColumnReader.class); - columnReaderFactory = new ColumnReaderFactory(UTC, new ParquetReaderOptions().withBatchColumnReaders(true)); - assertThat(columnReaderFactory.create(field, newSimpleAggregatedMemoryContext())) - .isInstanceOf(FlatColumnReader.class); - } + assertThat(columnReaderFactory.create(topLevelRepeatedPrimitiveField, newSimpleAggregatedMemoryContext())).isInstanceOf(NestedColumnReader.class); - @Test - public void testNestedColumnReaders() - { - PrimitiveField field = new PrimitiveField( + PrimitiveField topLevelOptionalPrimitiveField = new PrimitiveField( INTEGER, false, - new ColumnDescriptor(new String[] {"level1", "level2"}, new PrimitiveType(OPTIONAL, INT32, "test"), 1, 2), + new ColumnDescriptor(new String[] {"topLevelRequiredPrimitiveField test"}, primitiveType, 0, 1), 0); - ColumnReaderFactory columnReaderFactory = new ColumnReaderFactory(UTC, new ParquetReaderOptions().withBatchColumnReaders(false)); - assertThat(columnReaderFactory.create(field, newSimpleAggregatedMemoryContext())) - .isNotInstanceOf(AbstractColumnReader.class); - columnReaderFactory = new ColumnReaderFactory(UTC, new ParquetReaderOptions().withBatchColumnReaders(false).withBatchNestedColumnReaders(true)); - assertThat(columnReaderFactory.create(field, newSimpleAggregatedMemoryContext())) - .isNotInstanceOf(AbstractColumnReader.class); + assertThat(columnReaderFactory.create(topLevelOptionalPrimitiveField, newSimpleAggregatedMemoryContext())).isInstanceOf(FlatColumnReader.class); - columnReaderFactory = new ColumnReaderFactory(UTC, new ParquetReaderOptions().withBatchColumnReaders(true)); - assertThat(columnReaderFactory.create(field, newSimpleAggregatedMemoryContext())) - .isInstanceOf(NestedColumnReader.class); - columnReaderFactory = new ColumnReaderFactory(UTC, new ParquetReaderOptions().withBatchColumnReaders(true).withBatchNestedColumnReaders(true)); - assertThat(columnReaderFactory.create(field, newSimpleAggregatedMemoryContext())) - .isInstanceOf(NestedColumnReader.class); + PrimitiveField topLevelRequiredPrimitiveField = new PrimitiveField( + INTEGER, + true, + new ColumnDescriptor(new String[] {"topLevelRequiredPrimitiveField test"}, primitiveType, 0, 0), + 0); + assertThat(columnReaderFactory.create(topLevelRequiredPrimitiveField, newSimpleAggregatedMemoryContext())).isInstanceOf(FlatColumnReader.class); } } diff --git a/lib/trino-parquet/src/test/java/io/trino/parquet/reader/TestInt96Timestamp.java b/lib/trino-parquet/src/test/java/io/trino/parquet/reader/TestInt96Timestamp.java index cd0e32f93a42..dc131498ee8e 100644 --- a/lib/trino-parquet/src/test/java/io/trino/parquet/reader/TestInt96Timestamp.java +++ b/lib/trino-parquet/src/test/java/io/trino/parquet/reader/TestInt96Timestamp.java @@ -17,6 +17,7 @@ import io.airlift.slice.Slices; import io.trino.parquet.DataPage; import io.trino.parquet.DataPageV2; +import io.trino.parquet.ParquetDataSourceId; import io.trino.parquet.ParquetReaderOptions; import io.trino.parquet.PrimitiveField; import io.trino.plugin.base.type.DecodedTimestamp; @@ -107,11 +108,10 @@ public void testVariousTimestamps(TimestampType type, BiFunction columnNames = ImmutableList.of("columnA", "columnB"); - List types = ImmutableList.of(INTEGER, BIGINT); - - ParquetDataSource dataSource = new TestingParquetDataSource( - writeParquetFile( - ParquetWriterOptions.builder() - .setMaxBlockSize(DataSize.ofBytes(1000)) - .build(), - types, - columnNames, - generateInputPages(types, 100, 5)), - new ParquetReaderOptions()); - ParquetMetadata parquetMetadata = MetadataReader.readFooter(dataSource, Optional.empty()); - assertThat(parquetMetadata.getBlocks().size()).isGreaterThan(1); - // Verify file has only non-dictionary encodings as dictionary memory usage is already tested in TestFlatColumnReader#testMemoryUsage - parquetMetadata.getBlocks().forEach(block -> { - block.getColumns() - .forEach(columnChunkMetaData -> assertThat(columnChunkMetaData.getEncodingStats().hasDictionaryEncodedPages()).isFalse()); - assertThat(block.getRowCount()).isEqualTo(100); - }); - - AggregatedMemoryContext memoryContext = newSimpleAggregatedMemoryContext(); - ParquetReader reader = createParquetReader(dataSource, parquetMetadata, memoryContext, types, columnNames); - - Page page = reader.nextPage(); - assertThat(page.getBlock(0)).isInstanceOf(LazyBlock.class); - assertThat(memoryContext.getBytes()).isEqualTo(0); - page.getBlock(0).getLoadedBlock(); - // Memory usage due to reading data and decoding parquet page of 1st block - long initialMemoryUsage = memoryContext.getBytes(); - assertThat(initialMemoryUsage).isGreaterThan(0); - - // Memory usage due to decoding parquet page of 2nd block - page.getBlock(1).getLoadedBlock(); - long currentMemoryUsage = memoryContext.getBytes(); - assertThat(currentMemoryUsage).isGreaterThan(initialMemoryUsage); - - // Memory usage does not change until next row group (1 page per row-group) - long rowGroupRowCount = parquetMetadata.getBlocks().get(0).getRowCount(); - int rowsRead = page.getPositionCount(); - while (rowsRead < rowGroupRowCount) { - rowsRead += reader.nextPage().getPositionCount(); - if (rowsRead <= rowGroupRowCount) { - assertThat(memoryContext.getBytes()).isEqualTo(currentMemoryUsage); - } - } - - // New row-group should release memory from old column readers, while using some memory for data read - reader.nextPage(); - assertThat(memoryContext.getBytes()).isBetween(1L, currentMemoryUsage - 1); - - reader.close(); - assertThat(memoryContext.getBytes()).isEqualTo(0); - } -} diff --git a/lib/trino-parquet/src/test/java/io/trino/parquet/reader/TestTimeMillis.java b/lib/trino-parquet/src/test/java/io/trino/parquet/reader/TestTimeMillis.java index cd781878c1db..7ae81a8de4d5 100644 --- a/lib/trino-parquet/src/test/java/io/trino/parquet/reader/TestTimeMillis.java +++ b/lib/trino-parquet/src/test/java/io/trino/parquet/reader/TestTimeMillis.java @@ -17,12 +17,12 @@ import com.google.common.io.Resources; import io.trino.parquet.ParquetDataSource; import io.trino.parquet.ParquetReaderOptions; +import io.trino.parquet.metadata.ParquetMetadata; import io.trino.spi.Page; import io.trino.spi.block.Block; import io.trino.spi.type.SqlTime; import io.trino.spi.type.TimeType; import io.trino.spi.type.Type; -import org.apache.parquet.hadoop.metadata.ParquetMetadata; import org.testng.annotations.DataProvider; import org.testng.annotations.Test; diff --git a/lib/trino-parquet/src/test/java/io/trino/parquet/reader/flat/TestFlatColumnReader.java b/lib/trino-parquet/src/test/java/io/trino/parquet/reader/flat/TestFlatColumnReader.java index 06457048249e..347831081368 100644 --- a/lib/trino-parquet/src/test/java/io/trino/parquet/reader/flat/TestFlatColumnReader.java +++ b/lib/trino-parquet/src/test/java/io/trino/parquet/reader/flat/TestFlatColumnReader.java @@ -17,6 +17,7 @@ import io.airlift.slice.Slices; import io.trino.parquet.DataPage; import io.trino.parquet.DataPageV1; +import io.trino.parquet.ParquetDataSourceId; import io.trino.parquet.ParquetEncoding; import io.trino.parquet.ParquetReaderOptions; import io.trino.parquet.PrimitiveField; @@ -63,7 +64,7 @@ public class TestFlatColumnReader @Override protected ColumnReader createColumnReader(PrimitiveField field) { - ColumnReaderFactory columnReaderFactory = new ColumnReaderFactory(UTC, new ParquetReaderOptions().withBatchColumnReaders(true)); + ColumnReaderFactory columnReaderFactory = new ColumnReaderFactory(UTC, new ParquetReaderOptions()); ColumnReader columnReader = columnReaderFactory.create(field, newSimpleAggregatedMemoryContext()); assertThat(columnReader).isInstanceOf(FlatColumnReader.class); return columnReader; @@ -137,7 +138,7 @@ private static PageReader getSimplePageReaderMock(ParquetEncoding encoding) encoding, encoding, PLAIN)); - return new PageReader(UNCOMPRESSED, pages.iterator(), false, false); + return new PageReader(new ParquetDataSourceId("test"), UNCOMPRESSED, pages.iterator(), false, false); } private static PageReader getNullOnlyPageReaderMock() @@ -154,6 +155,6 @@ private static PageReader getNullOnlyPageReaderMock() RLE, RLE, PLAIN)); - return new PageReader(UNCOMPRESSED, pages.iterator(), false, false); + return new PageReader(new ParquetDataSourceId("test"), UNCOMPRESSED, pages.iterator(), false, false); } } diff --git a/lib/trino-parquet/src/test/java/io/trino/parquet/writer/NullsProvider.java b/lib/trino-parquet/src/test/java/io/trino/parquet/writer/NullsProvider.java index af204bb669db..f0b383626468 100644 --- a/lib/trino-parquet/src/test/java/io/trino/parquet/writer/NullsProvider.java +++ b/lib/trino-parquet/src/test/java/io/trino/parquet/writer/NullsProvider.java @@ -79,6 +79,26 @@ Optional getNulls(int positionCount) abstract Optional getNulls(int positionCount); + Optional getNulls(int positionCount, Optional forcedNulls) + { + Optional nulls = getNulls(positionCount); + if (forcedNulls.isEmpty()) { + return nulls; + } + if (nulls.isEmpty()) { + return forcedNulls; + } + + boolean[] nullPositions = nulls.get(); + boolean[] forcedNullPositions = forcedNulls.get(); + for (int i = 0; i < positionCount; i++) { + if (forcedNullPositions[i]) { + nullPositions[i] = true; + } + } + return Optional.of(nullPositions); + } + @DataProvider public static Object[][] nullsProviders() { diff --git a/lib/trino-parquet/src/test/java/io/trino/parquet/writer/TestDefinitionLevelWriter.java b/lib/trino-parquet/src/test/java/io/trino/parquet/writer/TestDefinitionLevelWriter.java index f73064df1f3e..3cab34ede5f5 100644 --- a/lib/trino-parquet/src/test/java/io/trino/parquet/writer/TestDefinitionLevelWriter.java +++ b/lib/trino-parquet/src/test/java/io/trino/parquet/writer/TestDefinitionLevelWriter.java @@ -18,8 +18,8 @@ import io.trino.spi.block.Block; import io.trino.spi.block.ColumnarArray; import io.trino.spi.block.ColumnarMap; -import io.trino.spi.block.ColumnarRow; import io.trino.spi.block.LongArrayBlock; +import io.trino.spi.block.RowBlock; import org.testng.annotations.Test; import java.util.List; @@ -34,7 +34,7 @@ import static io.trino.parquet.writer.repdef.DefLevelWriterProvider.getRootDefinitionLevelWriter; import static io.trino.spi.block.ColumnarArray.toColumnarArray; import static io.trino.spi.block.ColumnarMap.toColumnarMap; -import static io.trino.spi.block.ColumnarRow.toColumnarRow; +import static io.trino.spi.block.RowBlock.getNullSuppressedRowFieldsFromBlock; import static java.util.Collections.nCopies; import static org.assertj.core.api.Assertions.assertThat; @@ -60,28 +60,30 @@ public void testWritePrimitiveDefinitionLevels(NullsProvider nullsProvider) @Test(dataProviderClass = NullsProvider.class, dataProvider = "nullsProviders") public void testWriteRowDefinitionLevels(NullsProvider nullsProvider) { - Block rowBlock = createRowBlock(nullsProvider.getNulls(POSITIONS), POSITIONS); - ColumnarRow columnarRow = toColumnarRow(rowBlock); + RowBlock rowBlock = createRowBlock(nullsProvider.getNulls(POSITIONS), POSITIONS); + List fields = getNullSuppressedRowFieldsFromBlock(rowBlock); int fieldMaxDefinitionLevel = 2; // Write definition levels for all positions - for (int field = 0; field < columnarRow.getFieldCount(); field++) { - assertDefinitionLevels(columnarRow, ImmutableList.of(), field, fieldMaxDefinitionLevel); + for (int field = 0; field < fields.size(); field++) { + assertDefinitionLevels(rowBlock, fields, ImmutableList.of(), field, fieldMaxDefinitionLevel); } // Write definition levels for all positions one-at-a-time - for (int field = 0; field < columnarRow.getFieldCount(); field++) { + for (int field = 0; field < fields.size(); field++) { assertDefinitionLevels( - columnarRow, - nCopies(columnarRow.getPositionCount(), 1), + rowBlock, + fields, + nCopies(rowBlock.getPositionCount(), 1), field, fieldMaxDefinitionLevel); } // Write definition levels for all positions with different group sizes - for (int field = 0; field < columnarRow.getFieldCount(); field++) { + for (int field = 0; field < fields.size(); field++) { assertDefinitionLevels( - columnarRow, - generateGroupSizes(columnarRow.getPositionCount()), + rowBlock, + fields, + generateGroupSizes(rowBlock.getPositionCount()), field, fieldMaxDefinitionLevel); } @@ -178,7 +180,8 @@ private static void assertDefinitionLevels(Block block, List writePosit } private static void assertDefinitionLevels( - ColumnarRow columnarRow, + RowBlock block, + List nullSuppressedFields, List writePositionCounts, int field, int maxDefinitionLevel) @@ -187,8 +190,8 @@ private static void assertDefinitionLevels( TestingValuesWriter valuesWriter = new TestingValuesWriter(); DefinitionLevelWriter fieldRootDefLevelWriter = getRootDefinitionLevelWriter( ImmutableList.of( - DefLevelWriterProviders.of(columnarRow, maxDefinitionLevel - 1), - DefLevelWriterProviders.of(columnarRow.getField(field), maxDefinitionLevel)), + DefLevelWriterProviders.of(block, maxDefinitionLevel - 1), + DefLevelWriterProviders.of(nullSuppressedFields.get(field), maxDefinitionLevel)), valuesWriter); ValuesCount fieldValuesCount; if (writePositionCounts.isEmpty()) { @@ -209,12 +212,12 @@ private static void assertDefinitionLevels( int maxDefinitionValuesCount = 0; ImmutableList.Builder expectedDefLevelsBuilder = ImmutableList.builder(); int fieldOffset = 0; - for (int position = 0; position < columnarRow.getPositionCount(); position++) { - if (columnarRow.isNull(position)) { + for (int position = 0; position < block.getPositionCount(); position++) { + if (block.isNull(position)) { expectedDefLevelsBuilder.add(maxDefinitionLevel - 2); continue; } - Block fieldBlock = columnarRow.getField(field); + Block fieldBlock = nullSuppressedFields.get(field); if (fieldBlock.isNull(fieldOffset)) { expectedDefLevelsBuilder.add(maxDefinitionLevel - 1); } @@ -224,7 +227,7 @@ private static void assertDefinitionLevels( } fieldOffset++; } - assertThat(fieldValuesCount.totalValuesCount()).isEqualTo(columnarRow.getPositionCount()); + assertThat(fieldValuesCount.totalValuesCount()).isEqualTo(block.getPositionCount()); assertThat(fieldValuesCount.maxDefinitionLevelValuesCount()).isEqualTo(maxDefinitionValuesCount); assertThat(valuesWriter.getWrittenValues()).isEqualTo(expectedDefLevelsBuilder.build()); } diff --git a/lib/trino-parquet/src/test/java/io/trino/parquet/writer/TestParquetWriter.java b/lib/trino-parquet/src/test/java/io/trino/parquet/writer/TestParquetWriter.java index 8c951759122c..c16208eec731 100644 --- a/lib/trino-parquet/src/test/java/io/trino/parquet/writer/TestParquetWriter.java +++ b/lib/trino-parquet/src/test/java/io/trino/parquet/writer/TestParquetWriter.java @@ -19,7 +19,11 @@ import io.trino.parquet.DataPage; import io.trino.parquet.DiskRange; import io.trino.parquet.ParquetDataSource; +import io.trino.parquet.ParquetDataSourceId; import io.trino.parquet.ParquetReaderOptions; +import io.trino.parquet.metadata.BlockMetadata; +import io.trino.parquet.metadata.ColumnChunkMetadata; +import io.trino.parquet.metadata.ParquetMetadata; import io.trino.parquet.reader.ChunkedInputStream; import io.trino.parquet.reader.MetadataReader; import io.trino.parquet.reader.PageReader; @@ -28,9 +32,7 @@ import io.trino.spi.type.Type; import org.apache.parquet.VersionParser; import org.apache.parquet.column.ColumnDescriptor; -import org.apache.parquet.hadoop.metadata.BlockMetaData; -import org.apache.parquet.hadoop.metadata.ColumnChunkMetaData; -import org.apache.parquet.hadoop.metadata.ParquetMetadata; +import org.apache.parquet.format.CompressionCodec; import org.apache.parquet.schema.PrimitiveType; import org.testng.annotations.Test; @@ -90,13 +92,14 @@ public void testWrittenPageSize() new ParquetReaderOptions()); ParquetMetadata parquetMetadata = MetadataReader.readFooter(dataSource, Optional.empty()); assertThat(parquetMetadata.getBlocks().size()).isEqualTo(1); - assertThat(parquetMetadata.getBlocks().get(0).getRowCount()).isEqualTo(100 * 1000); + assertThat(parquetMetadata.getBlocks().get(0).rowCount()).isEqualTo(100 * 1000); - ColumnChunkMetaData chunkMetaData = parquetMetadata.getBlocks().get(0).getColumns().get(0); + ColumnChunkMetadata chunkMetaData = parquetMetadata.getBlocks().get(0).columns().get(0); DiskRange range = new DiskRange(chunkMetaData.getStartingPos(), chunkMetaData.getTotalSize()); Map chunkReader = dataSource.planRead(ImmutableListMultimap.of(0, range), newSimpleAggregatedMemoryContext()); PageReader pageReader = PageReader.createPageReader( + new ParquetDataSourceId("test"), chunkReader.get(0), chunkMetaData, new ColumnDescriptor(new String[] {"columna"}, new PrimitiveType(REQUIRED, INT32, "columna"), 0, 0), @@ -138,10 +141,10 @@ public void testColumnReordering() ParquetMetadata parquetMetadata = MetadataReader.readFooter(dataSource, Optional.empty()); assertThat(parquetMetadata.getBlocks().size()).isGreaterThanOrEqualTo(10); - for (BlockMetaData blockMetaData : parquetMetadata.getBlocks()) { + for (BlockMetadata blockMetaData : parquetMetadata.getBlocks()) { // Verify that the columns are stored in the same order as the metadata - List offsets = blockMetaData.getColumns().stream() - .map(ColumnChunkMetaData::getFirstDataPageOffset) + List offsets = blockMetaData.columns().stream() + .map(ColumnChunkMetadata::getFirstDataPageOffset) .collect(toImmutableList()); assertThat(offsets).isSorted(); } @@ -160,7 +163,8 @@ public void testWriterMemoryAccounting() .setMaxPageSize(DataSize.ofBytes(1024)) .build(), types, - columnNames); + columnNames, + CompressionCodec.SNAPPY); List inputPages = generateInputPages(types, 1000, 100); long previousRetainedBytes = 0; diff --git a/lib/trino-parquet/src/test/java/io/trino/parquet/writer/TestRepetitionLevelWriter.java b/lib/trino-parquet/src/test/java/io/trino/parquet/writer/TestRepetitionLevelWriter.java index d35ff53ca756..6f484b365128 100644 --- a/lib/trino-parquet/src/test/java/io/trino/parquet/writer/TestRepetitionLevelWriter.java +++ b/lib/trino-parquet/src/test/java/io/trino/parquet/writer/TestRepetitionLevelWriter.java @@ -14,12 +14,10 @@ package io.trino.parquet.writer; import com.google.common.collect.ImmutableList; -import com.google.common.primitives.Booleans; import io.trino.parquet.writer.repdef.RepLevelWriterProviders; import io.trino.spi.block.Block; import io.trino.spi.block.ColumnarArray; import io.trino.spi.block.ColumnarMap; -import io.trino.spi.block.ColumnarRow; import io.trino.spi.type.MapType; import io.trino.spi.type.TypeOperators; import org.testng.annotations.Test; @@ -39,11 +37,10 @@ import static io.trino.spi.block.ArrayBlock.fromElementBlock; import static io.trino.spi.block.ColumnarArray.toColumnarArray; import static io.trino.spi.block.ColumnarMap.toColumnarMap; -import static io.trino.spi.block.ColumnarRow.toColumnarRow; import static io.trino.spi.block.MapBlock.fromKeyValueBlock; import static io.trino.spi.block.RowBlock.fromFieldBlocks; +import static io.trino.spi.block.RowBlock.getNullSuppressedRowFieldsFromBlock; import static io.trino.spi.type.BigintType.BIGINT; -import static java.lang.Math.toIntExact; import static java.util.Collections.nCopies; import static org.assertj.core.api.Assertions.assertThat; @@ -62,23 +59,25 @@ public void testWriteRowRepetitionLevels(NullsProvider nullsProvider) Block arrayBlock = fromElementBlock(POSITIONS, valueIsNull, arrayOffsets, rowBlock); ColumnarArray columnarArray = toColumnarArray(arrayBlock); - ColumnarRow columnarRow = toColumnarRow(columnarArray.getElementsBlock()); + Block row = columnarArray.getElementsBlock(); + List nullSuppressedFields = getNullSuppressedRowFieldsFromBlock(row); // Write Repetition levels for all positions - for (int field = 0; field < columnarRow.getFieldCount(); field++) { - assertRepetitionLevels(columnarArray, columnarRow, field, ImmutableList.of()); - assertRepetitionLevels(columnarArray, columnarRow, field, ImmutableList.of()); + for (int fieldIndex = 0; fieldIndex < nullSuppressedFields.size(); fieldIndex++) { + Block field = nullSuppressedFields.get(fieldIndex); + assertRepetitionLevels(columnarArray, row, field, ImmutableList.of()); + assertRepetitionLevels(columnarArray, row, field, ImmutableList.of()); // Write Repetition levels for all positions one-at-a-time assertRepetitionLevels( columnarArray, - columnarRow, + row, field, nCopies(columnarArray.getPositionCount(), 1)); // Write Repetition levels for all positions with different group sizes assertRepetitionLevels( columnarArray, - columnarRow, + row, field, generateGroupSizes(columnarArray.getPositionCount())); } @@ -118,10 +117,10 @@ public void testWriteMapRepetitionLevels(NullsProvider nullsProvider) public void testNestedStructRepetitionLevels(NullsProvider nullsProvider) { Block rowBlock = createNestedRowBlock(nullsProvider.getNulls(POSITIONS), POSITIONS); - ColumnarRow columnarRow = toColumnarRow(rowBlock); + List fieldBlocks = getNullSuppressedRowFieldsFromBlock(rowBlock); - for (int field = 0; field < columnarRow.getFieldCount(); field++) { - Block fieldBlock = columnarRow.getField(field); + for (int field = 0; field < fieldBlocks.size(); field++) { + Block fieldBlock = fieldBlocks.get(field); ColumnarMap columnarMap = toColumnarMap(fieldBlock); for (Block mapElements : ImmutableList.of(columnarMap.getKeysBlock(), columnarMap.getValuesBlock())) { ColumnarArray columnarArray = toColumnarArray(mapElements); @@ -130,23 +129,21 @@ public void testNestedStructRepetitionLevels(NullsProvider nullsProvider) assertRepetitionLevels(rowBlock, columnarMap, columnarArray, ImmutableList.of()); // Write Repetition levels for all positions one-at-a-time - assertRepetitionLevels(rowBlock, columnarMap, columnarArray, nCopies(columnarRow.getPositionCount(), 1)); + assertRepetitionLevels(rowBlock, columnarMap, columnarArray, nCopies(rowBlock.getPositionCount(), 1)); // Write Repetition levels for all positions with different group sizes - assertRepetitionLevels(rowBlock, columnarMap, columnarArray, generateGroupSizes(columnarRow.getPositionCount())); + assertRepetitionLevels(rowBlock, columnarMap, columnarArray, generateGroupSizes(rowBlock.getPositionCount())); } } } private static Block createNestedRowBlock(Optional rowIsNull, int positionCount) { - int fieldPositionCount = rowIsNull.map(nulls -> toIntExact(Booleans.asList(nulls).stream().filter(isNull -> !isNull).count())) - .orElse(positionCount); Block[] fieldBlocks = new Block[2]; // no nulls map block - fieldBlocks[0] = createMapOfArraysBlock(Optional.empty(), fieldPositionCount); + fieldBlocks[0] = createMapOfArraysBlock(rowIsNull, positionCount); // random nulls map block - fieldBlocks[1] = createMapOfArraysBlock(RANDOM_NULLS.getNulls(fieldPositionCount), fieldPositionCount); + fieldBlocks[1] = createMapOfArraysBlock(RANDOM_NULLS.getNulls(positionCount, rowIsNull), positionCount); return fromFieldBlocks(positionCount, rowIsNull, fieldBlocks); } @@ -162,8 +159,8 @@ private static Block createMapOfArraysBlock(Optional mapIsNull, int p private static void assertRepetitionLevels( ColumnarArray columnarArray, - ColumnarRow columnarRow, - int field, + Block row, + Block field, List writePositionCounts) { int maxRepetitionLevel = 1; @@ -172,8 +169,8 @@ private static void assertRepetitionLevels( RepetitionLevelWriter fieldRootRepLevelWriter = getRootRepetitionLevelWriter( ImmutableList.of( RepLevelWriterProviders.of(columnarArray, maxRepetitionLevel), - RepLevelWriterProviders.of(columnarRow), - RepLevelWriterProviders.of(columnarRow.getField(field))), + RepLevelWriterProviders.of(row), + RepLevelWriterProviders.of(field)), valuesWriter); if (writePositionCounts.isEmpty()) { fieldRootRepLevelWriter.writeRepetitionLevels(0); @@ -188,7 +185,7 @@ private static void assertRepetitionLevels( Iterator expectedRepetitionLevelsIter = RepLevelIterables.getIterator(ImmutableList.builder() .add(RepLevelIterables.of(columnarArray, maxRepetitionLevel)) .add(RepLevelIterables.of(columnarArray.getElementsBlock())) - .add(RepLevelIterables.of(columnarRow.getField(field))) + .add(RepLevelIterables.of(field)) .build()); assertThat(valuesWriter.getWrittenValues()).isEqualTo(ImmutableList.copyOf(expectedRepetitionLevelsIter)); } @@ -302,7 +299,7 @@ private static void assertRepetitionLevels( TestingValuesWriter valuesWriter = new TestingValuesWriter(); RepetitionLevelWriter fieldRootRepLevelWriter = getRootRepetitionLevelWriter( ImmutableList.of( - RepLevelWriterProviders.of(toColumnarRow(rowBlock)), + RepLevelWriterProviders.of(rowBlock), RepLevelWriterProviders.of(columnarMap, 1), RepLevelWriterProviders.of(columnarArray, 2), RepLevelWriterProviders.of(columnarArray.getElementsBlock())), diff --git a/lib/trino-parquet/src/test/java/io/trino/parquet/writer/TestingValuesWriter.java b/lib/trino-parquet/src/test/java/io/trino/parquet/writer/TestingValuesWriter.java index 53d6dd89b681..9370d01764f5 100644 --- a/lib/trino-parquet/src/test/java/io/trino/parquet/writer/TestingValuesWriter.java +++ b/lib/trino-parquet/src/test/java/io/trino/parquet/writer/TestingValuesWriter.java @@ -13,16 +13,16 @@ */ package io.trino.parquet.writer; +import io.trino.parquet.writer.valuewriter.ColumnDescriptorValuesWriter; import it.unimi.dsi.fastutil.ints.IntArrayList; import it.unimi.dsi.fastutil.ints.IntList; import org.apache.parquet.bytes.BytesInput; import org.apache.parquet.column.Encoding; -import org.apache.parquet.column.values.ValuesWriter; import java.util.List; class TestingValuesWriter - extends ValuesWriter + implements ColumnDescriptorValuesWriter { private final IntList values = new IntArrayList(); @@ -57,15 +57,17 @@ public long getAllocatedSize() } @Override - public String memUsageString(String prefix) + public void writeInteger(int v) { - throw new UnsupportedOperationException(); + values.add(v); } @Override - public void writeInteger(int v) + public void writeRepeatInteger(int value, int valueRepetitions) { - values.add(v); + for (int i = 0; i < valueRepetitions; i++) { + values.add(value); + } } List getWrittenValues() diff --git a/lib/trino-plugin-toolkit/src/main/java/io/trino/plugin/base/ClosingBinder.java b/lib/trino-plugin-toolkit/src/main/java/io/trino/plugin/base/ClosingBinder.java new file mode 100644 index 000000000000..bd2b40dea484 --- /dev/null +++ b/lib/trino-plugin-toolkit/src/main/java/io/trino/plugin/base/ClosingBinder.java @@ -0,0 +1,144 @@ +/* + * 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.base; + +import com.google.common.collect.ImmutableSet; +import com.google.inject.Binder; +import com.google.inject.BindingAnnotation; +import com.google.inject.Inject; +import com.google.inject.Injector; +import com.google.inject.Key; +import com.google.inject.Provider; +import com.google.inject.multibindings.Multibinder; +import io.trino.plugin.base.util.AutoCloseableCloser; +import jakarta.annotation.PreDestroy; + +import java.lang.annotation.Retention; +import java.lang.annotation.Target; +import java.util.Set; +import java.util.concurrent.ExecutorService; +import java.util.function.Consumer; + +import static com.google.common.base.Verify.verifyNotNull; +import static com.google.inject.multibindings.Multibinder.newSetBinder; +import static java.lang.annotation.ElementType.FIELD; +import static java.lang.annotation.ElementType.METHOD; +import static java.lang.annotation.ElementType.PARAMETER; +import static java.lang.annotation.RetentionPolicy.RUNTIME; +import static java.util.Objects.requireNonNull; + +public class ClosingBinder +{ + public static ClosingBinder closingBinder(Binder binder) + { + return new ClosingBinder(binder); + } + + private final Multibinder executors; + private final Multibinder closeables; + + private ClosingBinder(Binder binder) + { + executors = newSetBinder(binder, ExecutorService.class, ForCleanup.class); + closeables = newSetBinder(binder, AutoCloseable.class, ForCleanup.class); + binder.bind(Cleanup.class).asEagerSingleton(); + } + + public void registerExecutor(Class type) + { + registerExecutor(Key.get(type)); + } + + public void registerExecutor(Key key) + { + executors.addBinding().to(requireNonNull(key, "key is null")); + } + + public void registerCloseable(Class type) + { + registerCloseable(Key.get(type)); + } + + public void registerCloseable(Key key) + { + closeables.addBinding().to(key); + } + + public void registerResource(Class type, Consumer close) + { + registerResource(Key.get(type), close); + } + + public void registerResource(Key key, Consumer close) + { + closeables.addBinding().toProvider(new ResourceCloser(key, close)); + } + + private static class ResourceCloser + implements Provider + { + private final Key key; + private final Consumer close; + private Injector injector; + + private ResourceCloser(Key key, Consumer close) + { + this.key = requireNonNull(key, "key is null"); + this.close = requireNonNull(close, "close is null"); + } + + @Inject + public void setInjector(Injector injector) + { + this.injector = injector; + } + + @Override + public AutoCloseable get() + { + T object = injector.getInstance(key); + verifyNotNull(object, "null at key %s", key); + Consumer close = this.close; + return () -> close.accept(object); + } + } + + private record Cleanup( + @ForCleanup Set executors, + @ForCleanup Set closeables) + { + @Inject + private Cleanup + { + executors = ImmutableSet.copyOf(executors); + closeables = ImmutableSet.copyOf(closeables); + } + + @PreDestroy + public void shutdown() + throws Exception + { + try (var closer = AutoCloseableCloser.create()) { + // TODO should this await termination? + executors.forEach(executor -> closer.register(executor::shutdownNow)); + closeables.forEach(closer::register); + } + } + } + + @Retention(RUNTIME) + @Target({FIELD, PARAMETER, METHOD}) + @BindingAnnotation + private @interface ForCleanup {} +} diff --git a/lib/trino-plugin-toolkit/src/main/java/io/trino/plugin/base/util/AutoCloseableCloser.java b/lib/trino-plugin-toolkit/src/main/java/io/trino/plugin/base/util/AutoCloseableCloser.java new file mode 100644 index 000000000000..f3d42e1337a3 --- /dev/null +++ b/lib/trino-plugin-toolkit/src/main/java/io/trino/plugin/base/util/AutoCloseableCloser.java @@ -0,0 +1,100 @@ +/* + * 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.base.util; + +import com.google.common.collect.Lists; +import com.google.errorprone.annotations.concurrent.GuardedBy; + +import java.util.ArrayList; +import java.util.List; + +import static com.google.common.base.Throwables.throwIfInstanceOf; +import static com.google.common.base.Throwables.throwIfUnchecked; +import static java.util.Objects.requireNonNull; + +/** + * This class is inspired by com.google.common.io.Closer + */ +public final class AutoCloseableCloser + implements AutoCloseable +{ + @GuardedBy("this") + private boolean closed; + @GuardedBy("this") + private final ArrayList closeables = new ArrayList<>(4); + + private AutoCloseableCloser() {} + + public static AutoCloseableCloser create() + { + return new AutoCloseableCloser(); + } + + public C register(C closeable) + { + requireNonNull(closeable, "closeable is null"); + boolean registered = false; + synchronized (this) { + if (!closed) { + closeables.add(closeable); + registered = true; + } + } + if (!registered) { + IllegalStateException failure = new IllegalStateException("Already closed"); + try { + closeable.close(); + } + catch (Exception e) { + failure.addSuppressed(e); + } + throw failure; + } + return closeable; + } + + @Override + public void close() + throws Exception + { + List closeables; + synchronized (this) { + closed = true; + closeables = List.copyOf(Lists.reverse(this.closeables)); + this.closeables.clear(); + this.closeables.trimToSize(); + } + Throwable rootCause = null; + for (AutoCloseable closeable : closeables) { + try { + closeable.close(); + } + catch (Throwable t) { + if (rootCause == null) { + rootCause = t; + } + else if (rootCause != t) { + // Self-suppression not permitted + rootCause.addSuppressed(t); + } + } + } + if (rootCause != null) { + throwIfInstanceOf(rootCause, Exception.class); + throwIfUnchecked(rootCause); + // not possible + throw new AssertionError(rootCause); + } + } +} diff --git a/lib/trino-plugin-toolkit/src/main/java/io/trino/plugin/base/util/ExecutorUtil.java b/lib/trino-plugin-toolkit/src/main/java/io/trino/plugin/base/util/ExecutorUtil.java new file mode 100644 index 000000000000..c9c5c24465f3 --- /dev/null +++ b/lib/trino-plugin-toolkit/src/main/java/io/trino/plugin/base/util/ExecutorUtil.java @@ -0,0 +1,140 @@ +/* + * 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.base.util; + +import com.google.errorprone.annotations.ThreadSafe; +import com.google.errorprone.annotations.concurrent.GuardedBy; + +import java.util.ArrayList; +import java.util.Collection; +import java.util.List; +import java.util.concurrent.Callable; +import java.util.concurrent.CompletionService; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.Executor; +import java.util.concurrent.ExecutorCompletionService; +import java.util.concurrent.Future; + +import static com.google.common.collect.ImmutableList.toImmutableList; +import static java.util.Collections.nCopies; +import static java.util.Objects.requireNonNull; + +public final class ExecutorUtil +{ + private ExecutorUtil() {} + + /** + * Process tasks in executors and additionally in calling thread. + * Upon task execution failure, other tasks are canceled and interrupted, but not waited + * for. + *

+ * This method propagates {@link Context#current()} into tasks it starts within the executor. + *

+ * Note: using this method allows simple parallelization of tasks within executor, when sub-tasks + * are also scheduled in that executor, without risking starvation when pool is saturated. + * + * @throws ExecutionException if any task fails; exception cause is the first task failure + */ + public static List processWithAdditionalThreads(Collection> tasks, Executor executor) + throws ExecutionException + { + List> wrapped = tasks.stream() + .map(Task::new) + .collect(toImmutableList()); + CompletionService> completionService = new ExecutorCompletionService<>(executor); + List> futures = new ArrayList<>(wrapped.size()); + + try { + // schedule in the executor + for (int i = 0; i < wrapped.size(); i++) { + int index = i; + Task task = wrapped.get(i); + futures.add(completionService.submit(() -> { + if (!task.take()) { + return null; // will be ignored + } + return new TaskResult<>(index, task.callable.call()); + })); + } + + List results = new ArrayList<>(nCopies(wrapped.size(), null)); + int pending = wrapped.size(); + // process in the calling thread (in reverse order, as an optimization) + for (int i = wrapped.size() - 1; i >= 0; i--) { + // process ready results to fail fast on exceptions + for (Future> ready = completionService.poll(); ready != null; ready = completionService.poll()) { + TaskResult taskResult = ready.get(); + // Null result means task was processed by the calling thread + if (taskResult != null) { + results.set(taskResult.taskIndex(), taskResult.result()); + pending--; + } + } + Task task = wrapped.get(i); + if (!task.take()) { + continue; + } + try { + results.set(i, task.callable.call()); + pending--; + } + catch (Exception e) { + throw new ExecutionException(e); + } + } + + while (pending > 0) { + TaskResult taskResult = completionService.take().get(); + // Null result means task was processed by the calling thread + if (taskResult != null) { + results.set(taskResult.taskIndex(), taskResult.result()); + pending--; + } + } + + return results; + } + catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException("Interrupted", e); + } + finally { + futures.forEach(future -> future.cancel(true)); + } + } + + @ThreadSafe + private static final class Task + { + private final Callable callable; + @GuardedBy("this") + private boolean taken; + + public Task(Callable callable) + { + this.callable = requireNonNull(callable, "callable is null"); + } + + public synchronized boolean take() + { + if (taken) { + return false; + } + taken = true; + return true; + } + } + + private record TaskResult(int taskIndex, T result) {} +} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveMaterializedViewMetadataFactory.java b/lib/trino-plugin-toolkit/src/main/java/io/trino/plugin/base/util/UncheckedCloseable.java similarity index 77% rename from plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveMaterializedViewMetadataFactory.java rename to lib/trino-plugin-toolkit/src/main/java/io/trino/plugin/base/util/UncheckedCloseable.java index aea473edb441..2e91f3bf37ba 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveMaterializedViewMetadataFactory.java +++ b/lib/trino-plugin-toolkit/src/main/java/io/trino/plugin/base/util/UncheckedCloseable.java @@ -11,9 +11,11 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package io.trino.plugin.hive; +package io.trino.plugin.base.util; -public interface HiveMaterializedViewMetadataFactory +public interface UncheckedCloseable + extends AutoCloseable { - HiveMaterializedViewMetadata create(HiveMetastoreClosure hiveMetastoreClosure); + @Override + void close(); } diff --git a/plugin/trino-bigquery/pom.xml b/plugin/trino-bigquery/pom.xml index f558faf3af25..89e3f151c8f3 100644 --- a/plugin/trino-bigquery/pom.xml +++ b/plugin/trino-bigquery/pom.xml @@ -36,7 +36,6 @@ org.apache.commons commons-lang3 - 3.11 diff --git a/plugin/trino-delta-lake/pom.xml b/plugin/trino-delta-lake/pom.xml index d0c01736b5bf..724cf7744832 100644 --- a/plugin/trino-delta-lake/pom.xml +++ b/plugin/trino-delta-lake/pom.xml @@ -27,16 +27,6 @@ - - com.amazonaws - aws-java-sdk-core - - - - com.amazonaws - aws-java-sdk-glue - - com.fasterxml.jackson.core jackson-core @@ -194,11 +184,6 @@ parquet-format-structures - - org.apache.parquet - parquet-hadoop - - org.roaringbitmap RoaringBitmap @@ -263,6 +248,12 @@ runtime + + software.amazon.awssdk + glue + runtime + + com.azure azure-core diff --git a/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/AbstractDeltaLakePageSink.java b/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/AbstractDeltaLakePageSink.java index 8c75ea33ce11..178d2b1b58e7 100644 --- a/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/AbstractDeltaLakePageSink.java +++ b/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/AbstractDeltaLakePageSink.java @@ -443,7 +443,8 @@ private FileWriter createParquetFileWriter(Location path) .setMaxBlockSize(getParquetWriterBlockSize(session)) .setMaxPageSize(getParquetWriterPageSize(session)) .build(); - CompressionCodec compressionCodec = getCompressionCodec(session).getParquetCompressionCodec(); + CompressionCodec compressionCodec = getCompressionCodec(session).getParquetCompressionCodec() + .orElseThrow(); // validated on the session property level try { Closeable rollbackAction = () -> fileSystem.deleteFile(path); @@ -470,7 +471,6 @@ private FileWriter createParquetFileWriter(Location path) identityMapping, compressionCodec, trinoVersion, - false, Optional.empty(), Optional.empty()); } diff --git a/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/DeltaLakeMergeSink.java b/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/DeltaLakeMergeSink.java index ec1b074bd40d..128ddcce3fcf 100644 --- a/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/DeltaLakeMergeSink.java +++ b/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/DeltaLakeMergeSink.java @@ -55,6 +55,7 @@ import java.util.List; import java.util.Map; import java.util.Optional; +import java.util.OptionalLong; import java.util.concurrent.CompletableFuture; import java.util.function.Supplier; import java.util.stream.IntStream; @@ -360,7 +361,8 @@ private FileWriter createParquetFileWriter(Location path, List fileSystem.deleteFile(path); @@ -384,7 +386,6 @@ private FileWriter createParquetFileWriter(Location path, List deltaTableProperties(ConnectorSession session, String location, boolean external) { ImmutableMap.Builder properties = ImmutableMap.builder() - .put(PRESTO_QUERY_ID_NAME, session.getQueryId()) + .put(TRINO_QUERY_ID_NAME, session.getQueryId()) .put(LOCATION_PROPERTY, location) .put(TABLE_PROVIDER_PROPERTY, TABLE_PROVIDER_VALUE) // Set bogus table stats to prevent Hive 3.x from gathering these stats at table creation. @@ -3374,11 +3374,11 @@ private static DeltaLakeColumnHandle toColumnHandle(String originalName, Type ty private static Optional getQueryId(Database database) { - return Optional.ofNullable(database.getParameters().get(PRESTO_QUERY_ID_NAME)); + return Optional.ofNullable(database.getParameters().get(TRINO_QUERY_ID_NAME)); } private static Optional getQueryId(Table table) { - return Optional.ofNullable(table.getParameters().get(PRESTO_QUERY_ID_NAME)); + return Optional.ofNullable(table.getParameters().get(TRINO_QUERY_ID_NAME)); } } diff --git a/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/DeltaLakePageSourceProvider.java b/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/DeltaLakePageSourceProvider.java index 82a0d3bbdbfa..d61d2cf4d8ed 100644 --- a/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/DeltaLakePageSourceProvider.java +++ b/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/DeltaLakePageSourceProvider.java @@ -23,6 +23,8 @@ import io.trino.filesystem.TrinoInputFile; import io.trino.parquet.ParquetDataSource; import io.trino.parquet.ParquetReaderOptions; +import io.trino.parquet.metadata.FileMetadata; +import io.trino.parquet.metadata.ParquetMetadata; import io.trino.parquet.reader.MetadataReader; import io.trino.plugin.deltalake.transactionlog.DeltaLakeSchemaSupport.ColumnMappingMode; import io.trino.plugin.hive.FileFormatDataSourceStats; @@ -51,8 +53,6 @@ import io.trino.spi.predicate.TupleDomain; import io.trino.spi.type.StandardTypes; import io.trino.spi.type.TypeManager; -import org.apache.parquet.hadoop.metadata.FileMetaData; -import org.apache.parquet.hadoop.metadata.ParquetMetadata; import org.apache.parquet.schema.MessageType; import org.apache.parquet.schema.Type; import org.joda.time.DateTimeZone; @@ -63,6 +63,7 @@ import java.util.List; import java.util.Map; import java.util.Optional; +import java.util.OptionalLong; import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.collect.ImmutableList.toImmutableList; @@ -74,8 +75,6 @@ import static io.trino.plugin.deltalake.DeltaLakeColumnType.REGULAR; import static io.trino.plugin.deltalake.DeltaLakeSessionProperties.getParquetMaxReadBlockRowCount; import static io.trino.plugin.deltalake.DeltaLakeSessionProperties.getParquetMaxReadBlockSize; -import static io.trino.plugin.deltalake.DeltaLakeSessionProperties.isParquetOptimizedNestedReaderEnabled; -import static io.trino.plugin.deltalake.DeltaLakeSessionProperties.isParquetOptimizedReaderEnabled; import static io.trino.plugin.deltalake.DeltaLakeSessionProperties.isParquetUseColumnIndex; import static io.trino.plugin.deltalake.transactionlog.DeltaLakeSchemaSupport.extractSchema; import static io.trino.plugin.deltalake.transactionlog.DeltaLakeSchemaSupport.getColumnMappingMode; @@ -183,9 +182,7 @@ public ConnectorPageSource createPageSource( TrinoInputFile inputFile = fileSystemFactory.create(session).newInputFile(location, split.getFileSize()); ParquetReaderOptions options = parquetReaderOptions.withMaxReadBlockSize(getParquetMaxReadBlockSize(session)) .withMaxReadBlockRowCount(getParquetMaxReadBlockRowCount(session)) - .withUseColumnIndex(isParquetUseColumnIndex(session)) - .withBatchColumnReaders(isParquetOptimizedReaderEnabled(session)) - .withBatchNestedColumnReaders(isParquetOptimizedNestedReaderEnabled(session)); + .withUseColumnIndex(isParquetUseColumnIndex(session)); ColumnMappingMode columnMappingMode = getColumnMappingMode(table.getMetadataEntry()); Map parquetFieldIdToName = columnMappingMode == ColumnMappingMode.ID ? loadParquetIdAndNameMapping(inputFile, options) : ImmutableMap.of(); @@ -209,13 +206,14 @@ public ConnectorPageSource createPageSource( split.getStart(), split.getLength(), hiveColumnHandles.build(), - parquetPredicate, + List.of(parquetPredicate), true, parquetDateTimeZone, fileFormatDataSourceStats, options, Optional.empty(), - domainCompactionThreshold); + domainCompactionThreshold, + OptionalLong.empty()); Optional projectionsAdapter = pageSource.getReaderColumns().map(readerColumns -> new ReaderProjectionsAdapter( @@ -240,7 +238,7 @@ public Map loadParquetIdAndNameMapping(TrinoInputFile inputFile { try (ParquetDataSource dataSource = new TrinoParquetDataSource(inputFile, options, fileFormatDataSourceStats)) { ParquetMetadata parquetMetadata = MetadataReader.readFooter(dataSource, Optional.empty()); - FileMetaData fileMetaData = parquetMetadata.getFileMetaData(); + FileMetadata fileMetaData = parquetMetadata.getFileMetaData(); MessageType fileSchema = fileMetaData.getSchema(); return fileSchema.getFields().stream() diff --git a/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/DeltaLakeSessionProperties.java b/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/DeltaLakeSessionProperties.java index d1ecc0e2cbcc..b6142b06dd39 100644 --- a/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/DeltaLakeSessionProperties.java +++ b/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/DeltaLakeSessionProperties.java @@ -54,8 +54,6 @@ public final class DeltaLakeSessionProperties private static final String PARQUET_MAX_READ_BLOCK_SIZE = "parquet_max_read_block_size"; private static final String PARQUET_MAX_READ_BLOCK_ROW_COUNT = "parquet_max_read_block_row_count"; private static final String PARQUET_USE_COLUMN_INDEX = "parquet_use_column_index"; - private static final String PARQUET_OPTIMIZED_READER_ENABLED = "parquet_optimized_reader_enabled"; - private static final String PARQUET_OPTIMIZED_NESTED_READER_ENABLED = "parquet_optimized_nested_reader_enabled"; private static final String PARQUET_WRITER_BLOCK_SIZE = "parquet_writer_block_size"; private static final String PARQUET_WRITER_PAGE_SIZE = "parquet_writer_page_size"; private static final String TARGET_MAX_FILE_SIZE = "target_max_file_size"; @@ -124,16 +122,6 @@ public DeltaLakeSessionProperties( "Use Parquet column index", parquetReaderConfig.isUseColumnIndex(), false), - booleanProperty( - PARQUET_OPTIMIZED_READER_ENABLED, - "Use optimized Parquet reader", - parquetReaderConfig.isOptimizedReaderEnabled(), - false), - booleanProperty( - PARQUET_OPTIMIZED_NESTED_READER_ENABLED, - "Use optimized Parquet reader for nested columns", - parquetReaderConfig.isOptimizedNestedReaderEnabled(), - false), dataSizeProperty( PARQUET_WRITER_BLOCK_SIZE, "Parquet: Writer block size", @@ -250,16 +238,6 @@ public static boolean isParquetUseColumnIndex(ConnectorSession session) return session.getProperty(PARQUET_USE_COLUMN_INDEX, Boolean.class); } - public static boolean isParquetOptimizedReaderEnabled(ConnectorSession session) - { - return session.getProperty(PARQUET_OPTIMIZED_READER_ENABLED, Boolean.class); - } - - public static boolean isParquetOptimizedNestedReaderEnabled(ConnectorSession session) - { - return session.getProperty(PARQUET_OPTIMIZED_NESTED_READER_ENABLED, Boolean.class); - } - public static DataSize getParquetWriterBlockSize(ConnectorSession session) { return session.getProperty(PARQUET_WRITER_BLOCK_SIZE, DataSize.class); diff --git a/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/DeltaLakeWriter.java b/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/DeltaLakeWriter.java index 8a6a9366a40e..4616be663ebf 100644 --- a/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/DeltaLakeWriter.java +++ b/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/DeltaLakeWriter.java @@ -21,6 +21,9 @@ import io.trino.filesystem.TrinoFileSystem; import io.trino.filesystem.TrinoInputFile; import io.trino.parquet.ParquetReaderOptions; +import io.trino.parquet.metadata.BlockMetadata; +import io.trino.parquet.metadata.ColumnChunkMetadata; +import io.trino.parquet.metadata.ParquetMetadata; import io.trino.parquet.reader.MetadataReader; import io.trino.plugin.deltalake.DataFileInfo.DataFileType; import io.trino.plugin.deltalake.transactionlog.statistics.DeltaLakeJsonFileStatistics; @@ -43,9 +46,6 @@ import io.trino.spi.type.TimestampWithTimeZoneType; import io.trino.spi.type.Type; import org.apache.parquet.column.statistics.Statistics; -import org.apache.parquet.hadoop.metadata.BlockMetaData; -import org.apache.parquet.hadoop.metadata.ColumnChunkMetaData; -import org.apache.parquet.hadoop.metadata.ParquetMetadata; import java.io.Closeable; import java.io.IOException; @@ -208,9 +208,9 @@ private static DeltaLakeJsonFileStatistics readStatistics(TrinoInputFile inputFi new FileFormatDataSourceStats())) { ParquetMetadata parquetMetadata = MetadataReader.readFooter(trinoParquetDataSource, Optional.empty()); - ImmutableMultimap.Builder metadataForColumn = ImmutableMultimap.builder(); - for (BlockMetaData blockMetaData : parquetMetadata.getBlocks()) { - for (ColumnChunkMetaData columnChunkMetaData : blockMetaData.getColumns()) { + ImmutableMultimap.Builder metadataForColumn = ImmutableMultimap.builder(); + for (BlockMetadata blockMetaData : parquetMetadata.getBlocks()) { + for (ColumnChunkMetadata columnChunkMetaData : blockMetaData.columns()) { if (columnChunkMetaData.getPath().size() != 1) { continue; // Only base column stats are supported } @@ -224,7 +224,7 @@ private static DeltaLakeJsonFileStatistics readStatistics(TrinoInputFile inputFi } @VisibleForTesting - static DeltaLakeJsonFileStatistics mergeStats(Multimap metadataForColumn, Map typeForColumn, long rowCount) + static DeltaLakeJsonFileStatistics mergeStats(Multimap metadataForColumn, Map typeForColumn, long rowCount) { Map>> statsForColumn = metadataForColumn.keySet().stream() .collect(toImmutableMap(identity(), key -> mergeMetadataList(metadataForColumn.get(key)))); @@ -240,14 +240,14 @@ static DeltaLakeJsonFileStatistics mergeStats(Multimap> mergeMetadataList(Collection metadataList) + private static Optional> mergeMetadataList(Collection metadataList) { if (hasInvalidStatistics(metadataList)) { return Optional.empty(); } return metadataList.stream() - .>map(ColumnChunkMetaData::getStatistics) + .>map(ColumnChunkMetadata::getStatistics) .reduce((statsA, statsB) -> { statsA.mergeStatistics(statsB); return statsA; diff --git a/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/InternalDeltaLakeConnectorFactory.java b/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/InternalDeltaLakeConnectorFactory.java index ff8d74d90c00..7b9c2d78c358 100644 --- a/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/InternalDeltaLakeConnectorFactory.java +++ b/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/InternalDeltaLakeConnectorFactory.java @@ -101,7 +101,7 @@ public static Connector createConnector( new DeltaLakeSecurityModule(), fileSystemFactory .map(factory -> (Module) binder -> binder.bind(TrinoFileSystemFactory.class).toInstance(factory)) - .orElseGet(FileSystemModule::new), + .orElseGet(() -> new FileSystemModule(catalogName, context.getNodeManager(), context.getOpenTelemetry(), false)), binder -> { binder.bind(OpenTelemetry.class).toInstance(context.getOpenTelemetry()); binder.bind(Tracer.class).toInstance(context.getTracer()); diff --git a/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/functions/tablechanges/TableChangesFunctionProcessor.java b/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/functions/tablechanges/TableChangesFunctionProcessor.java index de65c3bd5459..40553fa67aa5 100644 --- a/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/functions/tablechanges/TableChangesFunctionProcessor.java +++ b/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/functions/tablechanges/TableChangesFunctionProcessor.java @@ -38,6 +38,7 @@ import java.util.Map; import java.util.Optional; import java.util.OptionalInt; +import java.util.OptionalLong; import static com.google.common.base.Verify.verify; import static com.google.common.collect.ImmutableList.toImmutableList; @@ -46,8 +47,6 @@ import static io.trino.plugin.deltalake.DeltaLakeColumnType.REGULAR; import static io.trino.plugin.deltalake.DeltaLakeSessionProperties.getParquetMaxReadBlockRowCount; import static io.trino.plugin.deltalake.DeltaLakeSessionProperties.getParquetMaxReadBlockSize; -import static io.trino.plugin.deltalake.DeltaLakeSessionProperties.isParquetOptimizedNestedReaderEnabled; -import static io.trino.plugin.deltalake.DeltaLakeSessionProperties.isParquetOptimizedReaderEnabled; import static io.trino.plugin.deltalake.DeltaLakeSessionProperties.isParquetUseColumnIndex; import static io.trino.plugin.deltalake.functions.tablechanges.TableChangesFileType.CDF_FILE; import static io.trino.spi.function.table.TableFunctionProcessorState.Finished.FINISHED; @@ -176,9 +175,7 @@ private static DeltaLakePageSource createDeltaLakePageSource( parquetReaderOptions = parquetReaderOptions .withMaxReadBlockSize(getParquetMaxReadBlockSize(session)) .withMaxReadBlockRowCount(getParquetMaxReadBlockRowCount(session)) - .withUseColumnIndex(isParquetUseColumnIndex(session)) - .withBatchColumnReaders(isParquetOptimizedReaderEnabled(session)) - .withBatchNestedColumnReaders(isParquetOptimizedNestedReaderEnabled(session)); + .withUseColumnIndex(isParquetUseColumnIndex(session)); List splitColumns = switch (split.fileType()) { case CDF_FILE -> ImmutableList.builder().addAll(handle.columns()) @@ -199,13 +196,14 @@ private static DeltaLakePageSource createDeltaLakePageSource( 0, split.fileSize(), splitColumns.stream().filter(column -> column.getColumnType() == REGULAR).map(DeltaLakeColumnHandle::toHiveColumnHandle).collect(toImmutableList()), - TupleDomain.all(), // TODO add predicate pushdown https://github.com/trinodb/trino/issues/16990 + List.of(TupleDomain.all()), // TODO add predicate pushdown https://github.com/trinodb/trino/issues/16990 true, parquetDateTimeZone, fileFormatDataSourceStats, parquetReaderOptions, Optional.empty(), - domainCompactionThreshold); + domainCompactionThreshold, + OptionalLong.empty()); verify(pageSource.getReaderColumns().isEmpty(), "Unexpected reader columns: %s", pageSource.getReaderColumns().orElse(null)); diff --git a/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/metastore/HiveMetastoreBackedDeltaLakeMetastore.java b/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/metastore/HiveMetastoreBackedDeltaLakeMetastore.java index a8aff2285a8e..74be94e09385 100644 --- a/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/metastore/HiveMetastoreBackedDeltaLakeMetastore.java +++ b/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/metastore/HiveMetastoreBackedDeltaLakeMetastore.java @@ -63,7 +63,10 @@ public List getAllTables(String databaseName) // it would be nice to filter out non-Delta tables; however, we can not call // metastore.getTablesWithParameter(schema, TABLE_PROVIDER_PROP, TABLE_PROVIDER_VALUE), because that property // contains a dot and must be compared case-insensitive - return delegate.getAllTables(databaseName); + return delegate.getTables(databaseName) + .stream() + .map(table -> table.tableName().toString()) + .toList(); } @Override diff --git a/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/metastore/glue/DeltaLakeGlueMetastoreModule.java b/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/metastore/glue/DeltaLakeGlueMetastoreModule.java index fa400129eea6..0275115d9994 100644 --- a/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/metastore/glue/DeltaLakeGlueMetastoreModule.java +++ b/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/metastore/glue/DeltaLakeGlueMetastoreModule.java @@ -13,12 +13,12 @@ */ package io.trino.plugin.deltalake.metastore.glue; -import com.amazonaws.services.glue.model.Table; import com.google.inject.Binder; import com.google.inject.Key; import com.google.inject.TypeLiteral; import io.airlift.configuration.AbstractConfigurationAwareModule; import io.trino.plugin.deltalake.AllowDeltaLakeManagedTableRename; +import io.trino.plugin.hive.metastore.Table; import io.trino.plugin.hive.metastore.glue.ForGlueHiveMetastore; import io.trino.plugin.hive.metastore.glue.GlueMetastoreModule; diff --git a/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/metastore/glue/DeltaLakeGlueMetastoreTableFilterProvider.java b/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/metastore/glue/DeltaLakeGlueMetastoreTableFilterProvider.java index fd3e8dbd6472..78e48e388724 100644 --- a/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/metastore/glue/DeltaLakeGlueMetastoreTableFilterProvider.java +++ b/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/metastore/glue/DeltaLakeGlueMetastoreTableFilterProvider.java @@ -13,15 +13,13 @@ */ package io.trino.plugin.deltalake.metastore.glue; -import com.amazonaws.services.glue.model.Table; import com.google.inject.Inject; import com.google.inject.Provider; +import io.trino.plugin.hive.metastore.Table; import io.trino.plugin.hive.util.HiveUtil; import java.util.function.Predicate; -import static io.trino.plugin.hive.metastore.glue.converter.GlueToTrinoConverter.getTableParameters; - public class DeltaLakeGlueMetastoreTableFilterProvider implements Provider> { @@ -37,7 +35,7 @@ public DeltaLakeGlueMetastoreTableFilterProvider(DeltaLakeGlueMetastoreConfig co public Predicate get() { if (hideNonDeltaLakeTables) { - return table -> HiveUtil.isDeltaLakeTable(getTableParameters(table)); + return table -> HiveUtil.isDeltaLakeTable(table); } return table -> true; } diff --git a/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/transactionlog/DeltaLakeParquetStatisticsUtils.java b/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/transactionlog/DeltaLakeParquetStatisticsUtils.java index 573c99f32ff7..b9f62ad9ae20 100644 --- a/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/transactionlog/DeltaLakeParquetStatisticsUtils.java +++ b/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/transactionlog/DeltaLakeParquetStatisticsUtils.java @@ -16,6 +16,7 @@ import com.google.common.collect.ImmutableMap; import io.airlift.log.Logger; import io.airlift.slice.Slice; +import io.trino.parquet.metadata.ColumnChunkMetadata; import io.trino.plugin.base.type.DecodedTimestamp; import io.trino.spi.block.Block; import io.trino.spi.block.ColumnarRow; @@ -36,7 +37,6 @@ import org.apache.parquet.column.statistics.IntStatistics; import org.apache.parquet.column.statistics.LongStatistics; import org.apache.parquet.column.statistics.Statistics; -import org.apache.parquet.hadoop.metadata.ColumnChunkMetaData; import org.apache.parquet.schema.LogicalTypeAnnotation; import java.math.BigDecimal; @@ -87,7 +87,7 @@ public class DeltaLakeParquetStatisticsUtils private DeltaLakeParquetStatisticsUtils() {} - public static boolean hasInvalidStatistics(Collection metadataList) + public static boolean hasInvalidStatistics(Collection metadataList) { return metadataList.stream() .anyMatch(metadata -> diff --git a/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/transactionlog/checkpoint/CheckpointEntryIterator.java b/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/transactionlog/checkpoint/CheckpointEntryIterator.java index a8904fb2b6f0..0f684d771436 100644 --- a/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/transactionlog/checkpoint/CheckpointEntryIterator.java +++ b/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/transactionlog/checkpoint/CheckpointEntryIterator.java @@ -186,13 +186,14 @@ public CheckpointEntryIterator( 0, fileSize, columns, - tupleDomain, + List.of(tupleDomain), true, DateTimeZone.UTC, stats, parquetReaderOptions, Optional.empty(), - domainCompactionThreshold); + domainCompactionThreshold, + OptionalLong.empty()); verify(pageSource.getReaderColumns().isEmpty(), "All columns expected to be base columns"); diff --git a/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/transactionlog/checkpoint/CheckpointWriter.java b/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/transactionlog/checkpoint/CheckpointWriter.java index 99c8a92b0eaf..2e6a5c819197 100644 --- a/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/transactionlog/checkpoint/CheckpointWriter.java +++ b/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/transactionlog/checkpoint/CheckpointWriter.java @@ -138,7 +138,6 @@ public void write(CheckpointEntries entries, TrinoOutputFile outputFile) parquetWriterOptions, CompressionCodec.SNAPPY, trinoVersion, - false, Optional.of(DateTimeZone.UTC), Optional.empty()); diff --git a/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/BaseDeltaLakeConnectorSmokeTest.java b/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/BaseDeltaLakeConnectorSmokeTest.java index d6178f20ab3f..abf2f3f4f96b 100644 --- a/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/BaseDeltaLakeConnectorSmokeTest.java +++ b/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/BaseDeltaLakeConnectorSmokeTest.java @@ -154,7 +154,7 @@ protected QueryRunner createQueryRunner() this.metastore = new BridgingHiveMetastore( testingThriftHiveMetastoreBuilder() .metastoreClient(hiveMinioDataLake.getHiveHadoop().getHiveMetastoreEndpoint()) - .build()); + .build(this::closeAfterClass)); DistributedQueryRunner queryRunner = createDeltaLakeQueryRunner(); try { @@ -188,7 +188,7 @@ protected QueryRunner createQueryRunner() registerTableFromResources(table, resourcePath, queryRunner); }); - queryRunner.installPlugin(new TestingHivePlugin()); + queryRunner.installPlugin(new TestingHivePlugin(queryRunner.getCoordinator().getBaseDataDir().resolve("hive_data"))); queryRunner.createCatalog( "hive", diff --git a/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/BaseDeltaLakeSharedMetastoreViewsTest.java b/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/BaseDeltaLakeSharedMetastoreViewsTest.java index 0735f42b9581..c1da333464fc 100644 --- a/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/BaseDeltaLakeSharedMetastoreViewsTest.java +++ b/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/BaseDeltaLakeSharedMetastoreViewsTest.java @@ -64,7 +64,7 @@ protected QueryRunner createQueryRunner() queryRunner.installPlugin(new TestingDeltaLakePlugin(Optional.of(new TestingDeltaLakeMetastoreModule(metastore)), Optional.empty(), EMPTY_MODULE)); queryRunner.createCatalog(DELTA_CATALOG_NAME, "delta_lake"); - queryRunner.installPlugin(new TestingHivePlugin(metastore)); + queryRunner.installPlugin(new TestingHivePlugin(dataDirectory, metastore)); ImmutableMap hiveProperties = ImmutableMap.builder() .put("hive.allow-drop-table", "true") diff --git a/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/TestDeltaLakeBasic.java b/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/TestDeltaLakeBasic.java index 450d6d1b42e3..ab710894c9c3 100644 --- a/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/TestDeltaLakeBasic.java +++ b/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/TestDeltaLakeBasic.java @@ -27,6 +27,8 @@ import io.trino.filesystem.hdfs.HdfsFileSystemFactory; import io.trino.filesystem.local.LocalInputFile; import io.trino.parquet.ParquetReaderOptions; +import io.trino.parquet.metadata.FileMetadata; +import io.trino.parquet.metadata.ParquetMetadata; import io.trino.parquet.reader.MetadataReader; import io.trino.plugin.deltalake.transactionlog.AddFileEntry; import io.trino.plugin.deltalake.transactionlog.DeltaLakeTransactionLogEntry; @@ -36,8 +38,6 @@ import io.trino.plugin.hive.parquet.TrinoParquetDataSource; import io.trino.testing.AbstractTestQueryFramework; import io.trino.testing.QueryRunner; -import org.apache.parquet.hadoop.metadata.FileMetaData; -import org.apache.parquet.hadoop.metadata.ParquetMetadata; import org.apache.parquet.schema.PrimitiveType; import org.testng.annotations.BeforeClass; import org.testng.annotations.DataProvider; @@ -280,7 +280,7 @@ public void testOptimizeWithColumnMappingMode(String columnMappingMode) ParquetMetadata parquetMetadata = MetadataReader.readFooter( new TrinoParquetDataSource(inputFile, new ParquetReaderOptions(), new FileFormatDataSourceStats()), Optional.empty()); - FileMetaData fileMetaData = parquetMetadata.getFileMetaData(); + FileMetadata fileMetaData = parquetMetadata.getFileMetaData(); PrimitiveType physicalType = getOnlyElement(fileMetaData.getSchema().getColumns().iterator()).getPrimitiveType(); assertThat(physicalType.getName()).isEqualTo(physicalName); if (columnMappingMode.equals("id")) { diff --git a/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/TestDeltaLakeCreateSchemaInternalRetry.java b/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/TestDeltaLakeCreateSchemaInternalRetry.java index bcb14e7853b5..6d6bf8d1b315 100644 --- a/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/TestDeltaLakeCreateSchemaInternalRetry.java +++ b/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/TestDeltaLakeCreateSchemaInternalRetry.java @@ -15,6 +15,7 @@ import com.google.common.collect.ImmutableMap; import io.trino.Session; +import io.trino.filesystem.local.LocalFileSystemFactory; import io.trino.plugin.deltalake.metastore.TestingDeltaLakeMetastoreModule; import io.trino.plugin.hive.NodeVersion; import io.trino.plugin.hive.SchemaAlreadyExistsException; @@ -38,8 +39,7 @@ import static com.google.common.io.RecursiveDeleteOption.ALLOW_INSECURE; import static com.google.inject.util.Modules.EMPTY_MODULE; import static io.trino.plugin.deltalake.DeltaLakeConnectorFactory.CONNECTOR_NAME; -import static io.trino.plugin.hive.HiveMetadata.PRESTO_QUERY_ID_NAME; -import static io.trino.plugin.hive.HiveTestUtils.HDFS_ENVIRONMENT; +import static io.trino.plugin.hive.HiveMetadata.TRINO_QUERY_ID_NAME; import static io.trino.testing.TestingNames.randomNameSuffix; import static io.trino.testing.TestingSession.testSessionBuilder; @@ -63,9 +63,10 @@ protected QueryRunner createQueryRunner() DistributedQueryRunner queryRunner = DistributedQueryRunner.builder(session).build(); this.dataDirectory = queryRunner.getCoordinator().getBaseDataDir().resolve("delta_lake_data").toString(); + LocalFileSystemFactory fileSystemFactory = new LocalFileSystemFactory(Path.of(dataDirectory)); this.metastore = new FileHiveMetastore( new NodeVersion("testversion"), - HDFS_ENVIRONMENT, + fileSystemFactory, new HiveMetastoreConfig().isHideDeltaLakeTables(), new FileHiveMetastoreConfig() .setCatalogDirectory(dataDirectory) @@ -77,7 +78,7 @@ public synchronized void createDatabase(Database database) if (database.getDatabaseName().equals(TEST_SCHEMA_DIFFERENT_SESSION)) { // By modifying query id test simulates that schema was created from different session. database = Database.builder(database) - .setParameters(ImmutableMap.of(PRESTO_QUERY_ID_NAME, "new_query_id")) + .setParameters(ImmutableMap.of(TRINO_QUERY_ID_NAME, "new_query_id")) .build(); } // Simulate retry mechanism with timeout failure. diff --git a/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/TestDeltaLakeFlushMetadataCacheProcedure.java b/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/TestDeltaLakeFlushMetadataCacheProcedure.java index e53821e17596..15cb794dd264 100644 --- a/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/TestDeltaLakeFlushMetadataCacheProcedure.java +++ b/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/TestDeltaLakeFlushMetadataCacheProcedure.java @@ -44,7 +44,7 @@ protected QueryRunner createQueryRunner() metastore = new BridgingHiveMetastore( testingThriftHiveMetastoreBuilder() .metastoreClient(hiveMinioDataLake.getHiveHadoop().getHiveMetastoreEndpoint()) - .build()); + .build(this::closeAfterClass)); return createS3DeltaLakeQueryRunner( DELTA_CATALOG, diff --git a/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/TestDeltaLakePerTransactionMetastoreCache.java b/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/TestDeltaLakePerTransactionMetastoreCache.java index 9a0ab509d217..11d540af4e28 100644 --- a/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/TestDeltaLakePerTransactionMetastoreCache.java +++ b/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/TestDeltaLakePerTransactionMetastoreCache.java @@ -39,7 +39,7 @@ import java.util.Optional; import static io.trino.plugin.deltalake.DeltaLakeQueryRunner.DELTA_CATALOG; -import static io.trino.plugin.hive.metastore.CountingAccessHiveMetastore.Method.GET_TABLE; +import static io.trino.plugin.hive.metastore.MetastoreMethod.GET_TABLE; import static io.trino.plugin.hive.metastore.file.TestingFileHiveMetastore.createTestingFileHiveMetastore; import static io.trino.testing.TestingNames.randomNameSuffix; import static io.trino.testing.TestingSession.testSessionBuilder; diff --git a/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/TestDeltaLakeSharedFileMetastoreWithTableRedirections.java b/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/TestDeltaLakeSharedFileMetastoreWithTableRedirections.java index 683837239094..fbd0a5cc3952 100644 --- a/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/TestDeltaLakeSharedFileMetastoreWithTableRedirections.java +++ b/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/TestDeltaLakeSharedFileMetastoreWithTableRedirections.java @@ -55,7 +55,7 @@ protected QueryRunner createQueryRunner() queryRunner.createCatalog("delta_with_redirections", CONNECTOR_NAME, deltaLakeProperties); queryRunner.execute("CREATE SCHEMA " + schema); - queryRunner.installPlugin(new TestingHivePlugin()); + queryRunner.installPlugin(new TestingHivePlugin(dataDirectory)); queryRunner.createCatalog( "hive_with_redirections", diff --git a/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/TestDeltaLakeSharedGlueMetastoreViews.java b/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/TestDeltaLakeSharedGlueMetastoreViews.java index c9d2b942e75c..1dca66a12ed6 100644 --- a/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/TestDeltaLakeSharedGlueMetastoreViews.java +++ b/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/TestDeltaLakeSharedGlueMetastoreViews.java @@ -17,7 +17,7 @@ import java.nio.file.Path; -import static io.trino.plugin.hive.metastore.glue.GlueHiveMetastore.createTestingGlueHiveMetastore; +import static io.trino.plugin.hive.metastore.glue.TestingGlueHiveMetastore.createTestingGlueHiveMetastore; /** * Requires AWS credentials, which can be provided any way supported by the DefaultProviderChain @@ -29,6 +29,6 @@ public class TestDeltaLakeSharedGlueMetastoreViews @Override protected HiveMetastore createTestMetastore(Path dataDirectory) { - return createTestingGlueHiveMetastore(dataDirectory); + return createTestingGlueHiveMetastore(dataDirectory, this::closeAfterClass); } } diff --git a/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/TestDeltaLakeSharedGlueMetastoreWithTableRedirections.java b/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/TestDeltaLakeSharedGlueMetastoreWithTableRedirections.java index 8febb700f6ea..12d8b6657ac9 100644 --- a/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/TestDeltaLakeSharedGlueMetastoreWithTableRedirections.java +++ b/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/TestDeltaLakeSharedGlueMetastoreWithTableRedirections.java @@ -24,7 +24,7 @@ import java.nio.file.Path; -import static io.trino.plugin.hive.metastore.glue.GlueHiveMetastore.createTestingGlueHiveMetastore; +import static io.trino.plugin.hive.metastore.glue.TestingGlueHiveMetastore.createTestingGlueHiveMetastore; import static io.trino.testing.TestingSession.testSessionBuilder; import static java.lang.String.format; @@ -66,8 +66,8 @@ protected QueryRunner createQueryRunner() .put("delta.hive-catalog-name", "hive_with_redirections") .buildOrThrow()); - this.glueMetastore = createTestingGlueHiveMetastore(dataDirectory); - queryRunner.installPlugin(new TestingHivePlugin(glueMetastore)); + this.glueMetastore = createTestingGlueHiveMetastore(dataDirectory, this::closeAfterClass); + queryRunner.installPlugin(new TestingHivePlugin(queryRunner.getCoordinator().getBaseDataDir().resolve("hive_data"), glueMetastore)); queryRunner.createCatalog( "hive_with_redirections", "hive", diff --git a/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/TestDeltaLakeSharedHiveMetastoreWithViews.java b/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/TestDeltaLakeSharedHiveMetastoreWithViews.java index 666a00909116..46445b87923b 100644 --- a/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/TestDeltaLakeSharedHiveMetastoreWithViews.java +++ b/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/TestDeltaLakeSharedHiveMetastoreWithViews.java @@ -53,7 +53,7 @@ protected QueryRunner createQueryRunner() hiveMinioDataLake.getHiveHadoop()); queryRunner.execute("CREATE SCHEMA " + schema + " WITH (location = 's3://" + bucketName + "/" + schema + "')"); - queryRunner.installPlugin(new TestingHivePlugin()); + queryRunner.installPlugin(new TestingHivePlugin(queryRunner.getCoordinator().getBaseDataDir().resolve("hive_data"))); Map s3Properties = ImmutableMap.builder() .put("hive.s3.aws-access-key", MINIO_ACCESS_KEY) .put("hive.s3.aws-secret-key", MINIO_SECRET_KEY) diff --git a/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/TestDeltaLakeTableWithCustomLocationUsingGlueMetastore.java b/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/TestDeltaLakeTableWithCustomLocationUsingGlueMetastore.java index d9c2128884f3..2cb845027c61 100644 --- a/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/TestDeltaLakeTableWithCustomLocationUsingGlueMetastore.java +++ b/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/TestDeltaLakeTableWithCustomLocationUsingGlueMetastore.java @@ -25,7 +25,7 @@ import static com.google.common.io.MoreFiles.deleteRecursively; import static com.google.common.io.RecursiveDeleteOption.ALLOW_INSECURE; import static io.trino.plugin.deltalake.DeltaLakeConnectorFactory.CONNECTOR_NAME; -import static io.trino.plugin.hive.metastore.glue.GlueHiveMetastore.createTestingGlueHiveMetastore; +import static io.trino.plugin.hive.metastore.glue.TestingGlueHiveMetastore.createTestingGlueHiveMetastore; import static io.trino.testing.TestingSession.testSessionBuilder; public class TestDeltaLakeTableWithCustomLocationUsingGlueMetastore @@ -57,7 +57,7 @@ protected QueryRunner createQueryRunner() .put("hive.metastore.glue.default-warehouse-dir", metastoreDir.toURI().toString()) .buildOrThrow()); - metastore = createTestingGlueHiveMetastore(metastoreDir.toPath()); + metastore = createTestingGlueHiveMetastore(metastoreDir.toPath(), this::closeAfterClass); queryRunner.execute("CREATE SCHEMA " + SCHEMA + " WITH (location = '" + metastoreDir.toURI() + "')"); return queryRunner; diff --git a/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/TestDeltaLakeWriter.java b/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/TestDeltaLakeWriter.java index c2a075507932..67cdd3838707 100644 --- a/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/TestDeltaLakeWriter.java +++ b/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/TestDeltaLakeWriter.java @@ -18,11 +18,11 @@ import com.google.common.collect.ImmutableMultimap; import com.google.common.collect.ImmutableSet; import com.google.common.collect.Multimap; +import io.trino.parquet.metadata.ColumnChunkMetadata; import io.trino.plugin.deltalake.transactionlog.statistics.DeltaLakeFileStatistics; import io.trino.spi.type.VarcharType; import org.apache.parquet.column.EncodingStats; import org.apache.parquet.column.statistics.Statistics; -import org.apache.parquet.hadoop.metadata.ColumnChunkMetaData; import org.apache.parquet.hadoop.metadata.ColumnPath; import org.apache.parquet.hadoop.metadata.CompressionCodecName; import org.apache.parquet.schema.PrimitiveType; @@ -53,7 +53,7 @@ public void testMergeIntStatistics() { String columnName = "t_int"; PrimitiveType intType = new PrimitiveType(Type.Repetition.REQUIRED, PrimitiveType.PrimitiveTypeName.INT32, columnName); - List metadata = ImmutableList.of( + List metadata = ImmutableList.of( createMetaData(columnName, intType, 10, Statistics.getBuilderForReading(intType).withMin(getIntByteArray(-100)).withMax(getIntByteArray(250)).withNumNulls(6).build()), createMetaData(columnName, intType, 10, @@ -72,7 +72,7 @@ public void testMergeFloatStatistics() { String columnName = "t_float"; PrimitiveType type = new PrimitiveType(Type.Repetition.REQUIRED, PrimitiveType.PrimitiveTypeName.FLOAT, columnName); - List metadata = ImmutableList.of( + List metadata = ImmutableList.of( createMetaData(columnName, type, 10, Statistics.getBuilderForReading(type).withMin(getFloatByteArray(0.01f)).withMax(getFloatByteArray(1.0f)).withNumNulls(6).build()), createMetaData(columnName, type, 10, @@ -91,7 +91,7 @@ public void testMergeFloatNaNStatistics() { String columnName = "t_float"; PrimitiveType type = new PrimitiveType(Type.Repetition.REQUIRED, PrimitiveType.PrimitiveTypeName.FLOAT, columnName); - List metadata = ImmutableList.of( + List metadata = ImmutableList.of( createMetaData(columnName, type, 10, Statistics.getBuilderForReading(type).withMin(getFloatByteArray(0.01f)).withMax(getFloatByteArray(1.0f)).withNumNulls(6).build()), createMetaData(columnName, type, 10, @@ -112,7 +112,7 @@ public void testMergeDoubleNaNStatistics() { String columnName = "t_double"; PrimitiveType type = new PrimitiveType(Type.Repetition.REQUIRED, PrimitiveType.PrimitiveTypeName.DOUBLE, columnName); - List metadata = ImmutableList.of( + List metadata = ImmutableList.of( createMetaData(columnName, type, 10, Statistics.getBuilderForReading(type).withMin(getDoubleByteArray(0.01f)).withMax(getDoubleByteArray(1.0f)).withNumNulls(6).build()), createMetaData(columnName, type, 10, @@ -133,7 +133,7 @@ public void testMergeStringStatistics() { String columnName = "t_string"; PrimitiveType type = new PrimitiveType(Type.Repetition.REQUIRED, PrimitiveType.PrimitiveTypeName.BINARY, columnName); - List metadata = ImmutableList.of( + List metadata = ImmutableList.of( createMetaData(columnName, type, 10, Statistics.getBuilderForReading(type).withMin("aba".getBytes(UTF_8)).withMax("ab⌘".getBytes(UTF_8)).withNumNulls(6).build()), createMetaData(columnName, type, 10, @@ -152,7 +152,7 @@ public void testMergeStringUnicodeStatistics() { String columnName = "t_string"; PrimitiveType type = new PrimitiveType(Type.Repetition.REQUIRED, PrimitiveType.PrimitiveTypeName.BINARY, columnName); - List metadata = ImmutableList.of( + List metadata = ImmutableList.of( createMetaData(columnName, type, 10, Statistics.getBuilderForReading(type).withMin("aba".getBytes(UTF_8)).withMax("ab\uFAD8".getBytes(UTF_8)).withNumNulls(6).build()), createMetaData(columnName, type, 10, @@ -166,9 +166,9 @@ public void testMergeStringUnicodeStatistics() assertEquals(fileStats.getNullCount(columnName), Optional.of(12L)); } - private ColumnChunkMetaData createMetaData(String columnName, PrimitiveType columnType, long valueCount, Statistics statistics) + private ColumnChunkMetadata createMetaData(String columnName, PrimitiveType columnType, long valueCount, Statistics statistics) { - return ColumnChunkMetaData.get( + return ColumnChunkMetadata.get( ColumnPath.fromDotString(columnName), columnType, CompressionCodecName.SNAPPY, @@ -182,9 +182,9 @@ private ColumnChunkMetaData createMetaData(String columnName, PrimitiveType colu 0); } - private Multimap buildMultimap(String columnName, List metadata) + private Multimap buildMultimap(String columnName, List metadata) { - return ImmutableMultimap.builder() + return ImmutableMultimap.builder() .putAll(columnName, metadata) .build(); } diff --git a/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/metastore/TestDeltaLakeMetastoreAccessOperations.java b/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/metastore/TestDeltaLakeMetastoreAccessOperations.java index fb4ea507a3b5..aff723d72ba4 100644 --- a/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/metastore/TestDeltaLakeMetastoreAccessOperations.java +++ b/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/metastore/TestDeltaLakeMetastoreAccessOperations.java @@ -34,10 +34,10 @@ import java.io.File; import java.util.Optional; -import static io.trino.plugin.hive.metastore.CountingAccessHiveMetastore.Method.CREATE_TABLE; -import static io.trino.plugin.hive.metastore.CountingAccessHiveMetastore.Method.DROP_TABLE; -import static io.trino.plugin.hive.metastore.CountingAccessHiveMetastore.Method.GET_DATABASE; -import static io.trino.plugin.hive.metastore.CountingAccessHiveMetastore.Method.GET_TABLE; +import static io.trino.plugin.hive.metastore.MetastoreMethod.CREATE_TABLE; +import static io.trino.plugin.hive.metastore.MetastoreMethod.DROP_TABLE; +import static io.trino.plugin.hive.metastore.MetastoreMethod.GET_DATABASE; +import static io.trino.plugin.hive.metastore.MetastoreMethod.GET_TABLE; import static io.trino.plugin.hive.metastore.file.TestingFileHiveMetastore.createTestingFileHiveMetastore; import static io.trino.testing.TestingSession.testSessionBuilder; import static java.util.Objects.requireNonNull; diff --git a/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/metastore/glue/TestDeltaLakeCleanUpGlueMetastore.java b/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/metastore/glue/TestDeltaLakeCleanUpGlueMetastore.java index d6716dd463d3..44b4719738ab 100644 --- a/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/metastore/glue/TestDeltaLakeCleanUpGlueMetastore.java +++ b/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/metastore/glue/TestDeltaLakeCleanUpGlueMetastore.java @@ -13,54 +13,43 @@ */ package io.trino.plugin.deltalake.metastore.glue; -import com.amazonaws.services.glue.AWSGlueAsync; -import com.amazonaws.services.glue.AWSGlueAsyncClientBuilder; -import com.amazonaws.services.glue.model.DeleteDatabaseRequest; -import com.amazonaws.services.glue.model.EntityNotFoundException; -import com.amazonaws.services.glue.model.GetDatabasesRequest; -import com.amazonaws.services.glue.model.GetDatabasesResult; import io.airlift.log.Logger; -import io.trino.plugin.hive.aws.AwsApiCallStats; import org.testng.annotations.Test; +import software.amazon.awssdk.services.glue.GlueClient; +import software.amazon.awssdk.services.glue.model.Database; +import software.amazon.awssdk.services.glue.model.EntityNotFoundException; +import software.amazon.awssdk.services.glue.model.GetDatabasesResponse; +import java.time.Duration; +import java.time.Instant; import java.util.List; -import static com.google.common.collect.ImmutableList.toImmutableList; -import static io.trino.plugin.hive.metastore.glue.AwsSdkUtil.getPaginatedResults; -import static java.lang.System.currentTimeMillis; -import static java.util.concurrent.TimeUnit.DAYS; - public class TestDeltaLakeCleanUpGlueMetastore { private static final Logger log = Logger.get(TestDeltaLakeCleanUpGlueMetastore.class); private static final String TEST_DATABASE_NAME_PREFIX = "test_"; + private static final Duration CLEANUP_THRESHOLD = Duration.ofDays(1); @Test public void cleanupOrphanedDatabases() { - AWSGlueAsync glueClient = AWSGlueAsyncClientBuilder.defaultClient(); - long creationTimeMillisThreshold = currentTimeMillis() - DAYS.toMillis(1); - List orphanedDatabases = getPaginatedResults( - glueClient::getDatabases, - new GetDatabasesRequest(), - GetDatabasesRequest::setNextToken, - GetDatabasesResult::getNextToken, - new AwsApiCallStats()) - .map(GetDatabasesResult::getDatabaseList) + GlueClient glueClient = GlueClient.create(); + Instant creationTimeThreshold = Instant.now().minus(CLEANUP_THRESHOLD); + List orphanedDatabases = glueClient.getDatabasesPaginator(x -> {}).stream() + .map(GetDatabasesResponse::databaseList) .flatMap(List::stream) - .filter(glueDatabase -> glueDatabase.getName().startsWith(TEST_DATABASE_NAME_PREFIX) && - glueDatabase.getCreateTime().getTime() <= creationTimeMillisThreshold) - .map(com.amazonaws.services.glue.model.Database::getName) - .collect(toImmutableList()); + .filter(database -> database.name().startsWith(TEST_DATABASE_NAME_PREFIX)) + .filter(database -> database.createTime().isBefore(creationTimeThreshold)) + .map(Database::name) + .toList(); if (!orphanedDatabases.isEmpty()) { log.info("Found %s %s* databases that look orphaned, removing", orphanedDatabases.size(), TEST_DATABASE_NAME_PREFIX); orphanedDatabases.forEach(database -> { try { log.info("Deleting %s database", database); - glueClient.deleteDatabase(new DeleteDatabaseRequest() - .withName(database)); + glueClient.deleteDatabase(builder -> builder.name(database)); } catch (EntityNotFoundException e) { log.info("Database [%s] not found, could be removed by other cleanup process", database); diff --git a/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/metastore/glue/TestDeltaLakeConcurrentModificationGlueMetastore.java b/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/metastore/glue/TestDeltaLakeConcurrentModificationGlueMetastore.java deleted file mode 100644 index c75006287c57..000000000000 --- a/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/metastore/glue/TestDeltaLakeConcurrentModificationGlueMetastore.java +++ /dev/null @@ -1,130 +0,0 @@ -/* - * 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.deltalake.metastore.glue; - -import com.amazonaws.auth.DefaultAWSCredentialsProviderChain; -import com.amazonaws.services.glue.AWSGlueAsync; -import com.amazonaws.services.glue.model.ConcurrentModificationException; -import com.google.common.collect.ImmutableSet; -import io.trino.Session; -import io.trino.plugin.deltalake.TestingDeltaLakePlugin; -import io.trino.plugin.deltalake.metastore.TestingDeltaLakeMetastoreModule; -import io.trino.plugin.hive.metastore.glue.DefaultGlueColumnStatisticsProviderFactory; -import io.trino.plugin.hive.metastore.glue.GlueHiveMetastore; -import io.trino.plugin.hive.metastore.glue.GlueHiveMetastoreConfig; -import io.trino.plugin.hive.metastore.glue.GlueMetastoreStats; -import io.trino.spi.TrinoException; -import io.trino.testing.AbstractTestQueryFramework; -import io.trino.testing.DistributedQueryRunner; -import io.trino.testing.QueryRunner; -import org.testng.annotations.AfterClass; -import org.testng.annotations.Test; - -import java.io.IOException; -import java.lang.reflect.InvocationTargetException; -import java.nio.file.Path; -import java.util.Optional; -import java.util.concurrent.atomic.AtomicBoolean; - -import static com.google.common.io.MoreFiles.deleteRecursively; -import static com.google.common.io.RecursiveDeleteOption.ALLOW_INSECURE; -import static com.google.common.reflect.Reflection.newProxy; -import static com.google.common.util.concurrent.MoreExecutors.directExecutor; -import static com.google.inject.util.Modules.EMPTY_MODULE; -import static io.trino.plugin.hive.HiveErrorCode.HIVE_METASTORE_ERROR; -import static io.trino.plugin.hive.HiveTestUtils.HDFS_ENVIRONMENT; -import static io.trino.plugin.hive.metastore.glue.GlueClientUtil.createAsyncGlueClient; -import static io.trino.testing.TestingNames.randomNameSuffix; -import static io.trino.testing.TestingSession.testSessionBuilder; -import static org.testng.Assert.assertFalse; - -@Test(singleThreaded = true) -public class TestDeltaLakeConcurrentModificationGlueMetastore - extends AbstractTestQueryFramework -{ - private static final String CATALOG_NAME = "test_delta_lake_concurrent"; - private static final String SCHEMA = "test_delta_lake_glue_concurrent_" + randomNameSuffix(); - private Path dataDirectory; - private GlueHiveMetastore metastore; - private final AtomicBoolean failNextGlueDeleteTableCall = new AtomicBoolean(false); - - @Override - protected QueryRunner createQueryRunner() - throws Exception - { - Session deltaLakeSession = testSessionBuilder() - .setCatalog(CATALOG_NAME) - .setSchema(SCHEMA) - .build(); - - DistributedQueryRunner queryRunner = DistributedQueryRunner.builder(deltaLakeSession).build(); - - dataDirectory = queryRunner.getCoordinator().getBaseDataDir().resolve("data_delta_concurrent"); - GlueMetastoreStats stats = new GlueMetastoreStats(); - GlueHiveMetastoreConfig glueConfig = new GlueHiveMetastoreConfig() - .setDefaultWarehouseDir(dataDirectory.toUri().toString()); - - AWSGlueAsync glueClient = createAsyncGlueClient(glueConfig, DefaultAWSCredentialsProviderChain.getInstance(), ImmutableSet.of(), stats.newRequestMetricsCollector()); - AWSGlueAsync proxiedGlueClient = newProxy(AWSGlueAsync.class, (proxy, method, args) -> { - Object result; - try { - if (method.getName().equals("deleteTable") && failNextGlueDeleteTableCall.get()) { - // Simulate concurrent modifications on the table that is about to be dropped - failNextGlueDeleteTableCall.set(false); - throw new TrinoException(HIVE_METASTORE_ERROR, new ConcurrentModificationException("Test-simulated metastore concurrent modification exception")); - } - result = method.invoke(glueClient, args); - } - catch (InvocationTargetException e) { - throw e.getCause(); - } - return result; - }); - - metastore = new GlueHiveMetastore( - HDFS_ENVIRONMENT, - glueConfig, - directExecutor(), - new DefaultGlueColumnStatisticsProviderFactory(directExecutor(), directExecutor()), - proxiedGlueClient, - stats, - table -> true); - - queryRunner.installPlugin(new TestingDeltaLakePlugin(Optional.of(new TestingDeltaLakeMetastoreModule(metastore)), Optional.empty(), EMPTY_MODULE)); - queryRunner.createCatalog(CATALOG_NAME, "delta_lake"); - queryRunner.execute("CREATE SCHEMA " + SCHEMA); - return queryRunner; - } - - @Test - public void testDropTableWithConcurrentModifications() - { - String tableName = "test_glue_table_" + randomNameSuffix(); - assertUpdate("CREATE TABLE " + tableName + " AS SELECT 1 AS data", 1); - - failNextGlueDeleteTableCall.set(true); - assertUpdate("DROP TABLE " + tableName); - assertFalse(getQueryRunner().tableExists(getSession(), tableName)); - } - - @AfterClass(alwaysRun = true) - public void cleanup() - throws IOException - { - if (metastore != null) { - metastore.dropDatabase(SCHEMA, false); - deleteRecursively(dataDirectory, ALLOW_INSECURE); - } - } -} diff --git a/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/metastore/glue/TestDeltaLakeGlueMetastore.java b/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/metastore/glue/TestDeltaLakeGlueMetastore.java index 3a84a15292ed..266a2dde02d3 100644 --- a/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/metastore/glue/TestDeltaLakeGlueMetastore.java +++ b/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/metastore/glue/TestDeltaLakeGlueMetastore.java @@ -109,11 +109,11 @@ public void setUp() .put("delta.hide-non-delta-lake-tables", "true") .buildOrThrow(); + ConnectorContext context = new TestingConnectorContext(); Bootstrap app = new Bootstrap( // connector dependencies new JsonModule(), binder -> { - ConnectorContext context = new TestingConnectorContext(); binder.bind(CatalogName.class).toInstance(new CatalogName("test")); binder.bind(TypeManager.class).toInstance(context.getTypeManager()); binder.bind(NodeManager.class).toInstance(context.getNodeManager()); @@ -127,7 +127,7 @@ public void setUp() new DeltaLakeModule(), // test setup binder -> binder.bind(HdfsEnvironment.class).toInstance(HDFS_ENVIRONMENT), - new FileSystemModule()); + new FileSystemModule("test", context.getNodeManager(), context.getOpenTelemetry(), false)); Injector injector = app .doNotInitializeLogging() diff --git a/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/metastore/glue/TestDeltaLakeRegisterTableProcedureWithGlue.java b/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/metastore/glue/TestDeltaLakeRegisterTableProcedureWithGlue.java index 1aaf0f79f13a..2956e5ed9535 100644 --- a/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/metastore/glue/TestDeltaLakeRegisterTableProcedureWithGlue.java +++ b/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/metastore/glue/TestDeltaLakeRegisterTableProcedureWithGlue.java @@ -18,7 +18,7 @@ import java.nio.file.Path; -import static io.trino.plugin.hive.metastore.glue.GlueHiveMetastore.createTestingGlueHiveMetastore; +import static io.trino.plugin.hive.metastore.glue.TestingGlueHiveMetastore.createTestingGlueHiveMetastore; public class TestDeltaLakeRegisterTableProcedureWithGlue extends BaseDeltaLakeRegisterTableProcedureTest @@ -26,6 +26,6 @@ public class TestDeltaLakeRegisterTableProcedureWithGlue @Override protected HiveMetastore createTestMetastore(Path dataDirectory) { - return createTestingGlueHiveMetastore(dataDirectory); + return createTestingGlueHiveMetastore(dataDirectory, this::closeAfterClass); } } diff --git a/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/metastore/glue/TestDeltaLakeViewsGlueMetastore.java b/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/metastore/glue/TestDeltaLakeViewsGlueMetastore.java index 9e6fcbc90878..ee4dc0341b97 100644 --- a/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/metastore/glue/TestDeltaLakeViewsGlueMetastore.java +++ b/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/metastore/glue/TestDeltaLakeViewsGlueMetastore.java @@ -32,7 +32,7 @@ import static com.google.common.io.MoreFiles.deleteRecursively; import static com.google.common.io.RecursiveDeleteOption.ALLOW_INSECURE; import static com.google.inject.util.Modules.EMPTY_MODULE; -import static io.trino.plugin.hive.metastore.glue.GlueHiveMetastore.createTestingGlueHiveMetastore; +import static io.trino.plugin.hive.metastore.glue.TestingGlueHiveMetastore.createTestingGlueHiveMetastore; import static io.trino.testing.TestingNames.randomNameSuffix; import static io.trino.testing.TestingSession.testSessionBuilder; import static java.lang.String.format; @@ -47,7 +47,7 @@ public class TestDeltaLakeViewsGlueMetastore private HiveMetastore createTestMetastore(Path dataDirectory) { - return createTestingGlueHiveMetastore(dataDirectory); + return createTestingGlueHiveMetastore(dataDirectory, this::closeAfterClass); } @Override diff --git a/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/metastore/glue/TestDeltaS3AndGlueMetastoreTest.java b/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/metastore/glue/TestDeltaS3AndGlueMetastoreTest.java index cc12d940b2ee..b2d5f4c65681 100644 --- a/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/metastore/glue/TestDeltaS3AndGlueMetastoreTest.java +++ b/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/metastore/glue/TestDeltaS3AndGlueMetastoreTest.java @@ -25,7 +25,7 @@ import static com.google.common.collect.Iterables.getOnlyElement; import static io.trino.plugin.deltalake.DeltaLakeQueryRunner.DELTA_CATALOG; -import static io.trino.plugin.hive.metastore.glue.GlueHiveMetastore.createTestingGlueHiveMetastore; +import static io.trino.plugin.hive.metastore.glue.TestingGlueHiveMetastore.createTestingGlueHiveMetastore; import static java.util.Objects.requireNonNull; import static org.assertj.core.api.Assertions.assertThat; @@ -41,7 +41,7 @@ public TestDeltaS3AndGlueMetastoreTest() protected QueryRunner createQueryRunner() throws Exception { - metastore = createTestingGlueHiveMetastore(Path.of(schemaPath())); + metastore = createTestingGlueHiveMetastore(Path.of(schemaPath()), this::closeAfterClass); DistributedQueryRunner queryRunner = DeltaLakeQueryRunner.builder() .setCatalogName(DELTA_CATALOG) .setDeltaProperties(ImmutableMap.builder() diff --git a/plugin/trino-geospatial/src/test/java/io/trino/plugin/geospatial/TestSpatialJoins.java b/plugin/trino-geospatial/src/test/java/io/trino/plugin/geospatial/TestSpatialJoins.java index 8af1726a43b0..e9c721ef1e52 100644 --- a/plugin/trino-geospatial/src/test/java/io/trino/plugin/geospatial/TestSpatialJoins.java +++ b/plugin/trino-geospatial/src/test/java/io/trino/plugin/geospatial/TestSpatialJoins.java @@ -82,7 +82,7 @@ protected DistributedQueryRunner createQueryRunner() .setOwnerName(Optional.of("public")) .setOwnerType(Optional.of(PrincipalType.ROLE)) .build()); - queryRunner.installPlugin(new TestingHivePlugin(metastore)); + queryRunner.installPlugin(new TestingHivePlugin(queryRunner.getCoordinator().getBaseDataDir().resolve("hive_data"))); queryRunner.createCatalog("hive", "hive"); return queryRunner; diff --git a/plugin/trino-hive-hadoop2/src/test/java/io/trino/plugin/hive/AbstractTestHiveFileSystemAbfs.java b/plugin/trino-hive-hadoop2/src/test/java/io/trino/plugin/hive/AbstractTestHiveFileSystemAbfs.java deleted file mode 100644 index ad253f1166e8..000000000000 --- a/plugin/trino-hive-hadoop2/src/test/java/io/trino/plugin/hive/AbstractTestHiveFileSystemAbfs.java +++ /dev/null @@ -1,115 +0,0 @@ -/* - * 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.hive; - -import com.google.common.collect.ImmutableList; -import com.google.common.collect.ImmutableMap; -import com.google.common.collect.ImmutableSet; -import io.trino.hdfs.ConfigurationInitializer; -import io.trino.hdfs.DynamicHdfsConfiguration; -import io.trino.hdfs.HdfsConfig; -import io.trino.hdfs.HdfsConfiguration; -import io.trino.hdfs.HdfsConfigurationInitializer; -import io.trino.hdfs.azure.HiveAzureConfig; -import io.trino.hdfs.azure.TrinoAzureConfigurationInitializer; -import io.trino.plugin.hive.AbstractTestHive.Transaction; -import io.trino.spi.connector.ColumnMetadata; -import io.trino.spi.connector.ConnectorTableMetadata; -import io.trino.spi.connector.SchemaTableName; -import org.apache.hadoop.fs.Path; - -import java.util.Map; -import java.util.Optional; - -import static com.google.common.base.Preconditions.checkArgument; -import static io.trino.plugin.hive.HiveTableProperties.BUCKETED_BY_PROPERTY; -import static io.trino.plugin.hive.HiveTableProperties.BUCKET_COUNT_PROPERTY; -import static io.trino.plugin.hive.HiveTableProperties.EXTERNAL_LOCATION_PROPERTY; -import static io.trino.plugin.hive.HiveTableProperties.SKIP_FOOTER_LINE_COUNT; -import static io.trino.plugin.hive.HiveTableProperties.SKIP_HEADER_LINE_COUNT; -import static io.trino.plugin.hive.HiveTableProperties.SORTED_BY_PROPERTY; -import static io.trino.plugin.hive.HiveTableProperties.STORAGE_FORMAT_PROPERTY; -import static io.trino.spi.type.BigintType.BIGINT; -import static java.lang.String.format; -import static org.testng.util.Strings.isNullOrEmpty; - -public abstract class AbstractTestHiveFileSystemAbfs - extends AbstractTestHiveFileSystem -{ - protected String account; - protected String container; - protected String testDirectory; - - protected static String checkParameter(String value, String name) - { - checkArgument(!isNullOrEmpty(value), "expected non-empty %s", name); - return value; - } - - protected void setup(String host, int port, String databaseName, String container, String account, String testDirectory) - { - this.container = checkParameter(container, "container"); - this.account = checkParameter(account, "account"); - this.testDirectory = checkParameter(testDirectory, "test directory"); - super.setup( - checkParameter(host, "host"), - port, - checkParameter(databaseName, "database name"), - false, - createHdfsConfiguration()); - } - - @Override - protected void onSetupComplete() - { - ensureTableExists(table, "trino_test_external_fs", ImmutableMap.of()); - ensureTableExists(tableWithHeader, "trino_test_external_fs_with_header", ImmutableMap.of(SKIP_HEADER_LINE_COUNT, 1)); - ensureTableExists(tableWithHeaderAndFooter, "trino_test_external_fs_with_header_and_footer", ImmutableMap.of(SKIP_HEADER_LINE_COUNT, 2, SKIP_FOOTER_LINE_COUNT, 2)); - } - - private void ensureTableExists(SchemaTableName table, String tableDirectoryName, Map tableProperties) - { - try (Transaction transaction = newTransaction()) { - ConnectorTableMetadata tableMetadata = new ConnectorTableMetadata( - table, - ImmutableList.of(new ColumnMetadata("t_bigint", BIGINT)), - ImmutableMap.builder() - .putAll(tableProperties) - .put(STORAGE_FORMAT_PROPERTY, HiveStorageFormat.TEXTFILE) - .put(EXTERNAL_LOCATION_PROPERTY, getBasePath().toString() + "/" + tableDirectoryName) - .put(BUCKET_COUNT_PROPERTY, 0) - .put(BUCKETED_BY_PROPERTY, ImmutableList.of()) - .put(SORTED_BY_PROPERTY, ImmutableList.of()) - .buildOrThrow()); - if (!transaction.getMetadata().listTables(newSession(), Optional.of(table.getSchemaName())).contains(table)) { - transaction.getMetadata().createTable(newSession(), tableMetadata, false); - } - transaction.commit(); - } - } - - protected abstract HiveAzureConfig getConfig(); - - private HdfsConfiguration createHdfsConfiguration() - { - ConfigurationInitializer initializer = new TrinoAzureConfigurationInitializer(getConfig()); - return new DynamicHdfsConfiguration(new HdfsConfigurationInitializer(new HdfsConfig(), ImmutableSet.of(initializer)), ImmutableSet.of()); - } - - @Override - protected Path getBasePath() - { - return new Path(format("abfs://%s@%s.dfs.core.windows.net/%s/", container, account, testDirectory)); - } -} diff --git a/plugin/trino-hive-hadoop2/src/test/java/io/trino/plugin/hive/AbstractTestHiveFileSystemS3.java b/plugin/trino-hive-hadoop2/src/test/java/io/trino/plugin/hive/AbstractTestHiveFileSystemS3.java deleted file mode 100644 index 801ea4f667c0..000000000000 --- a/plugin/trino-hive-hadoop2/src/test/java/io/trino/plugin/hive/AbstractTestHiveFileSystemS3.java +++ /dev/null @@ -1,261 +0,0 @@ -/* - * 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.hive; - -import com.amazonaws.auth.AWSStaticCredentialsProvider; -import com.amazonaws.auth.BasicAWSCredentials; -import com.amazonaws.client.builder.AwsClientBuilder; -import com.amazonaws.services.s3.AmazonS3; -import com.amazonaws.services.s3.AmazonS3Client; -import com.amazonaws.services.s3.model.ObjectMetadata; -import com.amazonaws.services.s3.model.PutObjectRequest; -import com.google.common.collect.ImmutableList; -import com.google.common.collect.ImmutableSet; -import com.google.common.collect.Streams; -import com.google.common.net.MediaType; -import io.trino.filesystem.Location; -import io.trino.filesystem.TrinoFileSystem; -import io.trino.filesystem.hdfs.HdfsFileSystemFactory; -import io.trino.hdfs.ConfigurationInitializer; -import io.trino.hdfs.DynamicHdfsConfiguration; -import io.trino.hdfs.HdfsConfig; -import io.trino.hdfs.HdfsConfiguration; -import io.trino.hdfs.HdfsConfigurationInitializer; -import io.trino.hdfs.HdfsNamenodeStats; -import io.trino.hdfs.TrinoHdfsFileSystemStats; -import io.trino.hdfs.s3.HiveS3Config; -import io.trino.hdfs.s3.TrinoS3ConfigurationInitializer; -import io.trino.plugin.hive.fs.FileSystemDirectoryLister; -import io.trino.plugin.hive.fs.HiveFileIterator; -import io.trino.plugin.hive.fs.TrinoFileStatus; -import io.trino.plugin.hive.metastore.Column; -import io.trino.plugin.hive.metastore.StorageFormat; -import io.trino.plugin.hive.metastore.Table; -import org.apache.hadoop.fs.FileSystem; -import org.apache.hadoop.fs.Path; -import org.testng.annotations.Test; - -import java.util.Arrays; -import java.util.List; -import java.util.Optional; - -import static com.google.common.base.Preconditions.checkArgument; -import static io.trino.plugin.hive.HiveTestUtils.SESSION; -import static io.trino.plugin.hive.HiveType.HIVE_LONG; -import static io.trino.plugin.hive.HiveType.HIVE_STRING; -import static java.io.InputStream.nullInputStream; -import static java.lang.String.format; -import static org.assertj.core.api.Assertions.assertThat; -import static org.testng.Assert.assertFalse; -import static org.testng.util.Strings.isNullOrEmpty; - -public abstract class AbstractTestHiveFileSystemS3 - extends AbstractTestHiveFileSystem -{ - private static final MediaType DIRECTORY_MEDIA_TYPE = MediaType.create("application", "x-directory"); - - private String awsAccessKey; - private String awsSecretKey; - private String writableBucket; - private String testDirectory; - private AmazonS3 s3Client; - - protected void setup( - String host, - int port, - String databaseName, - String s3endpoint, - String awsAccessKey, - String awsSecretKey, - String writableBucket, - String testDirectory, - boolean s3SelectPushdownEnabled) - { - checkArgument(!isNullOrEmpty(host), "Expected non empty host"); - checkArgument(!isNullOrEmpty(databaseName), "Expected non empty databaseName"); - checkArgument(!isNullOrEmpty(awsAccessKey), "Expected non empty awsAccessKey"); - checkArgument(!isNullOrEmpty(awsSecretKey), "Expected non empty awsSecretKey"); - checkArgument(!isNullOrEmpty(s3endpoint), "Expected non empty s3endpoint"); - checkArgument(!isNullOrEmpty(writableBucket), "Expected non empty writableBucket"); - checkArgument(!isNullOrEmpty(testDirectory), "Expected non empty testDirectory"); - this.awsAccessKey = awsAccessKey; - this.awsSecretKey = awsSecretKey; - this.writableBucket = writableBucket; - this.testDirectory = testDirectory; - - s3Client = AmazonS3Client.builder() - .withEndpointConfiguration(new AwsClientBuilder.EndpointConfiguration(s3endpoint, null)) - .withCredentials(new AWSStaticCredentialsProvider(new BasicAWSCredentials(awsAccessKey, awsSecretKey))) - .build(); - - setup(host, port, databaseName, s3SelectPushdownEnabled, createHdfsConfiguration()); - } - - private HdfsConfiguration createHdfsConfiguration() - { - ConfigurationInitializer s3Config = new TrinoS3ConfigurationInitializer(new HiveS3Config() - .setS3AwsAccessKey(awsAccessKey) - .setS3AwsSecretKey(awsSecretKey)); - HdfsConfigurationInitializer initializer = new HdfsConfigurationInitializer(new HdfsConfig(), ImmutableSet.of(s3Config)); - return new DynamicHdfsConfiguration(initializer, ImmutableSet.of()); - } - - @Override - protected Path getBasePath() - { - // HDP 3.1 does not understand s3:// out of the box. - return new Path(format("s3a://%s/%s/", writableBucket, testDirectory)); - } - - @Test - public void testIgnoreHadoopFolderMarker() - throws Exception - { - Path basePath = getBasePath(); - FileSystem fs = hdfsEnvironment.getFileSystem(TESTING_CONTEXT, basePath); - - String markerFileName = "test_table_$folder$"; - Path filePath = new Path(basePath, markerFileName); - fs.create(filePath).close(); - - assertFalse(Arrays.stream(fs.listStatus(basePath)).anyMatch(file -> file.getPath().getName().equalsIgnoreCase(markerFileName))); - } - - /** - * Tests the same functionality like {@link #testFileIteratorPartitionedListing()} with the - * setup done by native {@link AmazonS3} - */ - @Test - public void testFileIteratorPartitionedListingNativeS3Client() - throws Exception - { - Table.Builder tableBuilder = Table.builder() - .setDatabaseName(table.getSchemaName()) - .setTableName(table.getTableName()) - .setDataColumns(ImmutableList.of(new Column("data", HIVE_LONG, Optional.empty()))) - .setPartitionColumns(ImmutableList.of(new Column("part", HIVE_STRING, Optional.empty()))) - .setOwner(Optional.empty()) - .setTableType("fake"); - tableBuilder.getStorageBuilder() - .setStorageFormat(StorageFormat.fromHiveStorageFormat(HiveStorageFormat.CSV)); - Table fakeTable = tableBuilder.build(); - - Path basePath = new Path(getBasePath(), "test-file-iterator-partitioned-listing-native-setup"); - FileSystem fs = hdfsEnvironment.getFileSystem(TESTING_CONTEXT, basePath); - TrinoFileSystem trinoFileSystem = new HdfsFileSystemFactory(hdfsEnvironment, new TrinoHdfsFileSystemStats()).create(SESSION); - fs.mkdirs(basePath); - String basePrefix = basePath.toUri().getPath().substring(1); - - // Expected file system tree: - // test-file-iterator-partitioned-listing-native-setup/ - // .hidden/ - // nested-file-in-hidden.txt - // part=simple/ - // _hidden-file.txt - // plain-file.txt - // part=nested/ - // parent/ - // _nested-hidden-file.txt - // nested-file.txt - // part=plus+sign/ - // plus-file.txt - // part=percent%sign/ - // percent-file.txt - // part=url%20encoded/ - // url-encoded-file.txt - // part=level1|level2/ - // pipe-file.txt - // parent1/ - // parent2/ - // deeply-nested-file.txt - // part=level1 | level2/ - // pipe-blanks-file.txt - // empty-directory/ - // .hidden-in-base.txt - - createFile(writableBucket, format("%s/.hidden/nested-file-in-hidden.txt", basePrefix)); - createFile(writableBucket, format("%s/part=simple/_hidden-file.txt", basePrefix)); - createFile(writableBucket, format("%s/part=simple/plain-file.txt", basePrefix)); - createFile(writableBucket, format("%s/part=nested/parent/_nested-hidden-file.txt", basePrefix)); - createFile(writableBucket, format("%s/part=nested/parent/nested-file.txt", basePrefix)); - createFile(writableBucket, format("%s/part=plus+sign/plus-file.txt", basePrefix)); - createFile(writableBucket, format("%s/part=percent%%sign/percent-file.txt", basePrefix)); - createFile(writableBucket, format("%s/part=url%%20encoded/url-encoded-file.txt", basePrefix)); - createFile(writableBucket, format("%s/part=level1|level2/pipe-file.txt", basePrefix)); - createFile(writableBucket, format("%s/part=level1|level2/parent1/parent2/deeply-nested-file.txt", basePrefix)); - createFile(writableBucket, format("%s/part=level1 | level2/pipe-blanks-file.txt", basePrefix)); - createDirectory(writableBucket, format("%s/empty-directory/", basePrefix)); - createFile(writableBucket, format("%s/.hidden-in-base.txt", basePrefix)); - - // List recursively through hive file iterator - HiveFileIterator recursiveIterator = new HiveFileIterator( - fakeTable, - Location.of(basePath.toString()), - trinoFileSystem, - new FileSystemDirectoryLister(), - new HdfsNamenodeStats(), - HiveFileIterator.NestedDirectoryPolicy.RECURSE); - - List recursiveListing = Streams.stream(recursiveIterator) - .map(TrinoFileStatus::getPath) - .toList(); - // Should not include directories, or files underneath hidden directories - assertThat(recursiveListing).containsExactlyInAnyOrder( - format("%s/part=simple/plain-file.txt", basePath), - format("%s/part=nested/parent/nested-file.txt", basePath), - format("%s/part=plus+sign/plus-file.txt", basePath), - format("%s/part=percent%%sign/percent-file.txt", basePath), - format("%s/part=url%%20encoded/url-encoded-file.txt", basePath), - format("%s/part=level1|level2/pipe-file.txt", basePath), - format("%s/part=level1|level2/parent1/parent2/deeply-nested-file.txt", basePath), - format("%s/part=level1 | level2/pipe-blanks-file.txt", basePath)); - - HiveFileIterator shallowIterator = new HiveFileIterator( - fakeTable, - Location.of(basePath.toString()), - trinoFileSystem, - new FileSystemDirectoryLister(), - new HdfsNamenodeStats(), - HiveFileIterator.NestedDirectoryPolicy.IGNORED); - List shallowListing = Streams.stream(shallowIterator) - .map(TrinoFileStatus::getPath) - .map(Path::new) - .toList(); - // Should not include any hidden files, folders, or nested files - assertThat(shallowListing).isEmpty(); - } - - protected void createDirectory(String bucketName, String key) - { - // create meta-data for your folder and set content-length to 0 - ObjectMetadata metadata = new ObjectMetadata(); - metadata.setContentLength(0); - metadata.setContentType(DIRECTORY_MEDIA_TYPE.toString()); - // create a PutObjectRequest passing the folder name suffixed by / - if (!key.endsWith("/")) { - key += "/"; - } - PutObjectRequest putObjectRequest = new PutObjectRequest(bucketName, key, nullInputStream(), metadata); - // send request to S3 to create folder - s3Client.putObject(putObjectRequest); - } - - protected void createFile(String bucketName, String key) - { - ObjectMetadata metadata = new ObjectMetadata(); - metadata.setContentLength(0); - PutObjectRequest putObjectRequest = new PutObjectRequest(bucketName, key, nullInputStream(), metadata); - s3Client.putObject(putObjectRequest); - } -} diff --git a/plugin/trino-hive-hadoop2/src/test/java/io/trino/plugin/hive/TestHive.java b/plugin/trino-hive-hadoop2/src/test/java/io/trino/plugin/hive/TestHive.java deleted file mode 100644 index 06c7b65ef34c..000000000000 --- a/plugin/trino-hive-hadoop2/src/test/java/io/trino/plugin/hive/TestHive.java +++ /dev/null @@ -1,184 +0,0 @@ -/* - * 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.hive; - -import com.google.common.collect.ImmutableList; -import com.google.common.net.HostAndPort; -import io.trino.spi.connector.ConnectorMetadata; -import io.trino.spi.connector.SchemaTableName; -import io.trino.spi.connector.SchemaTablePrefix; -import org.apache.hadoop.net.NetUtils; -import org.testng.SkipException; -import org.testng.annotations.BeforeClass; -import org.testng.annotations.Parameters; -import org.testng.annotations.Test; - -import static org.assertj.core.api.Assertions.assertThat; -import static org.assertj.core.api.Assertions.assertThatThrownBy; - -// staging directory is shared mutable state -@Test(singleThreaded = true) -public class TestHive - extends AbstractTestHive -{ - @Parameters({"test.metastore", "test.database"}) - @BeforeClass - public void initialize(String metastore, String database) - { - String hadoopMasterIp = System.getProperty("hadoop-master-ip"); - if (hadoopMasterIp != null) { - // Even though Hadoop is accessed by proxy, Hadoop still tries to resolve hadoop-master - // (e.g: in: NameNodeProxies.createProxy) - // This adds a static resolution for hadoop-master to docker container internal ip - NetUtils.addStaticResolution("hadoop-master", hadoopMasterIp); - } - - setup(HostAndPort.fromString(metastore), database); - } - - @Test - public void forceTestNgToRespectSingleThreaded() - { - // TODO: Remove after updating TestNG to 7.4.0+ (https://github.com/trinodb/trino/issues/8571) - // TestNG doesn't enforce @Test(singleThreaded = true) when tests are defined in base class. According to - // https://github.com/cbeust/testng/issues/2361#issuecomment-688393166 a workaround it to add a dummy test to the leaf test class. - } - - @Override - public void testHideDeltaLakeTables() - { - assertThatThrownBy(super::testHideDeltaLakeTables) - .hasMessageMatching("(?s)\n" + - "Expecting\n" + - " \\[.*\\b(\\w+.tmp_trino_test_trino_delta_lake_table_\\w+)\\b.*]\n" + - "not to contain\n" + - " \\[\\1]\n" + - "but found.*"); - - throw new SkipException("not supported"); - } - - @Test - public void testHiveViewsHaveNoColumns() - { - try (Transaction transaction = newTransaction()) { - ConnectorMetadata metadata = transaction.getMetadata(); - assertThat(listTableColumns(metadata, newSession(), new SchemaTablePrefix(view.getSchemaName(), view.getTableName()))) - .isEmpty(); - } - } - - @Test - public void testHiveViewTranslationError() - { - try (Transaction transaction = newTransaction()) { - assertThatThrownBy(() -> transaction.getMetadata().getView(newSession(), view)) - .isInstanceOf(HiveViewNotSupportedException.class) - .hasMessageContaining("Hive views are not supported"); - - // TODO: combine this with tests for successful translation (currently in TestHiveViews product test) - } - } - - @Override - public void testUpdateBasicPartitionStatistics() - throws Exception - { - SchemaTableName tableName = temporaryTable("update_basic_partition_statistics"); - try { - createDummyPartitionedTable(tableName, STATISTICS_PARTITIONED_TABLE_COLUMNS); - // When the table has partitions, but row count statistics are set to zero, we treat this case as empty - // statistics to avoid underestimation in the CBO. This scenario may be caused when other engines are - // used to ingest data into partitioned hive tables. - testUpdatePartitionStatistics( - tableName, - PartitionStatistics.empty(), - ImmutableList.of(BASIC_STATISTICS_1, BASIC_STATISTICS_2), - ImmutableList.of(BASIC_STATISTICS_2, BASIC_STATISTICS_1)); - } - finally { - dropTable(tableName); - } - } - - @Override - public void testUpdatePartitionColumnStatistics() - throws Exception - { - SchemaTableName tableName = temporaryTable("update_partition_column_statistics"); - try { - createDummyPartitionedTable(tableName, STATISTICS_PARTITIONED_TABLE_COLUMNS); - // When the table has partitions, but row count statistics are set to zero, we treat this case as empty - // statistics to avoid underestimation in the CBO. This scenario may be caused when other engines are - // used to ingest data into partitioned hive tables. - testUpdatePartitionStatistics( - tableName, - PartitionStatistics.empty(), - ImmutableList.of(STATISTICS_1_1, STATISTICS_1_2, STATISTICS_2), - ImmutableList.of(STATISTICS_1_2, STATISTICS_1_1, STATISTICS_2)); - } - finally { - dropTable(tableName); - } - } - - @Override - public void testUpdatePartitionColumnStatisticsEmptyOptionalFields() - throws Exception - { - SchemaTableName tableName = temporaryTable("update_partition_column_statistics"); - try { - createDummyPartitionedTable(tableName, STATISTICS_PARTITIONED_TABLE_COLUMNS); - // When the table has partitions, but row count statistics are set to zero, we treat this case as empty - // statistics to avoid underestimation in the CBO. This scenario may be caused when other engines are - // used to ingest data into partitioned hive tables. - testUpdatePartitionStatistics( - tableName, - PartitionStatistics.empty(), - ImmutableList.of(STATISTICS_EMPTY_OPTIONAL_FIELDS), - ImmutableList.of(STATISTICS_EMPTY_OPTIONAL_FIELDS)); - } - finally { - dropTable(tableName); - } - } - - @Override - public void testStorePartitionWithStatistics() - throws Exception - { - // When the table has partitions, but row count statistics are set to zero, we treat this case as empty - // statistics to avoid underestimation in the CBO. This scenario may be caused when other engines are - // used to ingest data into partitioned hive tables. - testStorePartitionWithStatistics(STATISTICS_PARTITIONED_TABLE_COLUMNS, STATISTICS_1, STATISTICS_2, STATISTICS_1_1, PartitionStatistics.empty()); - } - - @Override - public void testDataColumnProperties() - { - // Column properties are currently not supported in ThriftHiveMetastore - assertThatThrownBy(super::testDataColumnProperties) - .isInstanceOf(IllegalArgumentException.class) - .hasMessage("Persisting column properties is not supported: Column{name=id, type=bigint}"); - } - - @Override - public void testPartitionColumnProperties() - { - // Column properties are currently not supported in ThriftHiveMetastore - assertThatThrownBy(super::testPartitionColumnProperties) - .isInstanceOf(IllegalArgumentException.class) - .hasMessage("Persisting column properties is not supported: Column{name=part_key, type=varchar(256)}"); - } -} diff --git a/plugin/trino-hive-hadoop2/src/test/java/io/trino/plugin/hive/TestHiveFileSystemAbfsAccessKey.java b/plugin/trino-hive-hadoop2/src/test/java/io/trino/plugin/hive/TestHiveFileSystemAbfsAccessKey.java deleted file mode 100644 index cf9de14fb450..000000000000 --- a/plugin/trino-hive-hadoop2/src/test/java/io/trino/plugin/hive/TestHiveFileSystemAbfsAccessKey.java +++ /dev/null @@ -1,48 +0,0 @@ -/* - * 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.hive; - -import io.trino.hdfs.azure.HiveAzureConfig; -import org.testng.annotations.BeforeClass; -import org.testng.annotations.Parameters; - -public class TestHiveFileSystemAbfsAccessKey - extends AbstractTestHiveFileSystemAbfs -{ - private String accessKey; - - @Parameters({ - "hive.hadoop2.metastoreHost", - "hive.hadoop2.metastorePort", - "hive.hadoop2.databaseName", - "hive.hadoop2.abfs.container", - "hive.hadoop2.abfs.account", - "hive.hadoop2.abfs.accessKey", - "hive.hadoop2.abfs.testDirectory", - }) - @BeforeClass - public void setup(String host, int port, String databaseName, String container, String account, String accessKey, String testDirectory) - { - this.accessKey = checkParameter(accessKey, "access key"); - super.setup(host, port, databaseName, container, account, testDirectory); - } - - @Override - protected HiveAzureConfig getConfig() - { - return new HiveAzureConfig() - .setAbfsAccessKey(accessKey) - .setAbfsStorageAccount(account); - } -} diff --git a/plugin/trino-hive-hadoop2/src/test/java/io/trino/plugin/hive/TestHiveFileSystemAbfsOAuth.java b/plugin/trino-hive-hadoop2/src/test/java/io/trino/plugin/hive/TestHiveFileSystemAbfsOAuth.java deleted file mode 100644 index e2e5cf2c61da..000000000000 --- a/plugin/trino-hive-hadoop2/src/test/java/io/trino/plugin/hive/TestHiveFileSystemAbfsOAuth.java +++ /dev/null @@ -1,64 +0,0 @@ -/* - * 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.hive; - -import io.trino.hdfs.azure.HiveAzureConfig; -import org.testng.annotations.BeforeClass; -import org.testng.annotations.Parameters; - -public class TestHiveFileSystemAbfsOAuth - extends AbstractTestHiveFileSystemAbfs -{ - private String endpoint; - private String clientId; - private String secret; - - @Parameters({ - "hive.hadoop2.metastoreHost", - "hive.hadoop2.metastorePort", - "hive.hadoop2.databaseName", - "test.hive.azure.abfs.container", - "test.hive.azure.abfs.storage-account", - "test.hive.azure.abfs.test-directory", - "test.hive.azure.abfs.oauth.endpoint", - "test.hive.azure.abfs.oauth.client-id", - "test.hive.azure.abfs.oauth.secret", - }) - @BeforeClass - public void setup( - String host, - int port, - String databaseName, - String container, - String account, - String testDirectory, - String clientEndpoint, - String clientId, - String clientSecret) - { - this.endpoint = checkParameter(clientEndpoint, "endpoint"); - this.clientId = checkParameter(clientId, "client ID"); - this.secret = checkParameter(clientSecret, "secret"); - super.setup(host, port, databaseName, container, account, testDirectory); - } - - @Override - protected HiveAzureConfig getConfig() - { - return new HiveAzureConfig() - .setAbfsOAuthClientEndpoint(endpoint) - .setAbfsOAuthClientId(clientId) - .setAbfsOAuthClientSecret(secret); - } -} diff --git a/plugin/trino-hive-hadoop2/src/test/java/io/trino/plugin/hive/TestHiveFileSystemAdl.java b/plugin/trino-hive-hadoop2/src/test/java/io/trino/plugin/hive/TestHiveFileSystemAdl.java deleted file mode 100644 index fad889d28fac..000000000000 --- a/plugin/trino-hive-hadoop2/src/test/java/io/trino/plugin/hive/TestHiveFileSystemAdl.java +++ /dev/null @@ -1,162 +0,0 @@ -/* - * 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.hive; - -import com.google.common.collect.ImmutableSet; -import io.trino.hdfs.ConfigurationInitializer; -import io.trino.hdfs.DynamicHdfsConfiguration; -import io.trino.hdfs.HdfsConfig; -import io.trino.hdfs.HdfsConfiguration; -import io.trino.hdfs.HdfsConfigurationInitializer; -import io.trino.hdfs.azure.HiveAzureConfig; -import io.trino.hdfs.azure.TrinoAzureConfigurationInitializer; -import org.apache.hadoop.fs.FileSystem; -import org.apache.hadoop.fs.Path; -import org.testng.annotations.BeforeClass; -import org.testng.annotations.Parameters; -import org.testng.annotations.Test; - -import java.io.FileNotFoundException; -import java.util.UUID; - -import static com.google.common.base.Preconditions.checkArgument; -import static java.lang.String.format; -import static org.assertj.core.api.Assertions.assertThatThrownBy; -import static org.testng.Assert.assertFalse; -import static org.testng.Assert.assertTrue; -import static org.testng.util.Strings.isNullOrEmpty; - -public class TestHiveFileSystemAdl - extends AbstractTestHiveFileSystem -{ - private String dataLakeName; - private String clientId; - private String credential; - private String refreshUrl; - private String testDirectory; - - @Parameters({ - "hive.hadoop2.metastoreHost", - "hive.hadoop2.metastorePort", - "hive.hadoop2.databaseName", - "hive.hadoop2.adl.name", - "hive.hadoop2.adl.clientId", - "hive.hadoop2.adl.credential", - "hive.hadoop2.adl.refreshUrl", - "hive.hadoop2.adl.testDirectory", - }) - @BeforeClass - public void setup(String host, int port, String databaseName, String dataLakeName, String clientId, String credential, String refreshUrl, String testDirectory) - { - checkArgument(!isNullOrEmpty(host), "expected non empty host"); - checkArgument(!isNullOrEmpty(databaseName), "expected non empty databaseName"); - checkArgument(!isNullOrEmpty(dataLakeName), "expected non empty dataLakeName"); - checkArgument(!isNullOrEmpty(clientId), "expected non empty clientId"); - checkArgument(!isNullOrEmpty(credential), "expected non empty credential"); - checkArgument(!isNullOrEmpty(refreshUrl), "expected non empty refreshUrl"); - checkArgument(!isNullOrEmpty(testDirectory), "expected non empty testDirectory"); - - this.dataLakeName = dataLakeName; - this.clientId = clientId; - this.credential = credential; - this.refreshUrl = refreshUrl; - this.testDirectory = testDirectory; - - super.setup(host, port, databaseName, false, createHdfsConfiguration()); - } - - private HdfsConfiguration createHdfsConfiguration() - { - ConfigurationInitializer azureConfig = new TrinoAzureConfigurationInitializer(new HiveAzureConfig() - .setAdlClientId(clientId) - .setAdlCredential(credential) - .setAdlRefreshUrl(refreshUrl)); - return new DynamicHdfsConfiguration(new HdfsConfigurationInitializer(new HdfsConfig(), ImmutableSet.of(azureConfig)), ImmutableSet.of()); - } - - @Override - protected Path getBasePath() - { - return new Path(format("adl://%s.azuredatalakestore.net/%s/", dataLakeName, testDirectory)); - } - - @Override - @Test - public void testRename() - throws Exception - { - Path basePath = new Path(getBasePath(), UUID.randomUUID().toString()); - FileSystem fs = hdfsEnvironment.getFileSystem(TESTING_CONTEXT, basePath); - assertFalse(fs.exists(basePath)); - - // create file foo.txt - Path path = new Path(basePath, "foo.txt"); - assertTrue(fs.createNewFile(path)); - assertTrue(fs.exists(path)); - - // rename foo.txt to bar.txt when bar does not exist - Path newPath = new Path(basePath, "bar.txt"); - assertFalse(fs.exists(newPath)); - assertTrue(fs.rename(path, newPath)); - assertFalse(fs.exists(path)); - assertTrue(fs.exists(newPath)); - - // rename foo.txt to foo.txt when foo.txt does not exist - // This fails with error no such file in ADLFileSystem - assertThatThrownBy(() -> fs.rename(path, path)) - .isInstanceOf(FileNotFoundException.class); - - // create file foo.txt and rename to existing bar.txt - assertTrue(fs.createNewFile(path)); - assertFalse(fs.rename(path, newPath)); - - // rename foo.txt to foo.txt when foo.txt exists - // This returns true in ADLFileSystem - assertTrue(fs.rename(path, path)); - - // delete foo.txt - assertTrue(fs.delete(path, false)); - assertFalse(fs.exists(path)); - - // create directory source with file - Path source = new Path(basePath, "source"); - assertTrue(fs.createNewFile(new Path(source, "test.txt"))); - - // rename source to non-existing target - Path target = new Path(basePath, "target"); - assertFalse(fs.exists(target)); - assertTrue(fs.rename(source, target)); - assertFalse(fs.exists(source)); - assertTrue(fs.exists(target)); - - // create directory source with file - assertTrue(fs.createNewFile(new Path(source, "test.txt"))); - - // rename source to existing target - assertTrue(fs.rename(source, target)); - assertFalse(fs.exists(source)); - target = new Path(target, "source"); - assertTrue(fs.exists(target)); - assertTrue(fs.exists(new Path(target, "test.txt"))); - - // delete target - target = new Path(basePath, "target"); - assertTrue(fs.exists(target)); - assertTrue(fs.delete(target, true)); - assertFalse(fs.exists(target)); - - // cleanup - fs.delete(basePath, true); - } -} diff --git a/plugin/trino-hive-hadoop2/src/test/java/io/trino/plugin/hive/TestHiveFileSystemS3.java b/plugin/trino-hive-hadoop2/src/test/java/io/trino/plugin/hive/TestHiveFileSystemS3.java deleted file mode 100644 index c522c25d6a10..000000000000 --- a/plugin/trino-hive-hadoop2/src/test/java/io/trino/plugin/hive/TestHiveFileSystemS3.java +++ /dev/null @@ -1,37 +0,0 @@ -/* - * 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.hive; - -import org.testng.annotations.BeforeClass; -import org.testng.annotations.Parameters; - -public class TestHiveFileSystemS3 - extends AbstractTestHiveFileSystemS3 -{ - @Parameters({ - "hive.hadoop2.metastoreHost", - "hive.hadoop2.metastorePort", - "hive.hadoop2.databaseName", - "hive.hadoop2.s3.endpoint", - "hive.hadoop2.s3.awsAccessKey", - "hive.hadoop2.s3.awsSecretKey", - "hive.hadoop2.s3.writableBucket", - "hive.hadoop2.s3.testDirectory", - }) - @BeforeClass - public void setup(String host, int port, String databaseName, String s3endpoint, String awsAccessKey, String awsSecretKey, String writableBucket, String testDirectory) - { - super.setup(host, port, databaseName, s3endpoint, awsAccessKey, awsSecretKey, writableBucket, testDirectory, false); - } -} diff --git a/plugin/trino-hive-hadoop2/src/test/java/io/trino/plugin/hive/TestHiveFileSystemWasb.java b/plugin/trino-hive-hadoop2/src/test/java/io/trino/plugin/hive/TestHiveFileSystemWasb.java deleted file mode 100644 index 3ac98d86636a..000000000000 --- a/plugin/trino-hive-hadoop2/src/test/java/io/trino/plugin/hive/TestHiveFileSystemWasb.java +++ /dev/null @@ -1,80 +0,0 @@ -/* - * 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.hive; - -import com.google.common.collect.ImmutableSet; -import io.trino.hdfs.ConfigurationInitializer; -import io.trino.hdfs.DynamicHdfsConfiguration; -import io.trino.hdfs.HdfsConfig; -import io.trino.hdfs.HdfsConfiguration; -import io.trino.hdfs.HdfsConfigurationInitializer; -import io.trino.hdfs.azure.HiveAzureConfig; -import io.trino.hdfs.azure.TrinoAzureConfigurationInitializer; -import org.apache.hadoop.fs.Path; -import org.testng.annotations.BeforeClass; -import org.testng.annotations.Parameters; - -import static com.google.common.base.Preconditions.checkArgument; -import static java.lang.String.format; -import static org.testng.util.Strings.isNullOrEmpty; - -public class TestHiveFileSystemWasb - extends AbstractTestHiveFileSystem -{ - private String container; - private String account; - private String accessKey; - private String testDirectory; - - @Parameters({ - "hive.hadoop2.metastoreHost", - "hive.hadoop2.metastorePort", - "hive.hadoop2.databaseName", - "hive.hadoop2.wasb.container", - "hive.hadoop2.wasb.account", - "hive.hadoop2.wasb.accessKey", - "hive.hadoop2.wasb.testDirectory", - }) - @BeforeClass - public void setup(String host, int port, String databaseName, String container, String account, String accessKey, String testDirectory) - { - checkArgument(!isNullOrEmpty(host), "expected non empty host"); - checkArgument(!isNullOrEmpty(databaseName), "expected non empty databaseName"); - checkArgument(!isNullOrEmpty(container), "expected non empty container"); - checkArgument(!isNullOrEmpty(account), "expected non empty account"); - checkArgument(!isNullOrEmpty(accessKey), "expected non empty accessKey"); - checkArgument(!isNullOrEmpty(testDirectory), "expected non empty testDirectory"); - - this.container = container; - this.account = account; - this.accessKey = accessKey; - this.testDirectory = testDirectory; - - super.setup(host, port, databaseName, false, createHdfsConfiguration()); - } - - private HdfsConfiguration createHdfsConfiguration() - { - ConfigurationInitializer wasbConfig = new TrinoAzureConfigurationInitializer(new HiveAzureConfig() - .setWasbAccessKey(accessKey) - .setWasbStorageAccount(account)); - return new DynamicHdfsConfiguration(new HdfsConfigurationInitializer(new HdfsConfig(), ImmutableSet.of(wasbConfig)), ImmutableSet.of()); - } - - @Override - protected Path getBasePath() - { - return new Path(format("wasb://%s@%s.blob.core.windows.net/%s/", container, account, testDirectory)); - } -} diff --git a/plugin/trino-hive-hadoop2/src/test/java/io/trino/plugin/hive/s3select/S3SelectTestHelper.java b/plugin/trino-hive-hadoop2/src/test/java/io/trino/plugin/hive/s3select/S3SelectTestHelper.java deleted file mode 100644 index 00a8d48dbe0f..000000000000 --- a/plugin/trino-hive-hadoop2/src/test/java/io/trino/plugin/hive/s3select/S3SelectTestHelper.java +++ /dev/null @@ -1,293 +0,0 @@ -/* - * 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.hive.s3select; - -import com.google.common.collect.ImmutableMap; -import com.google.common.collect.ImmutableSet; -import com.google.common.net.HostAndPort; -import io.airlift.concurrent.BoundedExecutor; -import io.airlift.json.JsonCodec; -import io.airlift.stats.CounterStat; -import io.trino.filesystem.hdfs.HdfsFileSystemFactory; -import io.trino.hdfs.ConfigurationInitializer; -import io.trino.hdfs.DynamicHdfsConfiguration; -import io.trino.hdfs.HdfsConfig; -import io.trino.hdfs.HdfsConfiguration; -import io.trino.hdfs.HdfsConfigurationInitializer; -import io.trino.hdfs.HdfsEnvironment; -import io.trino.hdfs.HdfsNamenodeStats; -import io.trino.hdfs.authentication.NoHdfsAuthentication; -import io.trino.hdfs.s3.HiveS3Config; -import io.trino.hdfs.s3.TrinoS3ConfigurationInitializer; -import io.trino.plugin.base.CatalogName; -import io.trino.plugin.hive.AbstractTestHiveFileSystem.TestingHiveMetastore; -import io.trino.plugin.hive.DefaultHiveMaterializedViewMetadataFactory; -import io.trino.plugin.hive.GenericHiveRecordCursorProvider; -import io.trino.plugin.hive.HiveConfig; -import io.trino.plugin.hive.HiveLocationService; -import io.trino.plugin.hive.HiveMetadataFactory; -import io.trino.plugin.hive.HivePageSourceProvider; -import io.trino.plugin.hive.HivePartitionManager; -import io.trino.plugin.hive.HiveSplitManager; -import io.trino.plugin.hive.HiveTransactionManager; -import io.trino.plugin.hive.LocationService; -import io.trino.plugin.hive.NodeVersion; -import io.trino.plugin.hive.NoneHiveRedirectionsProvider; -import io.trino.plugin.hive.PartitionUpdate; -import io.trino.plugin.hive.PartitionsSystemTableProvider; -import io.trino.plugin.hive.PropertiesSystemTableProvider; -import io.trino.plugin.hive.aws.athena.PartitionProjectionService; -import io.trino.plugin.hive.fs.FileSystemDirectoryLister; -import io.trino.plugin.hive.fs.TransactionScopeCachingDirectoryListerFactory; -import io.trino.plugin.hive.metastore.HiveMetastoreConfig; -import io.trino.plugin.hive.metastore.HiveMetastoreFactory; -import io.trino.plugin.hive.metastore.thrift.BridgingHiveMetastore; -import io.trino.plugin.hive.security.SqlStandardAccessControlMetadata; -import io.trino.spi.connector.ColumnHandle; -import io.trino.spi.connector.ConnectorPageSourceProvider; -import io.trino.spi.connector.ConnectorSession; -import io.trino.spi.connector.ConnectorSplitManager; -import io.trino.spi.connector.SchemaTableName; -import io.trino.spi.type.TestingTypeManager; -import io.trino.testing.MaterializedResult; -import org.apache.hadoop.fs.Path; - -import java.io.IOException; -import java.util.List; -import java.util.concurrent.ExecutorService; -import java.util.concurrent.ScheduledExecutorService; -import java.util.stream.LongStream; - -import static com.google.common.base.Preconditions.checkArgument; -import static com.google.common.util.concurrent.MoreExecutors.newDirectExecutorService; -import static io.airlift.concurrent.Threads.daemonThreadsNamed; -import static io.trino.plugin.hive.HiveFileSystemTestUtils.filterTable; -import static io.trino.plugin.hive.HiveFileSystemTestUtils.getSplitsCount; -import static io.trino.plugin.hive.HiveTestUtils.HDFS_FILE_SYSTEM_STATS; -import static io.trino.plugin.hive.HiveTestUtils.getDefaultHiveFileWriterFactories; -import static io.trino.plugin.hive.HiveTestUtils.getDefaultHivePageSourceFactories; -import static io.trino.plugin.hive.HiveTestUtils.getDefaultHiveRecordCursorProviders; -import static io.trino.plugin.hive.TestingThriftHiveMetastoreBuilder.testingThriftHiveMetastoreBuilder; -import static io.trino.spi.connector.MetadataProvider.NOOP_METADATA_PROVIDER; -import static io.trino.spi.type.BigintType.BIGINT; -import static io.trino.type.InternalTypeManager.TESTING_TYPE_MANAGER; -import static java.lang.String.format; -import static java.util.concurrent.Executors.newCachedThreadPool; -import static java.util.concurrent.Executors.newScheduledThreadPool; -import static org.testng.util.Strings.isNullOrEmpty; - -public class S3SelectTestHelper -{ - private HdfsEnvironment hdfsEnvironment; - private LocationService locationService; - private TestingHiveMetastore metastoreClient; - private HiveMetadataFactory metadataFactory; - private HiveTransactionManager transactionManager; - private ConnectorSplitManager splitManager; - private ConnectorPageSourceProvider pageSourceProvider; - - private ExecutorService executorService; - private HiveConfig hiveConfig; - private ScheduledExecutorService heartbeatService; - - public S3SelectTestHelper(String host, - int port, - String databaseName, - String awsAccessKey, - String awsSecretKey, - String writableBucket, - String testDirectory, - HiveConfig hiveConfig) - { - checkArgument(!isNullOrEmpty(host), "Expected non empty host"); - checkArgument(!isNullOrEmpty(databaseName), "Expected non empty databaseName"); - checkArgument(!isNullOrEmpty(awsAccessKey), "Expected non empty awsAccessKey"); - checkArgument(!isNullOrEmpty(awsSecretKey), "Expected non empty awsSecretKey"); - checkArgument(!isNullOrEmpty(writableBucket), "Expected non empty writableBucket"); - checkArgument(!isNullOrEmpty(testDirectory), "Expected non empty testDirectory"); - - executorService = newCachedThreadPool(daemonThreadsNamed("s3select-tests-%s")); - heartbeatService = newScheduledThreadPool(1); - - ConfigurationInitializer s3Config = new TrinoS3ConfigurationInitializer(new HiveS3Config() - .setS3AwsAccessKey(awsAccessKey) - .setS3AwsSecretKey(awsSecretKey)); - HdfsConfigurationInitializer initializer = new HdfsConfigurationInitializer(new HdfsConfig(), ImmutableSet.of(s3Config)); - HdfsConfiguration hdfsConfiguration = new DynamicHdfsConfiguration(initializer, ImmutableSet.of()); - - this.hiveConfig = hiveConfig; - HivePartitionManager hivePartitionManager = new HivePartitionManager(this.hiveConfig); - - hdfsEnvironment = new HdfsEnvironment(hdfsConfiguration, new HdfsConfig(), new NoHdfsAuthentication()); - locationService = new HiveLocationService(hdfsEnvironment, hiveConfig); - JsonCodec partitionUpdateCodec = JsonCodec.jsonCodec(PartitionUpdate.class); - - metastoreClient = new TestingHiveMetastore( - new BridgingHiveMetastore( - testingThriftHiveMetastoreBuilder() - .metastoreClient(HostAndPort.fromParts(host, port)) - .hiveConfig(this.hiveConfig) - .hdfsEnvironment(hdfsEnvironment) - .build()), - new Path(format("s3a://%s/%s/", writableBucket, testDirectory)), - hdfsEnvironment); - metadataFactory = new HiveMetadataFactory( - new CatalogName("hive"), - this.hiveConfig, - new HiveMetastoreConfig(), - HiveMetastoreFactory.ofInstance(metastoreClient), - getDefaultHiveFileWriterFactories(hiveConfig, hdfsEnvironment), - new HdfsFileSystemFactory(hdfsEnvironment, HDFS_FILE_SYSTEM_STATS), - hdfsEnvironment, - hivePartitionManager, - newDirectExecutorService(), - heartbeatService, - TESTING_TYPE_MANAGER, - NOOP_METADATA_PROVIDER, - locationService, - partitionUpdateCodec, - new NodeVersion("test_version"), - new NoneHiveRedirectionsProvider(), - ImmutableSet.of( - new PartitionsSystemTableProvider(hivePartitionManager, TESTING_TYPE_MANAGER), - new PropertiesSystemTableProvider()), - new DefaultHiveMaterializedViewMetadataFactory(), - SqlStandardAccessControlMetadata::new, - new FileSystemDirectoryLister(), - new TransactionScopeCachingDirectoryListerFactory(hiveConfig), - new PartitionProjectionService(this.hiveConfig, ImmutableMap.of(), new TestingTypeManager()), - true); - transactionManager = new HiveTransactionManager(metadataFactory); - - splitManager = new HiveSplitManager( - transactionManager, - hivePartitionManager, - new HdfsFileSystemFactory(hdfsEnvironment, HDFS_FILE_SYSTEM_STATS), - new HdfsNamenodeStats(), - hdfsEnvironment, - new BoundedExecutor(executorService, this.hiveConfig.getMaxSplitIteratorThreads()), - new CounterStat(), - this.hiveConfig.getMaxOutstandingSplits(), - this.hiveConfig.getMaxOutstandingSplitsSize(), - this.hiveConfig.getMinPartitionBatchSize(), - this.hiveConfig.getMaxPartitionBatchSize(), - this.hiveConfig.getMaxInitialSplits(), - this.hiveConfig.getSplitLoaderConcurrency(), - this.hiveConfig.getMaxSplitsPerSecond(), - this.hiveConfig.getRecursiveDirWalkerEnabled(), - TESTING_TYPE_MANAGER, - this.hiveConfig.getMaxPartitionsPerScan()); - - pageSourceProvider = new HivePageSourceProvider( - TESTING_TYPE_MANAGER, - hdfsEnvironment, - this.hiveConfig, - getDefaultHivePageSourceFactories(hdfsEnvironment, this.hiveConfig), - getDefaultHiveRecordCursorProviders(this.hiveConfig, hdfsEnvironment), - new GenericHiveRecordCursorProvider(hdfsEnvironment, this.hiveConfig)); - } - - public S3SelectTestHelper(String host, - int port, - String databaseName, - String awsAccessKey, - String awsSecretKey, - String writableBucket, - String testDirectory) - { - this(host, port, databaseName, awsAccessKey, awsSecretKey, writableBucket, testDirectory, new HiveConfig().setS3SelectPushdownEnabled(true)); - } - - public HiveTransactionManager getTransactionManager() - { - return transactionManager; - } - - public ConnectorSplitManager getSplitManager() - { - return splitManager; - } - - public ConnectorPageSourceProvider getPageSourceProvider() - { - return pageSourceProvider; - } - - public HiveConfig getHiveConfig() - { - return hiveConfig; - } - - public void tearDown() - { - hdfsEnvironment = null; - locationService = null; - metastoreClient = null; - metadataFactory = null; - transactionManager = null; - splitManager = null; - pageSourceProvider = null; - hiveConfig = null; - if (executorService != null) { - executorService.shutdownNow(); - executorService = null; - } - if (heartbeatService != null) { - heartbeatService.shutdownNow(); - heartbeatService = null; - } - } - - int getTableSplitsCount(SchemaTableName table) - { - return getSplitsCount( - table, - getTransactionManager(), - getHiveConfig(), - getSplitManager()); - } - - MaterializedResult getFilteredTableResult(SchemaTableName table, ColumnHandle column) - { - try { - return filterTable( - table, - List.of(column), - getTransactionManager(), - getHiveConfig(), - getPageSourceProvider(), - getSplitManager()); - } - catch (IOException ignored) { - } - - return null; - } - - static MaterializedResult expectedResult(ConnectorSession session, int start, int end) - { - MaterializedResult.Builder builder = MaterializedResult.resultBuilder(session, BIGINT); - LongStream.rangeClosed(start, end).forEach(builder::row); - return builder.build(); - } - - static boolean isSplitCountInOpenInterval(int splitCount, - int lowerBound, - int upperBound) - { - // Split number may vary, the minimum number of splits being obtained with - // the first split of maxInitialSplitSize and the rest of maxSplitSize - return lowerBound < splitCount && splitCount < upperBound; - } -} diff --git a/plugin/trino-hive-hadoop2/src/test/java/io/trino/plugin/hive/s3select/TestHiveFileSystemS3SelectCsvPushdownWithSplits.java b/plugin/trino-hive-hadoop2/src/test/java/io/trino/plugin/hive/s3select/TestHiveFileSystemS3SelectCsvPushdownWithSplits.java deleted file mode 100644 index 2edc5bd71f0b..000000000000 --- a/plugin/trino-hive-hadoop2/src/test/java/io/trino/plugin/hive/s3select/TestHiveFileSystemS3SelectCsvPushdownWithSplits.java +++ /dev/null @@ -1,116 +0,0 @@ -/* - * 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.hive.s3select; - -import io.airlift.units.DataSize; -import io.trino.plugin.hive.HiveConfig; -import io.trino.spi.connector.ColumnHandle; -import io.trino.spi.connector.SchemaTableName; -import io.trino.testing.MaterializedResult; -import org.testng.annotations.BeforeClass; -import org.testng.annotations.DataProvider; -import org.testng.annotations.Parameters; -import org.testng.annotations.Test; - -import java.util.Optional; - -import static io.airlift.units.DataSize.Unit.KILOBYTE; -import static io.trino.plugin.hive.HiveColumnHandle.ColumnType.REGULAR; -import static io.trino.plugin.hive.HiveColumnHandle.createBaseColumn; -import static io.trino.plugin.hive.HiveFileSystemTestUtils.newSession; -import static io.trino.plugin.hive.HiveType.HIVE_INT; -import static io.trino.plugin.hive.s3select.S3SelectTestHelper.expectedResult; -import static io.trino.plugin.hive.s3select.S3SelectTestHelper.isSplitCountInOpenInterval; -import static io.trino.spi.type.BigintType.BIGINT; -import static io.trino.testing.QueryAssertions.assertEqualsIgnoreOrder; -import static org.testng.Assert.assertTrue; - -public class TestHiveFileSystemS3SelectCsvPushdownWithSplits -{ - private String host; - private int port; - private String databaseName; - private String awsAccessKey; - private String awsSecretKey; - private String writableBucket; - private String testDirectory; - - private SchemaTableName tableCsvWithSplits; - - @Parameters({ - "hive.hadoop2.metastoreHost", - "hive.hadoop2.metastorePort", - "hive.hadoop2.databaseName", - "hive.hadoop2.s3.awsAccessKey", - "hive.hadoop2.s3.awsSecretKey", - "hive.hadoop2.s3.writableBucket", - "hive.hadoop2.s3.testDirectory", - }) - @BeforeClass - public void setup(String host, int port, String databaseName, String awsAccessKey, String awsSecretKey, String writableBucket, String testDirectory) - { - this.host = host; - this.port = port; - this.databaseName = databaseName; - this.awsAccessKey = awsAccessKey; - this.awsSecretKey = awsSecretKey; - this.writableBucket = writableBucket; - this.testDirectory = testDirectory; - - tableCsvWithSplits = new SchemaTableName(databaseName, "trino_s3select_test_csv_scan_range_pushdown"); - } - - @DataProvider(name = "testSplitSize") - public static Object[][] splitSizeParametersProvider() - { - return new Object[][] {{3, 2, 15, 30}, {50, 30, 2, 4}}; - } - - @Test(dataProvider = "testSplitSize") - public void testQueryPushdownWithSplitSizeForCsv(int maxSplitSizeKB, - int maxInitialSplitSizeKB, - int minSplitCount, - int maxSplitCount) - { - S3SelectTestHelper s3SelectTestHelper = null; - try { - HiveConfig hiveConfig = new HiveConfig() - .setS3SelectPushdownEnabled(true) - .setMaxSplitSize(DataSize.of(maxSplitSizeKB, KILOBYTE)) - .setMaxInitialSplitSize(DataSize.of(maxInitialSplitSizeKB, KILOBYTE)); - s3SelectTestHelper = new S3SelectTestHelper( - host, - port, - databaseName, - awsAccessKey, - awsSecretKey, - writableBucket, - testDirectory, - hiveConfig); - - int tableSplitsCount = s3SelectTestHelper.getTableSplitsCount(tableCsvWithSplits); - assertTrue(isSplitCountInOpenInterval(tableSplitsCount, minSplitCount, maxSplitCount)); - - ColumnHandle indexColumn = createBaseColumn("index", 0, HIVE_INT, BIGINT, REGULAR, Optional.empty()); - MaterializedResult filteredTableResult = s3SelectTestHelper.getFilteredTableResult(tableCsvWithSplits, indexColumn); - assertEqualsIgnoreOrder(filteredTableResult, - expectedResult(newSession(s3SelectTestHelper.getHiveConfig()), 1, 300)); - } - finally { - if (s3SelectTestHelper != null) { - s3SelectTestHelper.tearDown(); - } - } - } -} diff --git a/plugin/trino-hive-hadoop2/src/test/java/io/trino/plugin/hive/s3select/TestHiveFileSystemS3SelectJsonPushdown.java b/plugin/trino-hive-hadoop2/src/test/java/io/trino/plugin/hive/s3select/TestHiveFileSystemS3SelectJsonPushdown.java deleted file mode 100644 index 260d03608d2c..000000000000 --- a/plugin/trino-hive-hadoop2/src/test/java/io/trino/plugin/hive/s3select/TestHiveFileSystemS3SelectJsonPushdown.java +++ /dev/null @@ -1,102 +0,0 @@ -/* - * 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.hive.s3select; - -import com.google.common.collect.ImmutableList; -import io.trino.spi.connector.ColumnHandle; -import io.trino.spi.connector.SchemaTableName; -import io.trino.testing.MaterializedResult; -import org.testng.annotations.AfterClass; -import org.testng.annotations.BeforeClass; -import org.testng.annotations.Parameters; -import org.testng.annotations.Test; - -import java.util.List; -import java.util.Optional; - -import static io.trino.plugin.hive.HiveColumnHandle.ColumnType.REGULAR; -import static io.trino.plugin.hive.HiveColumnHandle.createBaseColumn; -import static io.trino.plugin.hive.HiveFileSystemTestUtils.filterTable; -import static io.trino.plugin.hive.HiveFileSystemTestUtils.newSession; -import static io.trino.plugin.hive.HiveFileSystemTestUtils.readTable; -import static io.trino.plugin.hive.HiveType.HIVE_INT; -import static io.trino.spi.type.BigintType.BIGINT; -import static io.trino.testing.QueryAssertions.assertEqualsIgnoreOrder; - -public class TestHiveFileSystemS3SelectJsonPushdown -{ - private SchemaTableName tableJson; - - private S3SelectTestHelper s3SelectTestHelper; - - @Parameters({ - "hive.hadoop2.metastoreHost", - "hive.hadoop2.metastorePort", - "hive.hadoop2.databaseName", - "hive.hadoop2.s3.awsAccessKey", - "hive.hadoop2.s3.awsSecretKey", - "hive.hadoop2.s3.writableBucket", - "hive.hadoop2.s3.testDirectory", - }) - @BeforeClass - public void setup(String host, int port, String databaseName, String awsAccessKey, String awsSecretKey, String writableBucket, String testDirectory) - { - s3SelectTestHelper = new S3SelectTestHelper(host, port, databaseName, awsAccessKey, awsSecretKey, writableBucket, testDirectory); - tableJson = new SchemaTableName(databaseName, "trino_s3select_test_external_fs_json"); - } - - @Test - public void testGetRecordsJson() - throws Exception - { - assertEqualsIgnoreOrder( - readTable(tableJson, - s3SelectTestHelper.getTransactionManager(), - s3SelectTestHelper.getHiveConfig(), - s3SelectTestHelper.getPageSourceProvider(), - s3SelectTestHelper.getSplitManager()), - MaterializedResult.resultBuilder(newSession(s3SelectTestHelper.getHiveConfig()), BIGINT, BIGINT) - .row(2L, 4L).row(5L, 6L) // test_table.json - .row(7L, 23L).row(28L, 22L).row(13L, 10L) // test_table.json.gz - .row(1L, 19L).row(6L, 3L).row(24L, 22L).row(100L, 77L) // test_table.json.bz2 - .build()); - } - - @Test - public void testFilterRecordsJson() - throws Exception - { - List projectedColumns = ImmutableList.of( - createBaseColumn("col_1", 0, HIVE_INT, BIGINT, REGULAR, Optional.empty())); - - assertEqualsIgnoreOrder( - filterTable(tableJson, - projectedColumns, - s3SelectTestHelper.getTransactionManager(), - s3SelectTestHelper.getHiveConfig(), - s3SelectTestHelper.getPageSourceProvider(), - s3SelectTestHelper.getSplitManager()), - MaterializedResult.resultBuilder(newSession(s3SelectTestHelper.getHiveConfig()), BIGINT) - .row(2L).row(5L) // test_table.json - .row(7L).row(28L).row(13L) // test_table.json.gz - .row(1L).row(6L).row(24L).row(100L) // test_table.json.bz2 - .build()); - } - - @AfterClass(alwaysRun = true) - public void tearDown() - { - s3SelectTestHelper.tearDown(); - } -} diff --git a/plugin/trino-hive-hadoop2/src/test/java/io/trino/plugin/hive/s3select/TestHiveFileSystemS3SelectJsonPushdownWithSplits.java b/plugin/trino-hive-hadoop2/src/test/java/io/trino/plugin/hive/s3select/TestHiveFileSystemS3SelectJsonPushdownWithSplits.java deleted file mode 100644 index 1998ec9368da..000000000000 --- a/plugin/trino-hive-hadoop2/src/test/java/io/trino/plugin/hive/s3select/TestHiveFileSystemS3SelectJsonPushdownWithSplits.java +++ /dev/null @@ -1,116 +0,0 @@ -/* - * 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.hive.s3select; - -import io.airlift.units.DataSize; -import io.trino.plugin.hive.HiveConfig; -import io.trino.spi.connector.ColumnHandle; -import io.trino.spi.connector.SchemaTableName; -import io.trino.testing.MaterializedResult; -import org.testng.annotations.BeforeClass; -import org.testng.annotations.DataProvider; -import org.testng.annotations.Parameters; -import org.testng.annotations.Test; - -import java.util.Optional; - -import static io.airlift.units.DataSize.Unit.KILOBYTE; -import static io.trino.plugin.hive.HiveColumnHandle.ColumnType.REGULAR; -import static io.trino.plugin.hive.HiveColumnHandle.createBaseColumn; -import static io.trino.plugin.hive.HiveFileSystemTestUtils.newSession; -import static io.trino.plugin.hive.HiveType.HIVE_INT; -import static io.trino.plugin.hive.s3select.S3SelectTestHelper.expectedResult; -import static io.trino.plugin.hive.s3select.S3SelectTestHelper.isSplitCountInOpenInterval; -import static io.trino.spi.type.BigintType.BIGINT; -import static io.trino.testing.QueryAssertions.assertEqualsIgnoreOrder; -import static org.testng.Assert.assertTrue; - -public class TestHiveFileSystemS3SelectJsonPushdownWithSplits -{ - private String host; - private int port; - private String databaseName; - private String awsAccessKey; - private String awsSecretKey; - private String writableBucket; - private String testDirectory; - - private SchemaTableName tableJsonWithSplits; - - @Parameters({ - "hive.hadoop2.metastoreHost", - "hive.hadoop2.metastorePort", - "hive.hadoop2.databaseName", - "hive.hadoop2.s3.awsAccessKey", - "hive.hadoop2.s3.awsSecretKey", - "hive.hadoop2.s3.writableBucket", - "hive.hadoop2.s3.testDirectory", - }) - @BeforeClass - public void setup(String host, int port, String databaseName, String awsAccessKey, String awsSecretKey, String writableBucket, String testDirectory) - { - this.host = host; - this.port = port; - this.databaseName = databaseName; - this.awsAccessKey = awsAccessKey; - this.awsSecretKey = awsSecretKey; - this.writableBucket = writableBucket; - this.testDirectory = testDirectory; - - this.tableJsonWithSplits = new SchemaTableName(databaseName, "trino_s3select_test_json_scan_range_pushdown"); - } - - @DataProvider(name = "testSplitSize") - public static Object[][] splitSizeParametersProvider() - { - return new Object[][] {{15, 10, 6, 12}, {50, 30, 2, 4}}; - } - - @Test(dataProvider = "testSplitSize") - public void testQueryPushdownWithSplitSizeForJson(int maxSplitSizeKB, - int maxInitialSplitSizeKB, - int minSplitCount, - int maxSplitCount) - { - S3SelectTestHelper s3SelectTestHelper = null; - try { - HiveConfig hiveConfig = new HiveConfig() - .setS3SelectPushdownEnabled(true) - .setMaxSplitSize(DataSize.of(maxSplitSizeKB, KILOBYTE)) - .setMaxInitialSplitSize(DataSize.of(maxInitialSplitSizeKB, KILOBYTE)); - s3SelectTestHelper = new S3SelectTestHelper( - host, - port, - databaseName, - awsAccessKey, - awsSecretKey, - writableBucket, - testDirectory, - hiveConfig); - - int tableSplitsCount = s3SelectTestHelper.getTableSplitsCount(tableJsonWithSplits); - assertTrue(isSplitCountInOpenInterval(tableSplitsCount, minSplitCount, maxSplitCount)); - - ColumnHandle indexColumn = createBaseColumn("col_1", 0, HIVE_INT, BIGINT, REGULAR, Optional.empty()); - MaterializedResult filteredTableResult = s3SelectTestHelper.getFilteredTableResult(tableJsonWithSplits, indexColumn); - assertEqualsIgnoreOrder(filteredTableResult, - expectedResult(newSession(s3SelectTestHelper.getHiveConfig()), 1, 300)); - } - finally { - if (s3SelectTestHelper != null) { - s3SelectTestHelper.tearDown(); - } - } - } -} diff --git a/plugin/trino-hive-hadoop2/src/test/java/io/trino/plugin/hive/s3select/TestHiveFileSystemS3SelectPushdown.java b/plugin/trino-hive-hadoop2/src/test/java/io/trino/plugin/hive/s3select/TestHiveFileSystemS3SelectPushdown.java deleted file mode 100644 index eef3f86a1ad5..000000000000 --- a/plugin/trino-hive-hadoop2/src/test/java/io/trino/plugin/hive/s3select/TestHiveFileSystemS3SelectPushdown.java +++ /dev/null @@ -1,115 +0,0 @@ -/* - * 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.hive.s3select; - -import com.google.common.collect.ImmutableList; -import io.trino.plugin.hive.AbstractTestHiveFileSystemS3; -import io.trino.spi.connector.ColumnHandle; -import io.trino.spi.connector.SchemaTableName; -import io.trino.testing.MaterializedResult; -import org.testng.annotations.BeforeClass; -import org.testng.annotations.Parameters; -import org.testng.annotations.Test; - -import java.util.List; -import java.util.Optional; - -import static io.trino.plugin.hive.HiveColumnHandle.ColumnType.REGULAR; -import static io.trino.plugin.hive.HiveColumnHandle.createBaseColumn; -import static io.trino.plugin.hive.HiveType.HIVE_INT; -import static io.trino.spi.type.BigintType.BIGINT; -import static io.trino.testing.QueryAssertions.assertEqualsIgnoreOrder; - -public class TestHiveFileSystemS3SelectPushdown - extends AbstractTestHiveFileSystemS3 -{ - protected SchemaTableName tableWithPipeDelimiter; - protected SchemaTableName tableWithCommaDelimiter; - - @Parameters({ - "hive.hadoop2.metastoreHost", - "hive.hadoop2.metastorePort", - "hive.hadoop2.databaseName", - "hive.hadoop2.s3.endpoint", - "hive.hadoop2.s3.awsAccessKey", - "hive.hadoop2.s3.awsSecretKey", - "hive.hadoop2.s3.writableBucket", - "hive.hadoop2.s3.testDirectory", - }) - @BeforeClass - public void setup(String host, int port, String databaseName, String s3endpoint, String awsAccessKey, String awsSecretKey, String writableBucket, String testDirectory) - { - super.setup(host, port, databaseName, s3endpoint, awsAccessKey, awsSecretKey, writableBucket, testDirectory, true); - tableWithPipeDelimiter = new SchemaTableName(database, "trino_s3select_test_external_fs_with_pipe_delimiter"); - tableWithCommaDelimiter = new SchemaTableName(database, "trino_s3select_test_external_fs_with_comma_delimiter"); - } - - @Test - public void testGetRecordsWithPipeDelimiter() - throws Exception - { - assertEqualsIgnoreOrder( - readTable(tableWithPipeDelimiter), - MaterializedResult.resultBuilder(newSession(), BIGINT, BIGINT) - .row(1L, 2L).row(3L, 4L).row(55L, 66L) // test_table_with_pipe_delimiter.csv - .row(27L, 10L).row(8L, 2L).row(456L, 789L) // test_table_with_pipe_delimiter.csv.gzip - .row(22L, 11L).row(78L, 76L).row(1L, 2L).row(36L, 90L) // test_table_with_pipe_delimiter.csv.bz2 - .build()); - } - - @Test - public void testFilterRecordsWithPipeDelimiter() - throws Exception - { - List projectedColumns = ImmutableList.of( - createBaseColumn("t_bigint", 0, HIVE_INT, BIGINT, REGULAR, Optional.empty())); - - assertEqualsIgnoreOrder( - filterTable(tableWithPipeDelimiter, projectedColumns), - MaterializedResult.resultBuilder(newSession(), BIGINT) - .row(1L).row(3L).row(55L) // test_table_with_pipe_delimiter.csv - .row(27L).row(8L).row(456L) // test_table_with_pipe_delimiter.csv.gzip - .row(22L).row(78L).row(1L).row(36L) // test_table_with_pipe_delimiter.csv.bz2 - .build()); - } - - @Test - public void testGetRecordsWithCommaDelimiter() - throws Exception - { - assertEqualsIgnoreOrder( - readTable(tableWithCommaDelimiter), - MaterializedResult.resultBuilder(newSession(), BIGINT, BIGINT) - .row(7L, 1L).row(19L, 10L).row(1L, 345L) // test_table_with_comma_delimiter.csv - .row(27L, 10L).row(28L, 9L).row(90L, 94L) // test_table_with_comma_delimiter.csv.gzip - .row(11L, 24L).row(1L, 6L).row(21L, 12L).row(0L, 0L) // test_table_with_comma_delimiter.csv.bz2 - .build()); - } - - @Test - public void testFilterRecordsWithCommaDelimiter() - throws Exception - { - List projectedColumns = ImmutableList.of( - createBaseColumn("t_bigint", 0, HIVE_INT, BIGINT, REGULAR, Optional.empty())); - - assertEqualsIgnoreOrder( - filterTable(tableWithCommaDelimiter, projectedColumns), - MaterializedResult.resultBuilder(newSession(), BIGINT) - .row(7L).row(19L).row(1L) // test_table_with_comma_delimiter.csv - .row(27L).row(28L).row(90L) // test_table_with_comma_delimiter.csv.gzip - .row(11L).row(1L).row(21L).row(0L) // test_table_with_comma_delimiter.csv.bz2 - .build()); - } -} diff --git a/plugin/trino-hive/pom.xml b/plugin/trino-hive/pom.xml index ed6caaf034ac..bc4fe2e25c4e 100644 --- a/plugin/trino-hive/pom.xml +++ b/plugin/trino-hive/pom.xml @@ -29,6 +29,12 @@ com.amazonaws aws-java-sdk-core + + + org.apache.httpcomponents + httpclient + + @@ -36,16 +42,6 @@ aws-java-sdk-glue - - com.amazonaws - aws-java-sdk-s3 - - - - com.amazonaws - aws-java-sdk-sts - - com.fasterxml.jackson.core jackson-core @@ -76,11 +72,6 @@ failsafe - - io.airlift - aircompressor - - io.airlift bootstrap @@ -138,7 +129,7 @@ io.opentelemetry.instrumentation - opentelemetry-aws-sdk-1.11 + opentelemetry-aws-sdk-2.2 @@ -156,11 +147,6 @@ trino-filesystem-manager - - io.trino - trino-hdfs - - io.trino trino-hive-formats @@ -251,11 +237,6 @@ parquet-column - - org.apache.parquet - parquet-common - - org.apache.parquet parquet-format-structures @@ -281,6 +262,62 @@ jmxutils + + software.amazon.awssdk + apache-client + + + commons-logging + commons-logging + + + + + + software.amazon.awssdk + auth + + + + software.amazon.awssdk + aws-core + + + + software.amazon.awssdk + glue + + + + software.amazon.awssdk + http-client-spi + + + + software.amazon.awssdk + regions + + + + software.amazon.awssdk + retries-spi + + + + software.amazon.awssdk + sdk-core + + + + software.amazon.awssdk + sts + + + + software.amazon.awssdk + utils + + com.fasterxml.jackson.core jackson-annotations @@ -293,6 +330,12 @@ provided + + io.opentelemetry + opentelemetry-context + provided + + io.trino trino-spi @@ -317,12 +360,6 @@ runtime - - io.opentelemetry - opentelemetry-context - runtime - - io.trino trino-hadoop-toolkit @@ -335,6 +372,24 @@ runtime + + org.apache.parquet + parquet-common + runtime + + + + com.amazonaws + aws-java-sdk-s3 + test + + + + com.qubole.rubix + rubix-presto-shaded + test + + io.airlift testing @@ -402,6 +457,12 @@ test + + io.trino + trino-hdfs + test + + io.trino trino-main diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/BackgroundHiveSplitLoader.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/BackgroundHiveSplitLoader.java index a3f09bce745a..f3affd736e27 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/BackgroundHiveSplitLoader.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/BackgroundHiveSplitLoader.java @@ -18,18 +18,17 @@ import com.google.common.collect.ArrayListMultimap; import com.google.common.collect.ImmutableList; import com.google.common.collect.ListMultimap; +import com.google.common.collect.Multimaps; import com.google.common.collect.Streams; import com.google.common.io.CharStreams; import com.google.common.util.concurrent.Futures; import com.google.common.util.concurrent.ListenableFuture; import io.airlift.units.Duration; import io.trino.filesystem.FileEntry; +import io.trino.filesystem.FileIterator; import io.trino.filesystem.Location; import io.trino.filesystem.TrinoFileSystem; import io.trino.filesystem.TrinoFileSystemFactory; -import io.trino.hdfs.HdfsContext; -import io.trino.hdfs.HdfsEnvironment; -import io.trino.hdfs.HdfsNamenodeStats; import io.trino.plugin.hive.HiveSplit.BucketConversion; import io.trino.plugin.hive.HiveSplit.BucketValidation; import io.trino.plugin.hive.fs.DirectoryLister; @@ -39,7 +38,6 @@ import io.trino.plugin.hive.metastore.Partition; import io.trino.plugin.hive.metastore.StorageFormat; import io.trino.plugin.hive.metastore.Table; -import io.trino.plugin.hive.s3select.S3SelectPushdown; import io.trino.plugin.hive.util.AcidTables.AcidState; import io.trino.plugin.hive.util.AcidTables.ParsedDelta; import io.trino.plugin.hive.util.HiveBucketing.BucketingVersion; @@ -54,24 +52,15 @@ import io.trino.spi.connector.DynamicFilter; import io.trino.spi.predicate.TupleDomain; import io.trino.spi.type.TypeManager; -import org.apache.hadoop.conf.Configuration; import org.apache.hadoop.fs.FileStatus; import org.apache.hadoop.fs.FileSystem; import org.apache.hadoop.fs.Path; -import org.apache.hadoop.mapred.FileSplit; -import org.apache.hadoop.mapred.InputFormat; -import org.apache.hadoop.mapred.InputSplit; -import org.apache.hadoop.mapred.JobConf; -import org.apache.hadoop.mapred.JobConfigurable; -import org.apache.hadoop.mapred.TextInputFormat; -import org.apache.hadoop.mapreduce.MRConfig; -import org.apache.hadoop.util.StringUtils; import java.io.BufferedReader; import java.io.IOException; import java.io.InputStreamReader; +import java.io.Reader; import java.nio.charset.StandardCharsets; -import java.security.Principal; import java.util.ArrayList; import java.util.Arrays; import java.util.Deque; @@ -81,8 +70,6 @@ import java.util.Map; import java.util.Optional; import java.util.OptionalInt; -import java.util.Properties; -import java.util.Set; import java.util.concurrent.ConcurrentLinkedDeque; import java.util.concurrent.Executor; import java.util.concurrent.atomic.AtomicInteger; @@ -98,13 +85,10 @@ import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.base.Preconditions.checkState; -import static com.google.common.collect.ImmutableSet.toImmutableSet; -import static com.google.common.collect.Iterables.getOnlyElement; import static com.google.common.util.concurrent.Futures.immediateVoidFuture; import static com.google.common.util.concurrent.MoreExecutors.directExecutor; import static io.airlift.concurrent.MoreFutures.addExceptionCallback; import static io.airlift.concurrent.MoreFutures.toListenableFuture; -import static io.trino.hdfs.ConfigurationUtils.toJobConf; import static io.trino.plugin.hive.HiveErrorCode.HIVE_BAD_DATA; import static io.trino.plugin.hive.HiveErrorCode.HIVE_EXCEEDED_PARTITION_LIMIT; import static io.trino.plugin.hive.HiveErrorCode.HIVE_FILESYSTEM_ERROR; @@ -128,23 +112,23 @@ import static io.trino.plugin.hive.util.AcidTables.isFullAcidTable; import static io.trino.plugin.hive.util.AcidTables.isTransactionalTable; import static io.trino.plugin.hive.util.AcidTables.readAcidVersionFile; +import static io.trino.plugin.hive.util.HiveBucketing.getBucketingVersion; import static io.trino.plugin.hive.util.HiveClassNames.SYMLINK_TEXT_INPUT_FORMAT_CLASS; -import static io.trino.plugin.hive.util.HiveReaderUtil.getInputFormat; import static io.trino.plugin.hive.util.HiveUtil.checkCondition; -import static io.trino.plugin.hive.util.HiveUtil.getDeserializerClassName; import static io.trino.plugin.hive.util.HiveUtil.getFooterCount; import static io.trino.plugin.hive.util.HiveUtil.getHeaderCount; import static io.trino.plugin.hive.util.HiveUtil.getInputFormatName; import static io.trino.plugin.hive.util.HiveUtil.getPartitionKeyColumnHandles; +import static io.trino.plugin.hive.util.HiveUtil.getSerializationLibraryName; import static io.trino.plugin.hive.util.PartitionMatchSupplier.createPartitionMatchSupplier; import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED; import static java.lang.Integer.parseInt; import static java.lang.Math.max; import static java.lang.String.format; +import static java.nio.charset.StandardCharsets.UTF_8; import static java.util.Collections.max; import static java.util.Objects.requireNonNull; import static java.util.concurrent.TimeUnit.MILLISECONDS; -import static org.apache.hadoop.fs.Path.getPathWithoutSchemeAndAuthority; public class BackgroundHiveSplitLoader implements HiveSplitLoader @@ -170,15 +154,11 @@ public class BackgroundHiveSplitLoader private final long dynamicFilteringWaitTimeoutMillis; private final TypeManager typeManager; private final Optional tableBucketInfo; - private final HdfsEnvironment hdfsEnvironment; - private final HdfsContext hdfsContext; - private final HdfsNamenodeStats hdfsNamenodeStats; private final DirectoryLister directoryLister; private final TrinoFileSystemFactory fileSystemFactory; private final int loaderConcurrency; private final boolean recursiveDirWalkerEnabled; private final boolean ignoreAbsentPartitions; - private final boolean optimizeSymlinkListing; private final Executor executor; private final ConnectorSession session; private final ConcurrentLazyQueue partitions; @@ -220,14 +200,11 @@ public BackgroundHiveSplitLoader( Optional tableBucketInfo, ConnectorSession session, TrinoFileSystemFactory fileSystemFactory, - HdfsEnvironment hdfsEnvironment, - HdfsNamenodeStats hdfsNamenodeStats, DirectoryLister directoryLister, Executor executor, int loaderConcurrency, boolean recursiveDirWalkerEnabled, boolean ignoreAbsentPartitions, - boolean optimizeSymlinkListing, Optional validWriteIds, Optional maxSplitFileSize, int maxPartitions) @@ -242,18 +219,14 @@ public BackgroundHiveSplitLoader( checkArgument(loaderConcurrency > 0, "loaderConcurrency must be > 0, found: %s", loaderConcurrency); this.session = session; this.fileSystemFactory = requireNonNull(fileSystemFactory, "fileSystemFactory is null"); - this.hdfsEnvironment = hdfsEnvironment; - this.hdfsNamenodeStats = hdfsNamenodeStats; this.directoryLister = directoryLister; this.recursiveDirWalkerEnabled = recursiveDirWalkerEnabled; this.ignoreAbsentPartitions = ignoreAbsentPartitions; - this.optimizeSymlinkListing = optimizeSymlinkListing; requireNonNull(executor, "executor is null"); // direct executor is not supported in this implementation due to locking specifics checkExecutorIsNotDirectExecutor(executor); this.executor = executor; this.partitions = new ConcurrentLazyQueue<>(partitions); - this.hdfsContext = new HdfsContext(session); this.validWriteIds = requireNonNull(validWriteIds, "validWriteIds is null"); this.maxSplitFileSize = requireNonNull(maxSplitFileSize, "maxSplitFileSize is null"); this.maxPartitions = maxPartitions; @@ -422,7 +395,7 @@ private ListenableFuture loadPartition(HivePartitionMetadata partition) { HivePartition hivePartition = partition.getHivePartition(); String partitionName = hivePartition.getPartitionId(); - Properties schema = partition.getPartition() + Map schema = partition.getPartition() .map(value -> getHiveSchema(value, table)) .orElseGet(() -> getHiveSchema(table)); List partitionKeys = getPartitionKeys(table, partition.getPartition()); @@ -434,61 +407,37 @@ private ListenableFuture loadPartition(HivePartitionMetadata partition) return COMPLETED_FUTURE; } - Path path = new Path(getPartitionLocation(table, partition.getPartition())); - Configuration configuration = hdfsEnvironment.getConfiguration(hdfsContext, path); - FileSystem fs = hdfsEnvironment.getFileSystem(hdfsContext, path); - - boolean s3SelectPushdownEnabled = S3SelectPushdown.shouldEnablePushdownForTable(session, table, path.toString(), partition.getPartition()); - - // S3 Select pushdown works at the granularity of individual S3 objects for compressed files - // and finer granularity for uncompressed files using scan range feature. - boolean shouldEnableSplits = S3SelectPushdown.isSplittable(s3SelectPushdownEnabled, schema, path.toString()); + Location location = Location.of(getPartitionLocation(table, partition.getPartition())); // Skip header / footer lines are not splittable except for a special case when skip.header.line.count=1 - boolean splittable = shouldEnableSplits && getFooterCount(schema) == 0 && getHeaderCount(schema) <= 1; + boolean splittable = getFooterCount(schema) == 0 && getHeaderCount(schema) <= 1; if (SYMLINK_TEXT_INPUT_FORMAT_CLASS.equals(getInputFormatName(schema).orElse(null))) { if (tableBucketInfo.isPresent()) { throw new TrinoException(NOT_SUPPORTED, "Bucketed table in SymlinkTextInputFormat is not yet supported"); } - HiveStorageFormat targetStorageFormat = getSymlinkStorageFormat(getDeserializerClassName(schema)); - InputFormat targetInputFormat = getInputFormat(configuration, schema); - List targetPaths = hdfsEnvironment.doAs( - hdfsContext.getIdentity(), - () -> getTargetPathsFromSymlink(fs, path)); - Set parents = targetPaths.stream() - .map(Path::getParent) - .distinct() - .collect(toImmutableSet()); - if (optimizeSymlinkListing && parents.size() == 1 && !recursiveDirWalkerEnabled) { - Optional> manifestFileIterator = buildManifestFileIterator( - targetStorageFormat, - partitionName, - schema, - partitionKeys, - effectivePredicate, - partitionMatchSupplier, - s3SelectPushdownEnabled, - partition.getTableToPartitionMapping(), - getOnlyElement(parents), - targetPaths, - splittable); - if (manifestFileIterator.isPresent()) { - fileIterators.addLast(manifestFileIterator.get()); - return COMPLETED_FUTURE; - } - } - return createHiveSymlinkSplits( + HiveStorageFormat targetStorageFormat = getSymlinkStorageFormat(getSerializationLibraryName(schema)); + ListMultimap targets = getTargetLocationsByParentFromSymlink(location); + + InternalHiveSplitFactory splitFactory = new InternalHiveSplitFactory( partitionName, targetStorageFormat, - targetInputFormat, schema, partitionKeys, effectivePredicate, partitionMatchSupplier, - s3SelectPushdownEnabled, - partition.getTableToPartitionMapping(), - targetPaths); + partition.getHiveColumnCoercions(), + Optional.empty(), + Optional.empty(), + getMaxInitialSplitSize(session), + isForceLocalScheduling(session), + maxSplitFileSize); + + for (Map.Entry> entry : Multimaps.asMap(targets).entrySet()) { + fileIterators.addLast(buildManifestFileIterator(splitFactory, entry.getKey(), entry.getValue(), splittable)); + } + + return COMPLETED_FUTURE; } StorageFormat rawStorageFormat = partition.getPartition() @@ -502,7 +451,8 @@ private ListenableFuture loadPartition(HivePartitionMetadata partition) Optional partitionBucketProperty = partition.getPartition().get().getStorage().getBucketProperty(); if (tableBucketInfo.isPresent() && partitionBucketProperty.isPresent()) { int tableBucketCount = tableBucketInfo.get().getTableBucketCount(); - BucketingVersion bucketingVersion = partitionBucketProperty.get().getBucketingVersion(); // TODO can partition's bucketing_version be different from table's? + // Partition bucketing_version cannot be different from table + BucketingVersion bucketingVersion = getBucketingVersion(table.getParameters()); int partitionBucketCount = partitionBucketProperty.get().getBucketCount(); // Validation was done in HiveSplitManager#getPartitionMetadata. // Here, it's just trying to see if its needs the BucketConversion. @@ -522,35 +472,31 @@ private ListenableFuture loadPartition(HivePartitionMetadata partition) } InternalHiveSplitFactory splitFactory = new InternalHiveSplitFactory( - fs, partitionName, storageFormat, schema, partitionKeys, effectivePredicate, partitionMatchSupplier, - partition.getTableToPartitionMapping(), + partition.getHiveColumnCoercions(), bucketConversionRequiresWorkerParticipation ? bucketConversion : Optional.empty(), bucketValidation, getMaxInitialSplitSize(session), isForceLocalScheduling(session), - s3SelectPushdownEnabled, maxSplitFileSize); if (isTransactionalTable(table.getParameters())) { - return getTransactionalSplits(Location.of(path.toString()), splittable, bucketConversion, splitFactory); + return getTransactionalSplits(location, splittable, bucketConversion, splitFactory); } TrinoFileSystem trinoFileSystem = fileSystemFactory.create(session); - Location location = Location.of(path.toString()); // Bucketed partitions are fully loaded immediately since all files must be loaded to determine the file to bucket mapping if (tableBucketInfo.isPresent()) { List files = listBucketFiles(trinoFileSystem, location, splitFactory.getPartitionName()); return hiveSplitSource.addToQueue(getBucketedSplits(files, splitFactory, tableBucketInfo.get(), bucketConversion, splittable, Optional.empty())); } - Iterator splitIterator = createInternalHiveSplitIterator(trinoFileSystem, location, splitFactory, splittable, Optional.empty()); - fileIterators.addLast(splitIterator); + fileIterators.addLast(createInternalHiveSplitIterator(trinoFileSystem, location, splitFactory, splittable, Optional.empty())); return COMPLETED_FUTURE; } @@ -558,7 +504,7 @@ private ListenableFuture loadPartition(HivePartitionMetadata partition) private List listBucketFiles(TrinoFileSystem fs, Location location, String partitionName) { try { - HiveFileIterator fileIterator = new HiveFileIterator(table, location, fs, directoryLister, hdfsNamenodeStats, FAIL); + HiveFileIterator fileIterator = new HiveFileIterator(table, location, fs, directoryLister, FAIL); if (!fileIterator.hasNext() && !ignoreAbsentPartitions) { checkPartitionLocationExists(fs, location); } @@ -571,117 +517,47 @@ private List listBucketFiles(TrinoFileSystem fs, Location locat } } - private ListenableFuture createHiveSymlinkSplits( - String partitionName, - HiveStorageFormat storageFormat, - InputFormat targetInputFormat, - Properties schema, - List partitionKeys, - TupleDomain effectivePredicate, - BooleanSupplier partitionMatchSupplier, - boolean s3SelectPushdownEnabled, - TableToPartitionMapping tableToPartitionMapping, - List targetPaths) - throws IOException + @VisibleForTesting + Iterator buildManifestFileIterator(InternalHiveSplitFactory splitFactory, Location location, List paths, boolean splittable) { - ListenableFuture lastResult = COMPLETED_FUTURE; - for (Path targetPath : targetPaths) { - // the splits must be generated using the file system for the target path - // get the configuration for the target path -- it may be a different hdfs instance - FileSystem targetFilesystem = hdfsEnvironment.getFileSystem(hdfsContext, targetPath); - JobConf targetJob = toJobConf(targetFilesystem.getConf()); - targetJob.setInputFormat(TextInputFormat.class); - Optional principal = hdfsContext.getIdentity().getPrincipal(); - if (principal.isPresent()) { - targetJob.set(MRConfig.FRAMEWORK_NAME, MRConfig.CLASSIC_FRAMEWORK_NAME); - targetJob.set(MRConfig.MASTER_USER_NAME, principal.get().getName()); - } - if (targetInputFormat instanceof JobConfigurable) { - ((JobConfigurable) targetInputFormat).configure(targetJob); - } - targetJob.set(FILE_INPUT_FORMAT_INPUT_DIR, StringUtils.escapeString(targetPath.toString())); - InputSplit[] targetSplits = hdfsEnvironment.doAs( - hdfsContext.getIdentity(), - () -> targetInputFormat.getSplits(targetJob, 0)); - - InternalHiveSplitFactory splitFactory = new InternalHiveSplitFactory( - targetFilesystem, - partitionName, - storageFormat, - schema, - partitionKeys, - effectivePredicate, - partitionMatchSupplier, - tableToPartitionMapping, - Optional.empty(), - Optional.empty(), - getMaxInitialSplitSize(session), - isForceLocalScheduling(session), - s3SelectPushdownEnabled, - maxSplitFileSize); - lastResult = addSplitsToSource(targetSplits, splitFactory); - if (stopped) { - return COMPLETED_FUTURE; - } - } - return lastResult; + return createInternalHiveSplitIterator(splitFactory, splittable, Optional.empty(), verifiedFileStatusesStream(location, paths)); } - @VisibleForTesting - Optional> buildManifestFileIterator( - HiveStorageFormat targetStorageFormat, - String partitionName, - Properties schema, - List partitionKeys, - TupleDomain effectivePredicate, - BooleanSupplier partitionMatchSupplier, - boolean s3SelectPushdownEnabled, - TableToPartitionMapping tableToPartitionMapping, - Path parent, - List paths, - boolean splittable) - throws IOException + private Stream verifiedFileStatusesStream(Location location, List paths) { - FileSystem targetFilesystem = hdfsEnvironment.getFileSystem(hdfsContext, parent); TrinoFileSystem trinoFileSystem = fileSystemFactory.create(session); - Location location = Location.of(parent.toString()); + // Check if location is cached BEFORE using the directoryLister + boolean isCached = directoryLister.isCached(location); - Map fileStatuses = new HashMap<>(); - HiveFileIterator fileStatusIterator = new HiveFileIterator(table, location, trinoFileSystem, directoryLister, hdfsNamenodeStats, IGNORED); + Map fileStatuses = new HashMap<>(); + Iterator fileStatusIterator = new HiveFileIterator(table, location, trinoFileSystem, directoryLister, RECURSE); if (!fileStatusIterator.hasNext()) { checkPartitionLocationExists(trinoFileSystem, location); } - fileStatusIterator.forEachRemaining(status -> fileStatuses.put(getPathWithoutSchemeAndAuthority(new Path(status.getPath())), status)); - - List locatedFileStatuses = new ArrayList<>(); - for (Path path : paths) { - TrinoFileStatus status = fileStatuses.get(getPathWithoutSchemeAndAuthority(path)); - // This check will catch all directories in the manifest since HiveFileIterator will not return any directories. - // Some files may not be listed by HiveFileIterator - if those are included in the manifest this check will fail as well. - if (status == null) { - return Optional.empty(); - } + fileStatusIterator.forEachRemaining(status -> fileStatuses.put(Location.of(status.getPath()).path(), status)); - locatedFileStatuses.add(status); - } + // If file statuses came from cache verify that all are present + if (isCached) { + boolean missing = paths.stream() + .anyMatch(path -> !fileStatuses.containsKey(path.path())); + // Invalidate the cache and reload + if (missing) { + directoryLister.invalidate(location); - InternalHiveSplitFactory splitFactory = new InternalHiveSplitFactory( - targetFilesystem, - partitionName, - targetStorageFormat, - schema, - partitionKeys, - effectivePredicate, - partitionMatchSupplier, - tableToPartitionMapping, - Optional.empty(), - Optional.empty(), - getMaxInitialSplitSize(session), - isForceLocalScheduling(session), - s3SelectPushdownEnabled, - maxSplitFileSize); + fileStatuses.clear(); + fileStatusIterator = new HiveFileIterator(table, location, trinoFileSystem, directoryLister, RECURSE); + fileStatusIterator.forEachRemaining(status -> fileStatuses.put(Location.of(status.getPath()).path(), status)); + } + } - return Optional.of(createInternalHiveSplitIterator(splitFactory, splittable, Optional.empty(), locatedFileStatuses.stream())); + return paths.stream() + .map(path -> { + TrinoFileStatus status = fileStatuses.get(path.path()); + if (status == null) { + throw new TrinoException(HIVE_FILE_NOT_FOUND, "Manifest file from the location [%s] contains non-existent path: %s".formatted(location, path)); + } + return status; + }); } private ListenableFuture getTransactionalSplits(Location path, boolean splittable, Optional bucketConversion, InternalHiveSplitFactory splitFactory) @@ -785,25 +661,9 @@ private static Optional acidInfoForOriginalFiles(boolean fullAcid, Aci return fullAcid ? Optional.of(builder.buildWithRequiredOriginalFiles(getRequiredBucketNumber(location))) : Optional.empty(); } - private ListenableFuture addSplitsToSource(InputSplit[] targetSplits, InternalHiveSplitFactory splitFactory) - throws IOException - { - ListenableFuture lastResult = COMPLETED_FUTURE; - for (InputSplit inputSplit : targetSplits) { - Optional internalHiveSplit = splitFactory.createInternalHiveSplit((FileSplit) inputSplit); - if (internalHiveSplit.isPresent()) { - lastResult = hiveSplitSource.addToQueue(internalHiveSplit.get()); - } - if (stopped) { - return COMPLETED_FUTURE; - } - } - return lastResult; - } - private Iterator createInternalHiveSplitIterator(TrinoFileSystem fileSystem, Location location, InternalHiveSplitFactory splitFactory, boolean splittable, Optional acidInfo) { - Iterator iterator = new HiveFileIterator(table, location, fileSystem, directoryLister, hdfsNamenodeStats, recursiveDirWalkerEnabled ? RECURSE : IGNORED); + Iterator iterator = new HiveFileIterator(table, location, fileSystem, directoryLister, recursiveDirWalkerEnabled ? RECURSE : IGNORED); if (!iterator.hasNext() && !ignoreAbsentPartitions) { checkPartitionLocationExists(fileSystem, location); } @@ -1007,6 +867,32 @@ private static List getTargetPathsFromSymlink(FileSystem fileSystem, Path } } + private ListMultimap getTargetLocationsByParentFromSymlink(Location symlinkDir) + { + TrinoFileSystem fileSystem = fileSystemFactory.create(session); + try { + ListMultimap targets = ArrayListMultimap.create(); + FileIterator iterator = fileSystem.listFiles(symlinkDir); + while (iterator.hasNext()) { + Location location = iterator.next().location(); + String name = location.fileName(); + if (name.startsWith("_") || name.startsWith(".")) { + continue; + } + + try (Reader reader = new InputStreamReader(fileSystem.newInputFile(location).newStream(), UTF_8)) { + CharStreams.readLines(reader).stream() + .map(Location::of) + .forEach(target -> targets.put(target.parentDirectory(), target)); + } + } + return targets; + } + catch (IOException | IllegalArgumentException e) { + throw new TrinoException(HIVE_BAD_DATA, "Error parsing symlinks from: " + symlinkDir, e); + } + } + private static List getPartitionKeys(Table table, Optional partition) { if (partition.isEmpty()) { @@ -1037,21 +923,21 @@ public static class BucketSplitInfo private final int readBucketCount; private final IntPredicate bucketFilter; - public static Optional createBucketSplitInfo(Optional bucketHandle, Optional bucketFilter) + public static Optional createBucketSplitInfo(Optional tablePartitioning, Optional bucketFilter) { - requireNonNull(bucketHandle, "bucketHandle is null"); + requireNonNull(tablePartitioning, "bucketHandle is null"); requireNonNull(bucketFilter, "bucketFilter is null"); - if (bucketHandle.isEmpty()) { - checkArgument(bucketFilter.isEmpty(), "bucketHandle must be present if bucketFilter is present"); + if (tablePartitioning.isEmpty()) { + checkArgument(bucketFilter.isEmpty(), "tablePartitioning must be present if bucketFilter is present"); return Optional.empty(); } - BucketingVersion bucketingVersion = bucketHandle.get().getBucketingVersion(); - int tableBucketCount = bucketHandle.get().getTableBucketCount(); - int readBucketCount = bucketHandle.get().getReadBucketCount(); + BucketingVersion bucketingVersion = tablePartitioning.get().partitioningHandle().getBucketingVersion(); + int tableBucketCount = tablePartitioning.get().tableBucketCount(); + int readBucketCount = tablePartitioning.get().partitioningHandle().getBucketCount(); - List bucketColumns = bucketHandle.get().getColumns(); + List bucketColumns = tablePartitioning.get().columns(); IntPredicate predicate = bucketFilter .map(filter -> filter.getBucketsToKeep()::contains) .orElse(bucket -> true); diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/GenericHiveRecordCursor.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/GenericHiveRecordCursor.java deleted file mode 100644 index 70dbebd7ca49..000000000000 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/GenericHiveRecordCursor.java +++ /dev/null @@ -1,605 +0,0 @@ -/* - * 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.hive; - -import io.airlift.slice.Slice; -import io.airlift.slice.Slices; -import io.trino.hadoop.TextLineLengthLimitExceededException; -import io.trino.plugin.base.type.DecodedTimestamp; -import io.trino.plugin.base.type.TrinoTimestampEncoder; -import io.trino.spi.TrinoException; -import io.trino.spi.block.Block; -import io.trino.spi.connector.RecordCursor; -import io.trino.spi.type.CharType; -import io.trino.spi.type.DecimalType; -import io.trino.spi.type.Int128; -import io.trino.spi.type.LongTimestamp; -import io.trino.spi.type.TimestampType; -import io.trino.spi.type.Type; -import io.trino.spi.type.VarcharType; -import org.apache.hadoop.conf.Configuration; -import org.apache.hadoop.fs.Path; -import org.apache.hadoop.hive.common.type.Date; -import org.apache.hadoop.hive.common.type.HiveChar; -import org.apache.hadoop.hive.common.type.HiveDecimal; -import org.apache.hadoop.hive.common.type.HiveVarchar; -import org.apache.hadoop.hive.common.type.Timestamp; -import org.apache.hadoop.hive.serde2.Deserializer; -import org.apache.hadoop.hive.serde2.SerDeException; -import org.apache.hadoop.hive.serde2.io.HiveCharWritable; -import org.apache.hadoop.hive.serde2.io.HiveVarcharWritable; -import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspector; -import org.apache.hadoop.hive.serde2.objectinspector.PrimitiveObjectInspector; -import org.apache.hadoop.hive.serde2.objectinspector.StructField; -import org.apache.hadoop.hive.serde2.objectinspector.StructObjectInspector; -import org.apache.hadoop.io.BinaryComparable; -import org.apache.hadoop.io.BytesWritable; -import org.apache.hadoop.io.Text; -import org.apache.hadoop.io.Writable; -import org.apache.hadoop.mapred.RecordReader; - -import java.io.IOException; -import java.io.UncheckedIOException; -import java.math.BigInteger; -import java.util.Arrays; -import java.util.List; -import java.util.Properties; - -import static com.google.common.base.Preconditions.checkArgument; -import static com.google.common.base.Preconditions.checkState; -import static io.trino.plugin.base.type.TrinoTimestampEncoderFactory.createTimestampEncoder; -import static io.trino.plugin.base.util.Closables.closeAllSuppress; -import static io.trino.plugin.hive.HiveColumnHandle.ColumnType.REGULAR; -import static io.trino.plugin.hive.HiveErrorCode.HIVE_BAD_DATA; -import static io.trino.plugin.hive.HiveErrorCode.HIVE_CURSOR_ERROR; -import static io.trino.plugin.hive.util.HiveReaderUtil.getDeserializer; -import static io.trino.plugin.hive.util.HiveReaderUtil.getTableObjectInspector; -import static io.trino.plugin.hive.util.HiveUtil.isStructuralType; -import static io.trino.plugin.hive.util.SerDeUtils.getBlockObject; -import static io.trino.spi.type.BigintType.BIGINT; -import static io.trino.spi.type.BooleanType.BOOLEAN; -import static io.trino.spi.type.Chars.truncateToLengthAndTrimSpaces; -import static io.trino.spi.type.DateType.DATE; -import static io.trino.spi.type.Decimals.rescale; -import static io.trino.spi.type.DoubleType.DOUBLE; -import static io.trino.spi.type.IntegerType.INTEGER; -import static io.trino.spi.type.RealType.REAL; -import static io.trino.spi.type.SmallintType.SMALLINT; -import static io.trino.spi.type.TinyintType.TINYINT; -import static io.trino.spi.type.VarbinaryType.VARBINARY; -import static io.trino.spi.type.Varchars.truncateToLength; -import static java.lang.Float.floatToRawIntBits; -import static java.lang.Math.max; -import static java.lang.Math.min; -import static java.lang.String.format; -import static java.util.Objects.requireNonNull; -import static org.joda.time.DateTimeZone.UTC; - -public class GenericHiveRecordCursor - implements RecordCursor -{ - private final Path path; - private final RecordReader recordReader; - private final K key; - private final V value; - - private final Deserializer deserializer; - - private final Type[] types; - private final HiveType[] hiveTypes; - - private final StructObjectInspector rowInspector; - private final ObjectInspector[] fieldInspectors; - private final StructField[] structFields; - - private final boolean[] loaded; - private final boolean[] booleans; - private final long[] longs; - private final double[] doubles; - private final Slice[] slices; - private final Object[] objects; - private final boolean[] nulls; - private final TrinoTimestampEncoder[] timestampEncoders; - - private final long totalBytes; - - private long completedBytes; - private Object rowData; - private boolean closed; - - public GenericHiveRecordCursor( - Configuration configuration, - Path path, - RecordReader recordReader, - long totalBytes, - Properties splitSchema, - List columns) - { - requireNonNull(path, "path is null"); - requireNonNull(recordReader, "recordReader is null"); - checkArgument(totalBytes >= 0, "totalBytes is negative"); - requireNonNull(splitSchema, "splitSchema is null"); - requireNonNull(columns, "columns is null"); - - this.path = path; - this.recordReader = recordReader; - this.totalBytes = totalBytes; - this.key = recordReader.createKey(); - this.value = recordReader.createValue(); - - this.deserializer = getDeserializer(configuration, splitSchema); - this.rowInspector = getTableObjectInspector(deserializer); - - int size = columns.size(); - - this.types = new Type[size]; - this.hiveTypes = new HiveType[size]; - - this.structFields = new StructField[size]; - this.fieldInspectors = new ObjectInspector[size]; - - this.loaded = new boolean[size]; - this.booleans = new boolean[size]; - this.longs = new long[size]; - this.doubles = new double[size]; - this.slices = new Slice[size]; - this.objects = new Object[size]; - this.nulls = new boolean[size]; - this.timestampEncoders = new TrinoTimestampEncoder[size]; - - // initialize data columns - for (int i = 0; i < columns.size(); i++) { - HiveColumnHandle column = columns.get(i); - checkState(column.getColumnType() == REGULAR, "column type must be regular"); - - Type columnType = column.getType(); - types[i] = columnType; - if (columnType instanceof TimestampType) { - timestampEncoders[i] = createTimestampEncoder((TimestampType) columnType, UTC); - } - hiveTypes[i] = column.getHiveType(); - - StructField field = rowInspector.getStructFieldRef(column.getName()); - structFields[i] = field; - fieldInspectors[i] = field.getFieldObjectInspector(); - } - } - - @Override - public long getCompletedBytes() - { - if (!closed) { - updateCompletedBytes(); - } - return completedBytes; - } - - @Override - public long getReadTimeNanos() - { - return 0; - } - - private void updateCompletedBytes() - { - try { - @SuppressWarnings("NumericCastThatLosesPrecision") - long newCompletedBytes = (long) (totalBytes * recordReader.getProgress()); - completedBytes = min(totalBytes, max(completedBytes, newCompletedBytes)); - } - catch (IOException ignored) { - } - } - - @Override - public Type getType(int field) - { - return types[field]; - } - - @Override - public boolean advanceNextPosition() - { - try { - if (closed || !recordReader.next(key, value)) { - close(); - return false; - } - - // Only deserialize the value if atleast one column is required - if (types.length > 0) { - // reset loaded flags - Arrays.fill(loaded, false); - - // decode value - rowData = deserializer.deserialize(value); - } - - return true; - } - catch (IOException | SerDeException | RuntimeException e) { - closeAllSuppress(e, this); - if (e instanceof TextLineLengthLimitExceededException) { - throw new TrinoException(HIVE_BAD_DATA, "Line too long in text file: " + path, e); - } - throw new TrinoException(HIVE_CURSOR_ERROR, e); - } - } - - @Override - public boolean getBoolean(int fieldId) - { - checkState(!closed, "Cursor is closed"); - - validateType(fieldId, boolean.class); - if (!loaded[fieldId]) { - parseBooleanColumn(fieldId); - } - return booleans[fieldId]; - } - - private void parseBooleanColumn(int column) - { - loaded[column] = true; - - Object fieldData = rowInspector.getStructFieldData(rowData, structFields[column]); - - if (fieldData == null) { - nulls[column] = true; - } - else { - Object fieldValue = ((PrimitiveObjectInspector) fieldInspectors[column]).getPrimitiveJavaObject(fieldData); - checkState(fieldValue != null, "fieldValue should not be null"); - booleans[column] = (Boolean) fieldValue; - nulls[column] = false; - } - } - - @Override - public long getLong(int fieldId) - { - checkState(!closed, "Cursor is closed"); - - validateType(fieldId, long.class); - if (!loaded[fieldId]) { - parseLongColumn(fieldId); - } - return longs[fieldId]; - } - - private void parseLongColumn(int column) - { - loaded[column] = true; - - Object fieldData = rowInspector.getStructFieldData(rowData, structFields[column]); - - if (fieldData == null) { - nulls[column] = true; - } - else { - Object fieldValue = ((PrimitiveObjectInspector) fieldInspectors[column]).getPrimitiveJavaObject(fieldData); - checkState(fieldValue != null, "fieldValue should not be null"); - longs[column] = getLongExpressedValue(fieldValue, column); - nulls[column] = false; - } - } - - private long getLongExpressedValue(Object value, int column) - { - if (value instanceof Date) { - return ((Date) value).toEpochDay(); - } - if (value instanceof Timestamp) { - return shortTimestamp((Timestamp) value, column); - } - if (value instanceof Float) { - return floatToRawIntBits(((Float) value)); - } - return ((Number) value).longValue(); - } - - @Override - public double getDouble(int fieldId) - { - checkState(!closed, "Cursor is closed"); - - validateType(fieldId, double.class); - if (!loaded[fieldId]) { - parseDoubleColumn(fieldId); - } - return doubles[fieldId]; - } - - private void parseDoubleColumn(int column) - { - loaded[column] = true; - - Object fieldData = rowInspector.getStructFieldData(rowData, structFields[column]); - - if (fieldData == null) { - nulls[column] = true; - } - else { - Object fieldValue = ((PrimitiveObjectInspector) fieldInspectors[column]).getPrimitiveJavaObject(fieldData); - checkState(fieldValue != null, "fieldValue should not be null"); - doubles[column] = ((Number) fieldValue).doubleValue(); - nulls[column] = false; - } - } - - @Override - public Slice getSlice(int fieldId) - { - checkState(!closed, "Cursor is closed"); - - validateType(fieldId, Slice.class); - if (!loaded[fieldId]) { - parseStringColumn(fieldId); - } - return slices[fieldId]; - } - - private void parseStringColumn(int column) - { - loaded[column] = true; - - Object fieldData = rowInspector.getStructFieldData(rowData, structFields[column]); - - if (fieldData == null) { - nulls[column] = true; - } - else { - PrimitiveObjectInspector inspector = (PrimitiveObjectInspector) fieldInspectors[column]; - Slice value; - if (inspector.preferWritable()) { - value = parseStringFromPrimitiveWritableObjectValue(types[column], inspector.getPrimitiveWritableObject(fieldData)); - } - else { - value = parseStringFromPrimitiveJavaObjectValue(types[column], inspector.getPrimitiveJavaObject(fieldData)); - } - slices[column] = value; - nulls[column] = false; - } - } - - private static Slice trimStringToCharacterLimits(Type type, Slice value) - { - if (type instanceof VarcharType) { - return truncateToLength(value, type); - } - if (type instanceof CharType) { - return truncateToLengthAndTrimSpaces(value, type); - } - return value; - } - - private static Slice parseStringFromPrimitiveWritableObjectValue(Type type, Object fieldValue) - { - checkState(fieldValue != null, "fieldValue should not be null"); - BinaryComparable hiveValue; - if (fieldValue instanceof Text) { - hiveValue = (Text) fieldValue; - } - else if (fieldValue instanceof BytesWritable) { - hiveValue = (BytesWritable) fieldValue; - } - else if (fieldValue instanceof HiveVarcharWritable) { - hiveValue = ((HiveVarcharWritable) fieldValue).getTextValue(); - } - else if (fieldValue instanceof HiveCharWritable) { - hiveValue = ((HiveCharWritable) fieldValue).getTextValue(); - } - else { - throw new IllegalStateException("unsupported string field type: " + fieldValue.getClass().getName()); - } - // create a slice view over the hive value and trim to character limits - Slice value = trimStringToCharacterLimits(type, Slices.wrappedBuffer(hiveValue.getBytes(), 0, hiveValue.getLength())); - // store a copy of the bytes, since the hive reader can reuse the underlying buffer - return Slices.copyOf(value); - } - - private static Slice parseStringFromPrimitiveJavaObjectValue(Type type, Object fieldValue) - { - checkState(fieldValue != null, "fieldValue should not be null"); - Slice value; - if (fieldValue instanceof String) { - value = Slices.utf8Slice((String) fieldValue); - } - else if (fieldValue instanceof byte[]) { - value = Slices.wrappedBuffer((byte[]) fieldValue); - } - else if (fieldValue instanceof HiveVarchar) { - value = Slices.utf8Slice(((HiveVarchar) fieldValue).getValue()); - } - else if (fieldValue instanceof HiveChar) { - value = Slices.utf8Slice(((HiveChar) fieldValue).getValue()); - } - else { - throw new IllegalStateException("unsupported string field type: " + fieldValue.getClass().getName()); - } - value = trimStringToCharacterLimits(type, value); - // Copy the slice if the value was trimmed and is now smaller than the backing buffer - if (!value.isCompact()) { - return Slices.copyOf(value); - } - return value; - } - - private void parseDecimalColumn(int column) - { - loaded[column] = true; - - Object fieldData = rowInspector.getStructFieldData(rowData, structFields[column]); - - if (fieldData == null) { - nulls[column] = true; - } - else { - Object fieldValue = ((PrimitiveObjectInspector) fieldInspectors[column]).getPrimitiveJavaObject(fieldData); - checkState(fieldValue != null, "fieldValue should not be null"); - - HiveDecimal decimal = (HiveDecimal) fieldValue; - DecimalType columnType = (DecimalType) types[column]; - BigInteger unscaledDecimal = rescale(decimal.unscaledValue(), decimal.scale(), columnType.getScale()); - - if (columnType.isShort()) { - longs[column] = unscaledDecimal.longValue(); - } - else { - objects[column] = Int128.valueOf(unscaledDecimal); - } - nulls[column] = false; - } - } - - @Override - public Object getObject(int fieldId) - { - checkState(!closed, "Cursor is closed"); - - if (!loaded[fieldId]) { - parseObjectColumn(fieldId); - } - return objects[fieldId]; - } - - private void parseObjectColumn(int column) - { - loaded[column] = true; - - Object fieldData = rowInspector.getStructFieldData(rowData, structFields[column]); - - if (fieldData == null) { - nulls[column] = true; - } - else { - Type type = types[column]; - if (type.getJavaType() == Block.class) { - objects[column] = getBlockObject(type, fieldData, fieldInspectors[column]); - } - else if (type instanceof TimestampType) { - Timestamp timestamp = (Timestamp) ((PrimitiveObjectInspector) fieldInspectors[column]).getPrimitiveJavaObject(fieldData); - objects[column] = longTimestamp(timestamp, column); - } - else { - throw new IllegalStateException("Unsupported type: " + type); - } - nulls[column] = false; - } - } - - @Override - public boolean isNull(int fieldId) - { - checkState(!closed, "Cursor is closed"); - - if (!loaded[fieldId]) { - parseColumn(fieldId); - } - return nulls[fieldId]; - } - - private void parseColumn(int column) - { - Type type = types[column]; - if (BOOLEAN.equals(type)) { - parseBooleanColumn(column); - } - else if (BIGINT.equals(type)) { - parseLongColumn(column); - } - else if (INTEGER.equals(type)) { - parseLongColumn(column); - } - else if (SMALLINT.equals(type)) { - parseLongColumn(column); - } - else if (TINYINT.equals(type)) { - parseLongColumn(column); - } - else if (REAL.equals(type)) { - parseLongColumn(column); - } - else if (DOUBLE.equals(type)) { - parseDoubleColumn(column); - } - else if (type instanceof VarcharType || VARBINARY.equals(type)) { - parseStringColumn(column); - } - else if (type instanceof CharType) { - parseStringColumn(column); - } - else if (isStructuralType(type)) { - parseObjectColumn(column); - } - else if (DATE.equals(type)) { - parseLongColumn(column); - } - else if (type instanceof TimestampType) { - if (((TimestampType) type).isShort()) { - parseLongColumn(column); - } - else { - parseObjectColumn(column); - } - } - else if (type instanceof DecimalType) { - parseDecimalColumn(column); - } - else { - throw new UnsupportedOperationException("Unsupported column type: " + type); - } - } - - private void validateType(int fieldId, Class type) - { - if (!types[fieldId].getJavaType().equals(type)) { - // we don't use Preconditions.checkArgument because it requires boxing fieldId, which affects inner loop performance - throw new IllegalArgumentException(format("Expected field to be %s, actual %s (field %s)", type, types[fieldId], fieldId)); - } - } - - @Override - public void close() - { - // some hive input formats are broken and bad things can happen if you close them multiple times - if (closed) { - return; - } - closed = true; - - updateCompletedBytes(); - - try { - recordReader.close(); - } - catch (IOException e) { - throw new UncheckedIOException(e); - } - } - - private long shortTimestamp(Timestamp value, int column) - { - @SuppressWarnings("unchecked") - TrinoTimestampEncoder encoder = (TrinoTimestampEncoder) timestampEncoders[column]; - return encoder.getTimestamp(new DecodedTimestamp(value.toEpochSecond(), value.getNanos())); - } - - private LongTimestamp longTimestamp(Timestamp value, int column) - { - @SuppressWarnings("unchecked") - TrinoTimestampEncoder encoder = (TrinoTimestampEncoder) timestampEncoders[column]; - return encoder.getTimestamp(new DecodedTimestamp(value.toEpochSecond(), value.getNanos())); - } -} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/GenericHiveRecordCursorProvider.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/GenericHiveRecordCursorProvider.java deleted file mode 100644 index e809bb83278a..000000000000 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/GenericHiveRecordCursorProvider.java +++ /dev/null @@ -1,135 +0,0 @@ -/* - * 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.hive; - -import com.google.inject.Inject; -import io.airlift.units.DataSize; -import io.trino.filesystem.Location; -import io.trino.hdfs.HdfsEnvironment; -import io.trino.spi.TrinoException; -import io.trino.spi.connector.ConnectorSession; -import io.trino.spi.connector.RecordCursor; -import io.trino.spi.predicate.TupleDomain; -import io.trino.spi.type.TypeManager; -import org.apache.hadoop.conf.Configuration; -import org.apache.hadoop.fs.Path; -import org.apache.hadoop.io.Writable; -import org.apache.hadoop.mapred.RecordReader; -import org.apache.hadoop.mapreduce.lib.input.LineRecordReader; - -import java.io.IOException; -import java.util.List; -import java.util.Optional; -import java.util.Properties; - -import static com.google.common.base.Preconditions.checkArgument; -import static io.trino.plugin.hive.HiveErrorCode.HIVE_FILESYSTEM_ERROR; -import static io.trino.plugin.hive.HivePageSourceProvider.projectBaseColumns; -import static io.trino.plugin.hive.util.HiveReaderUtil.createRecordReader; -import static java.lang.Math.toIntExact; -import static java.util.Objects.requireNonNull; -import static java.util.stream.Collectors.toUnmodifiableList; - -public class GenericHiveRecordCursorProvider - implements HiveRecordCursorProvider -{ - private final HdfsEnvironment hdfsEnvironment; - private final int textMaxLineLengthBytes; - - @Inject - public GenericHiveRecordCursorProvider(HdfsEnvironment hdfsEnvironment, HiveConfig config) - { - this(hdfsEnvironment, config.getTextMaxLineLength()); - } - - public GenericHiveRecordCursorProvider(HdfsEnvironment hdfsEnvironment, DataSize textMaxLineLength) - { - this.hdfsEnvironment = requireNonNull(hdfsEnvironment, "hdfsEnvironment is null"); - this.textMaxLineLengthBytes = toIntExact(textMaxLineLength.toBytes()); - checkArgument(textMaxLineLengthBytes >= 1, "textMaxLineLength must be at least 1 byte"); - } - - @Override - public Optional createRecordCursor( - Configuration configuration, - ConnectorSession session, - Location location, - long start, - long length, - long fileSize, - Properties schema, - List columns, - TupleDomain effectivePredicate, - TypeManager typeManager, - boolean s3SelectPushdownEnabled) - { - configuration.setInt(LineRecordReader.MAX_LINE_LENGTH, textMaxLineLengthBytes); - - // make sure the FileSystem is created with the proper Configuration object - Path path = new Path(location.toString()); - try { - this.hdfsEnvironment.getFileSystem(session.getIdentity(), path, configuration); - } - catch (IOException e) { - throw new TrinoException(HIVE_FILESYSTEM_ERROR, "Failed getting FileSystem: " + path, e); - } - - Optional projections = projectBaseColumns(columns); - List readerColumns = projections - .map(ReaderColumns::get) - .map(columnHandles -> columnHandles.stream() - .map(HiveColumnHandle.class::cast) - .collect(toUnmodifiableList())) - .orElse(columns); - - RecordCursor cursor = hdfsEnvironment.doAs(session.getIdentity(), () -> { - RecordReader recordReader = createRecordReader( - configuration, - path, - start, - length, - schema, - readerColumns); - - try { - return new GenericHiveRecordCursor<>( - configuration, - path, - genericRecordReader(recordReader), - length, - schema, - readerColumns); - } - catch (Exception e) { - try { - recordReader.close(); - } - catch (IOException closeException) { - if (e != closeException) { - e.addSuppressed(closeException); - } - } - throw e; - } - }); - - return Optional.of(new ReaderRecordCursorWithProjections(cursor, projections)); - } - - @SuppressWarnings("unchecked") - private static RecordReader genericRecordReader(RecordReader recordReader) - { - return (RecordReader) recordReader; - } -} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveBasicStatistics.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveBasicStatistics.java index 0ae16566732c..d83bf1bacc58 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveBasicStatistics.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveBasicStatistics.java @@ -83,6 +83,11 @@ public OptionalLong getOnDiskDataSizeInBytes() return onDiskDataSizeInBytes; } + public HiveBasicStatistics withEmptyRowCount() + { + return new HiveBasicStatistics(fileCount, OptionalLong.empty(), inMemoryDataSizeInBytes, onDiskDataSizeInBytes); + } + @Override public boolean equals(Object o) { diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveBucketHandle.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveBucketHandle.java index a787569cc0bd..c2b953664669 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveBucketHandle.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveBucketHandle.java @@ -26,7 +26,6 @@ import static com.google.common.base.Preconditions.checkArgument; import static java.lang.String.format; import static java.util.Objects.requireNonNull; -import static java.util.stream.Collectors.toList; public class HiveBucketHandle { @@ -85,17 +84,6 @@ public List getSortedBy() return sortedBy; } - public HiveBucketProperty toTableBucketProperty() - { - return new HiveBucketProperty( - columns.stream() - .map(HiveColumnHandle::getName) - .collect(toList()), - bucketingVersion, - tableBucketCount, - sortedBy); - } - @Override public boolean equals(Object obj) { diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveBucketProperty.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveBucketProperty.java index 429492b85ad4..2b92bc0a2f45 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveBucketProperty.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveBucketProperty.java @@ -36,19 +36,16 @@ public class HiveBucketProperty { private final List bucketedBy; - private final BucketingVersion bucketingVersion; private final int bucketCount; private final List sortedBy; @JsonCreator public HiveBucketProperty( @JsonProperty("bucketedBy") List bucketedBy, - @JsonProperty("bucketingVersion") BucketingVersion bucketingVersion, @JsonProperty("bucketCount") int bucketCount, @JsonProperty("sortedBy") List sortedBy) { this.bucketedBy = ImmutableList.copyOf(requireNonNull(bucketedBy, "bucketedBy is null")); - this.bucketingVersion = requireNonNull(bucketingVersion, "bucketingVersion is null"); this.bucketCount = bucketCount; this.sortedBy = ImmutableList.copyOf(requireNonNull(sortedBy, "sortedBy is null")); } @@ -75,7 +72,7 @@ public static Optional fromStorageDescriptor(Map name.toLowerCase(ENGLISH)) .collect(toImmutableList()); - return Optional.of(new HiveBucketProperty(bucketColumnNames, bucketingVersion, storageDescriptor.getNumBuckets(), sortedBy)); + return Optional.of(new HiveBucketProperty(bucketColumnNames, storageDescriptor.getNumBuckets(), sortedBy)); } @JsonProperty @@ -84,12 +81,6 @@ public List getBucketedBy() return bucketedBy; } - @JsonProperty - public BucketingVersion getBucketingVersion() - { - return bucketingVersion; - } - @JsonProperty public int getBucketCount() { @@ -112,8 +103,7 @@ public boolean equals(Object o) return false; } HiveBucketProperty that = (HiveBucketProperty) o; - return bucketingVersion == that.bucketingVersion && - bucketCount == that.bucketCount && + return bucketCount == that.bucketCount && Objects.equals(bucketedBy, that.bucketedBy) && Objects.equals(sortedBy, that.sortedBy); } @@ -121,7 +111,7 @@ public boolean equals(Object o) @Override public int hashCode() { - return Objects.hash(bucketedBy, bucketingVersion, bucketCount, sortedBy); + return Objects.hash(bucketedBy, bucketCount, sortedBy); } @Override @@ -129,7 +119,6 @@ public String toString() { return toStringHelper(this) .add("bucketedBy", bucketedBy) - .add("bucketingVersion", bucketingVersion) .add("bucketCount", bucketCount) .add("sortedBy", sortedBy) .toString(); diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveColumnHandle.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveColumnHandle.java index ca27ecacf0b3..3e54e3425bb0 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveColumnHandle.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveColumnHandle.java @@ -15,6 +15,7 @@ import com.fasterxml.jackson.annotation.JsonCreator; import com.fasterxml.jackson.annotation.JsonProperty; +import com.google.common.collect.ImmutableMap; import io.airlift.slice.SizeOf; import io.trino.plugin.hive.metastore.Column; import io.trino.spi.connector.ColumnHandle; @@ -254,7 +255,7 @@ public String toString() public Column toMetastoreColumn() { - return new Column(name, getHiveType(), comment); + return new Column(name, getHiveType(), comment, ImmutableMap.of()); } public static HiveColumnHandle mergeRowIdColumnHandle() diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveColumnProperties.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveColumnProperties.java index 82eafa1c04c5..11f3da8cb590 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveColumnProperties.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveColumnProperties.java @@ -14,7 +14,7 @@ package io.trino.plugin.hive; import com.google.common.collect.ImmutableList; -import io.trino.plugin.hive.aws.athena.projection.ProjectionType; +import io.trino.plugin.hive.projection.ProjectionType; import io.trino.spi.session.PropertyMetadata; import io.trino.spi.type.ArrayType; @@ -22,13 +22,13 @@ import java.util.List; import static com.google.common.collect.ImmutableList.toImmutableList; -import static io.trino.plugin.hive.aws.athena.PartitionProjectionProperties.COLUMN_PROJECTION_DIGITS; -import static io.trino.plugin.hive.aws.athena.PartitionProjectionProperties.COLUMN_PROJECTION_FORMAT; -import static io.trino.plugin.hive.aws.athena.PartitionProjectionProperties.COLUMN_PROJECTION_INTERVAL; -import static io.trino.plugin.hive.aws.athena.PartitionProjectionProperties.COLUMN_PROJECTION_INTERVAL_UNIT; -import static io.trino.plugin.hive.aws.athena.PartitionProjectionProperties.COLUMN_PROJECTION_RANGE; -import static io.trino.plugin.hive.aws.athena.PartitionProjectionProperties.COLUMN_PROJECTION_TYPE; -import static io.trino.plugin.hive.aws.athena.PartitionProjectionProperties.COLUMN_PROJECTION_VALUES; +import static io.trino.plugin.hive.projection.PartitionProjectionProperties.COLUMN_PROJECTION_DIGITS; +import static io.trino.plugin.hive.projection.PartitionProjectionProperties.COLUMN_PROJECTION_FORMAT; +import static io.trino.plugin.hive.projection.PartitionProjectionProperties.COLUMN_PROJECTION_INTERVAL; +import static io.trino.plugin.hive.projection.PartitionProjectionProperties.COLUMN_PROJECTION_INTERVAL_UNIT; +import static io.trino.plugin.hive.projection.PartitionProjectionProperties.COLUMN_PROJECTION_RANGE; +import static io.trino.plugin.hive.projection.PartitionProjectionProperties.COLUMN_PROJECTION_TYPE; +import static io.trino.plugin.hive.projection.PartitionProjectionProperties.COLUMN_PROJECTION_VALUES; import static io.trino.spi.session.PropertyMetadata.enumProperty; import static io.trino.spi.session.PropertyMetadata.integerProperty; import static io.trino.spi.session.PropertyMetadata.stringProperty; diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveCompressionCodec.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveCompressionCodec.java index bc8307909b17..ac2ab60a555c 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveCompressionCodec.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveCompressionCodec.java @@ -25,7 +25,7 @@ public enum HiveCompressionCodec { NONE(null, CompressionKind.NONE, CompressionCodec.UNCOMPRESSED, AvroCompressionKind.NULL), SNAPPY(io.trino.hive.formats.compression.CompressionKind.SNAPPY, CompressionKind.SNAPPY, CompressionCodec.SNAPPY, AvroCompressionKind.SNAPPY), - LZ4(io.trino.hive.formats.compression.CompressionKind.LZ4, CompressionKind.LZ4, CompressionCodec.LZ4, null), + LZ4(io.trino.hive.formats.compression.CompressionKind.LZ4, CompressionKind.LZ4, null, null), ZSTD(io.trino.hive.formats.compression.CompressionKind.ZSTD, CompressionKind.ZSTD, CompressionCodec.ZSTD, AvroCompressionKind.ZSTANDARD), // Using DEFLATE for GZIP for Avro for now so Avro files can be written in default configuration // TODO(https://github.com/trinodb/trino/issues/12580) change GZIP to be unsupported for Avro when we change Trino default compression to be storage format aware @@ -33,7 +33,7 @@ public enum HiveCompressionCodec private final Optional hiveCompressionKind; private final CompressionKind orcCompressionKind; - private final CompressionCodec parquetCompressionCodec; + private final Optional parquetCompressionCodec; private final Optional avroCompressionKind; @@ -45,7 +45,7 @@ public enum HiveCompressionCodec { this.hiveCompressionKind = Optional.ofNullable(hiveCompressionKind); this.orcCompressionKind = requireNonNull(orcCompressionKind, "orcCompressionKind is null"); - this.parquetCompressionCodec = requireNonNull(parquetCompressionCodec, "parquetCompressionCodec is null"); + this.parquetCompressionCodec = Optional.ofNullable(parquetCompressionCodec); this.avroCompressionKind = Optional.ofNullable(avroCompressionKind); } @@ -59,7 +59,7 @@ public CompressionKind getOrcCompressionKind() return orcCompressionKind; } - public CompressionCodec getParquetCompressionCodec() + public Optional getParquetCompressionCodec() { return parquetCompressionCodec; } diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveCompressionCodecs.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveCompressionCodecs.java index b86ef0ce3d28..b046a60397fb 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveCompressionCodecs.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveCompressionCodecs.java @@ -28,7 +28,7 @@ public static HiveCompressionCodec selectCompressionCodec(ConnectorSession sessi HiveCompressionOption compressionOption = HiveSessionProperties.getCompressionCodec(session); return HiveStorageFormat.getHiveStorageFormat(storageFormat) .map(format -> selectCompressionCodec(compressionOption, format)) - .orElseGet(() -> selectCompressionCodec(compressionOption)); + .orElseGet(() -> toCompressionCodec(compressionOption)); } public static HiveCompressionCodec selectCompressionCodec(ConnectorSession session, HiveStorageFormat storageFormat) @@ -38,17 +38,18 @@ public static HiveCompressionCodec selectCompressionCodec(ConnectorSession sessi public static HiveCompressionCodec selectCompressionCodec(HiveCompressionOption compressionOption, HiveStorageFormat storageFormat) { - HiveCompressionCodec selectedCodec = selectCompressionCodec(compressionOption); + HiveCompressionCodec selectedCodec = toCompressionCodec(compressionOption); // perform codec vs format validation - if (storageFormat == HiveStorageFormat.AVRO && selectedCodec.getAvroCompressionKind().isEmpty()) { - throw new TrinoException(HIVE_UNSUPPORTED_FORMAT, "Compression codec %s not supported for AVRO".formatted(selectedCodec)); + if ((storageFormat == HiveStorageFormat.PARQUET && selectedCodec.getParquetCompressionCodec().isEmpty()) || + (storageFormat == HiveStorageFormat.AVRO && selectedCodec.getAvroCompressionKind().isEmpty())) { + throw new TrinoException(HIVE_UNSUPPORTED_FORMAT, "Compression codec %s not supported for %s".formatted(selectedCodec, storageFormat.humanName())); } return selectedCodec; } - private static HiveCompressionCodec selectCompressionCodec(HiveCompressionOption compressionOption) + public static HiveCompressionCodec toCompressionCodec(HiveCompressionOption compressionOption) { return switch (compressionOption) { case NONE -> HiveCompressionCodec.NONE; diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveConfig.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveConfig.java index ab85c5c7b8fb..07330abb19a2 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveConfig.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveConfig.java @@ -177,6 +177,10 @@ public class HiveConfig private boolean partitionProjectionEnabled; + private S3StorageClassFilter s3StorageClassFilter = S3StorageClassFilter.READ_ALL; + + private int metadataParallelism = 8; + public boolean isSingleStatementWritesOnly() { return singleStatementWritesOnly; @@ -1289,4 +1293,31 @@ public HiveConfig setPartitionProjectionEnabled(boolean enabledAthenaPartitionPr this.partitionProjectionEnabled = enabledAthenaPartitionProjection; return this; } + + public S3StorageClassFilter getS3StorageClassFilter() + { + return s3StorageClassFilter; + } + + @Config("hive.s3.storage-class-filter") + @ConfigDescription("Filter based on storage class of S3 object") + public HiveConfig setS3StorageClassFilter(S3StorageClassFilter s3StorageClassFilter) + { + this.s3StorageClassFilter = s3StorageClassFilter; + return this; + } + + @Min(1) + public int getMetadataParallelism() + { + return metadataParallelism; + } + + @ConfigDescription("Limits metadata enumeration calls parallelism") + @Config("hive.metadata.parallelism") + public HiveConfig setMetadataParallelism(int metadataParallelism) + { + this.metadataParallelism = metadataParallelism; + return this; + } } diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveConnector.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveConnector.java index 862c380f0a40..5350afc9ed6f 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveConnector.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveConnector.java @@ -29,7 +29,6 @@ import io.trino.spi.connector.ConnectorSplitManager; import io.trino.spi.connector.ConnectorTransactionHandle; import io.trino.spi.connector.TableProcedureMetadata; -import io.trino.spi.eventlistener.EventListener; import io.trino.spi.procedure.Procedure; import io.trino.spi.session.PropertyMetadata; import io.trino.spi.transaction.IsolationLevel; @@ -55,13 +54,11 @@ public class HiveConnector private final ConnectorNodePartitioningProvider nodePartitioningProvider; private final Set procedures; private final Set tableProcedures; - private final Set eventListeners; private final List> sessionProperties; private final List> schemaProperties; private final List> tableProperties; private final List> columnProperties; private final List> analyzeProperties; - private final List> materializedViewProperties; private final Optional accessControl; private final ClassLoader classLoader; @@ -79,13 +76,11 @@ public HiveConnector( ConnectorNodePartitioningProvider nodePartitioningProvider, Set procedures, Set tableProcedures, - Set eventListeners, Set sessionPropertiesProviders, List> schemaProperties, List> tableProperties, List> columnProperties, List> analyzeProperties, - List> materializedViewProperties, Optional accessControl, boolean singleStatementWritesOnly, ClassLoader classLoader) @@ -99,7 +94,6 @@ public HiveConnector( this.nodePartitioningProvider = requireNonNull(nodePartitioningProvider, "nodePartitioningProvider is null"); this.procedures = ImmutableSet.copyOf(requireNonNull(procedures, "procedures is null")); this.tableProcedures = ImmutableSet.copyOf(requireNonNull(tableProcedures, "tableProcedures is null")); - this.eventListeners = ImmutableSet.copyOf(requireNonNull(eventListeners, "eventListeners is null")); this.sessionProperties = sessionPropertiesProviders.stream() .flatMap(sessionPropertiesProvider -> sessionPropertiesProvider.getSessionProperties().stream()) .collect(toImmutableList()); @@ -107,7 +101,6 @@ public HiveConnector( this.tableProperties = ImmutableList.copyOf(requireNonNull(tableProperties, "tableProperties is null")); this.columnProperties = ImmutableList.copyOf(requireNonNull(columnProperties, "columnProperties is null")); this.analyzeProperties = ImmutableList.copyOf(requireNonNull(analyzeProperties, "analyzeProperties is null")); - this.materializedViewProperties = requireNonNull(materializedViewProperties, "materializedViewProperties is null"); this.accessControl = requireNonNull(accessControl, "accessControl is null"); this.singleStatementWritesOnly = singleStatementWritesOnly; this.classLoader = requireNonNull(classLoader, "classLoader is null"); @@ -181,18 +174,6 @@ public List> getColumnProperties() return this.columnProperties; } - @Override - public List> getMaterializedViewProperties() - { - return materializedViewProperties; - } - - @Override - public Iterable getEventListeners() - { - return eventListeners; - } - @Override public ConnectorAccessControl getAccessControl() { diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveFileWriterFactory.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveFileWriterFactory.java index 4efac69d9a8e..05c7feddd4c9 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveFileWriterFactory.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveFileWriterFactory.java @@ -19,9 +19,9 @@ import io.trino.spi.connector.ConnectorSession; import java.util.List; +import java.util.Map; import java.util.Optional; import java.util.OptionalInt; -import java.util.Properties; public interface HiveFileWriterFactory { @@ -30,7 +30,7 @@ Optional createFileWriter( List inputColumnNames, StorageFormat storageFormat, HiveCompressionCodec compressionCodec, - Properties schema, + Map schema, ConnectorSession session, OptionalInt bucketNumber, AcidTransaction transaction, diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveInsertTableHandle.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveInsertTableHandle.java index 2a8a10272f56..9e7ce50c69e1 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveInsertTableHandle.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveInsertTableHandle.java @@ -33,7 +33,7 @@ public HiveInsertTableHandle( @JsonProperty("inputColumns") List inputColumns, @JsonProperty("pageSinkMetadata") HivePageSinkMetadata pageSinkMetadata, @JsonProperty("locationHandle") LocationHandle locationHandle, - @JsonProperty("bucketProperty") Optional bucketProperty, + @JsonProperty("bucketInfo") Optional bucketInfo, @JsonProperty("tableStorageFormat") HiveStorageFormat tableStorageFormat, @JsonProperty("partitionStorageFormat") HiveStorageFormat partitionStorageFormat, @JsonProperty("transaction") AcidTransaction transaction, @@ -45,7 +45,7 @@ public HiveInsertTableHandle( inputColumns, pageSinkMetadata, locationHandle, - bucketProperty, + bucketInfo, tableStorageFormat, partitionStorageFormat, transaction, diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveLocationService.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveLocationService.java index 6da35e45ba4f..a790c177284d 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveLocationService.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveLocationService.java @@ -15,15 +15,14 @@ import com.google.inject.Inject; import io.trino.filesystem.Location; -import io.trino.hdfs.HdfsContext; -import io.trino.hdfs.HdfsEnvironment; +import io.trino.filesystem.TrinoFileSystem; +import io.trino.filesystem.TrinoFileSystemFactory; import io.trino.plugin.hive.LocationHandle.WriteMode; import io.trino.plugin.hive.metastore.Partition; import io.trino.plugin.hive.metastore.SemiTransactionalHiveMetastore; import io.trino.plugin.hive.metastore.Table; import io.trino.spi.TrinoException; import io.trino.spi.connector.ConnectorSession; -import org.apache.hadoop.fs.Path; import java.util.Optional; @@ -33,10 +32,8 @@ import static io.trino.plugin.hive.LocationHandle.WriteMode.STAGE_AND_MOVE_TO_TARGET_DIRECTORY; import static io.trino.plugin.hive.util.AcidTables.isTransactionalTable; import static io.trino.plugin.hive.util.HiveWriteUtils.createTemporaryPath; +import static io.trino.plugin.hive.util.HiveWriteUtils.directoryExists; import static io.trino.plugin.hive.util.HiveWriteUtils.getTableDefaultLocation; -import static io.trino.plugin.hive.util.HiveWriteUtils.isHdfsEncrypted; -import static io.trino.plugin.hive.util.HiveWriteUtils.isS3FileSystem; -import static io.trino.plugin.hive.util.HiveWriteUtils.pathExists; import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED; import static java.lang.String.format; import static java.util.Objects.requireNonNull; @@ -44,14 +41,14 @@ public class HiveLocationService implements LocationService { - private final HdfsEnvironment hdfsEnvironment; + private final TrinoFileSystemFactory fileSystemFactory; private final boolean temporaryStagingDirectoryEnabled; private final String temporaryStagingDirectoryPath; @Inject - public HiveLocationService(HdfsEnvironment hdfsEnvironment, HiveConfig hiveConfig) + public HiveLocationService(TrinoFileSystemFactory fileSystemFactory, HiveConfig hiveConfig) { - this.hdfsEnvironment = requireNonNull(hdfsEnvironment, "hdfsEnvironment is null"); + this.fileSystemFactory = requireNonNull(fileSystemFactory, "fileSystemFactory is null"); this.temporaryStagingDirectoryEnabled = hiveConfig.isTemporaryStagingDirectoryEnabled(); this.temporaryStagingDirectoryPath = hiveConfig.getTemporaryStagingDirectoryPath(); } @@ -59,11 +56,11 @@ public HiveLocationService(HdfsEnvironment hdfsEnvironment, HiveConfig hiveConfi @Override public Location forNewTable(SemiTransactionalHiveMetastore metastore, ConnectorSession session, String schemaName, String tableName) { - HdfsContext context = new HdfsContext(session); - Location targetPath = getTableDefaultLocation(context, metastore, hdfsEnvironment, schemaName, tableName); + TrinoFileSystem fileSystem = fileSystemFactory.create(session); + Location targetPath = getTableDefaultLocation(metastore, fileSystem, schemaName, tableName); // verify the target directory for table - if (pathExists(context, hdfsEnvironment, new Path(targetPath.toString()))) { + if (directoryExists(fileSystem, targetPath).orElse(false)) { throw new TrinoException(HIVE_PATH_ALREADY_EXISTS, format("Target directory for table '%s.%s' already exists: %s", schemaName, tableName, targetPath)); } return targetPath; @@ -72,18 +69,20 @@ public Location forNewTable(SemiTransactionalHiveMetastore metastore, ConnectorS @Override public LocationHandle forNewTableAsSelect(SemiTransactionalHiveMetastore metastore, ConnectorSession session, String schemaName, String tableName, Optional externalLocation) { - HdfsContext context = new HdfsContext(session); - Location targetPath = externalLocation.orElseGet(() -> getTableDefaultLocation(context, metastore, hdfsEnvironment, schemaName, tableName)); + TrinoFileSystem fileSystem = fileSystemFactory.create(session); + Location targetPath = externalLocation.orElseGet(() -> getTableDefaultLocation(metastore, fileSystem, schemaName, tableName)); // verify the target directory for the table - if (pathExists(context, hdfsEnvironment, new Path(targetPath.toString()))) { + if (directoryExists(fileSystem, targetPath).orElse(false)) { throw new TrinoException(HIVE_PATH_ALREADY_EXISTS, format("Target directory for table '%s.%s' already exists: %s", schemaName, tableName, targetPath)); } - // TODO detect when existing table's location is a on a different file system than the temporary directory - if (shouldUseTemporaryDirectory(context, new Path(targetPath.toString()), externalLocation.isPresent())) { - Location writePath = createTemporaryPath(context, hdfsEnvironment, new Path(targetPath.toString()), temporaryStagingDirectoryPath); - return new LocationHandle(targetPath, writePath, STAGE_AND_MOVE_TO_TARGET_DIRECTORY); + // Skip using temporary directory if the destination is external. Target may be on a different file system. + if (temporaryStagingDirectoryEnabled && externalLocation.isEmpty()) { + Optional writePath = createTemporaryPath(fileSystem, session.getIdentity(), targetPath, temporaryStagingDirectoryPath); + if (writePath.isPresent()) { + return new LocationHandle(targetPath, writePath.get(), STAGE_AND_MOVE_TO_TARGET_DIRECTORY); + } } return new LocationHandle(targetPath, targetPath, DIRECT_TO_TARGET_NEW_DIRECTORY); } @@ -91,12 +90,14 @@ public LocationHandle forNewTableAsSelect(SemiTransactionalHiveMetastore metasto @Override public LocationHandle forExistingTable(SemiTransactionalHiveMetastore metastore, ConnectorSession session, Table table) { - HdfsContext context = new HdfsContext(session); + TrinoFileSystem fileSystem = fileSystemFactory.create(session); Location targetPath = Location.of(table.getStorage().getLocation()); - if (shouldUseTemporaryDirectory(context, new Path(targetPath.toString()), false) && !isTransactionalTable(table.getParameters())) { - Location writePath = createTemporaryPath(context, hdfsEnvironment, new Path(targetPath.toString()), temporaryStagingDirectoryPath); - return new LocationHandle(targetPath, writePath, STAGE_AND_MOVE_TO_TARGET_DIRECTORY); + if (temporaryStagingDirectoryEnabled && !isTransactionalTable(table.getParameters())) { + Optional writePath = createTemporaryPath(fileSystem, session.getIdentity(), targetPath, temporaryStagingDirectoryPath); + if (writePath.isPresent()) { + return new LocationHandle(targetPath, writePath.get(), STAGE_AND_MOVE_TO_TARGET_DIRECTORY); + } } return new LocationHandle(targetPath, targetPath, DIRECT_TO_TARGET_EXISTING_DIRECTORY); } @@ -109,17 +110,6 @@ public LocationHandle forOptimize(SemiTransactionalHiveMetastore metastore, Conn return new LocationHandle(targetPath, targetPath, DIRECT_TO_TARGET_EXISTING_DIRECTORY); } - private boolean shouldUseTemporaryDirectory(HdfsContext context, Path path, boolean hasExternalLocation) - { - return temporaryStagingDirectoryEnabled - // skip using temporary directory for S3 - && !isS3FileSystem(context, hdfsEnvironment, path) - // skip using temporary directory if destination is encrypted; it's not possible to move a file between encryption zones - && !isHdfsEncrypted(context, hdfsEnvironment, path) - // Skip using temporary directory if destination is external. Target may be on a different file system. - && !hasExternalLocation; - } - @Override public WriteInfo getQueryWriteInfo(LocationHandle locationHandle) { diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveMaterializedViewMetadata.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveMaterializedViewMetadata.java deleted file mode 100644 index f47ad0dba86c..000000000000 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveMaterializedViewMetadata.java +++ /dev/null @@ -1,49 +0,0 @@ -/* - * 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.hive; - -import io.trino.spi.connector.ConnectorMaterializedViewDefinition; -import io.trino.spi.connector.ConnectorSession; -import io.trino.spi.connector.MaterializedViewFreshness; -import io.trino.spi.connector.SchemaTableName; - -import java.util.List; -import java.util.Map; -import java.util.Optional; -import java.util.concurrent.CompletableFuture; - -public interface HiveMaterializedViewMetadata -{ - void createMaterializedView(ConnectorSession session, SchemaTableName viewName, ConnectorMaterializedViewDefinition definition, boolean replace, boolean ignoreExisting); - - void dropMaterializedView(ConnectorSession session, SchemaTableName viewName); - - List listMaterializedViews(ConnectorSession session, Optional schemaName); - - Map getMaterializedViews(ConnectorSession session, Optional schemaName); - - Optional getMaterializedView(ConnectorSession session, SchemaTableName viewName); - - MaterializedViewFreshness getMaterializedViewFreshness(ConnectorSession session, SchemaTableName name); - - boolean delegateMaterializedViewRefreshToConnector(ConnectorSession session, SchemaTableName viewName); - - CompletableFuture refreshMaterializedView(ConnectorSession session, SchemaTableName viewName); - - void renameMaterializedView(ConnectorSession session, SchemaTableName existingViewName, SchemaTableName newViewName); - - void setMaterializedViewProperties(ConnectorSession session, SchemaTableName viewName, Map> properties); - - void setMaterializedViewColumnComment(ConnectorSession session, SchemaTableName viewName, String columnName, Optional comment); -} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveMetadata.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveMetadata.java index 38c23f6cddff..e29cc18e4ba3 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveMetadata.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveMetadata.java @@ -24,37 +24,40 @@ import com.google.common.collect.Iterables; import com.google.common.collect.Maps; import com.google.common.collect.Sets; +import dev.failsafe.Failsafe; +import dev.failsafe.RetryPolicy; import io.airlift.json.JsonCodec; import io.airlift.log.Logger; import io.airlift.slice.Slice; import io.airlift.units.DataSize; +import io.trino.filesystem.FileIterator; import io.trino.filesystem.Location; +import io.trino.filesystem.TrinoFileSystem; import io.trino.filesystem.TrinoFileSystemFactory; -import io.trino.hdfs.HdfsContext; -import io.trino.hdfs.HdfsEnvironment; import io.trino.plugin.base.CatalogName; import io.trino.plugin.base.projection.ApplyProjectionUtil; import io.trino.plugin.base.projection.ApplyProjectionUtil.ProjectedColumnRepresentation; import io.trino.plugin.hive.HiveSessionProperties.InsertExistingPartitionsBehavior; +import io.trino.plugin.hive.HiveWritableTableHandle.BucketInfo; import io.trino.plugin.hive.LocationService.WriteInfo; import io.trino.plugin.hive.acid.AcidOperation; import io.trino.plugin.hive.acid.AcidTransaction; -import io.trino.plugin.hive.aws.athena.PartitionProjectionService; import io.trino.plugin.hive.fs.DirectoryLister; import io.trino.plugin.hive.metastore.Column; import io.trino.plugin.hive.metastore.Database; import io.trino.plugin.hive.metastore.HiveColumnStatistics; import io.trino.plugin.hive.metastore.HivePrincipal; +import io.trino.plugin.hive.metastore.MetastoreUtil; import io.trino.plugin.hive.metastore.Partition; import io.trino.plugin.hive.metastore.PrincipalPrivileges; import io.trino.plugin.hive.metastore.SemiTransactionalHiveMetastore; import io.trino.plugin.hive.metastore.SortingColumn; import io.trino.plugin.hive.metastore.StorageFormat; import io.trino.plugin.hive.metastore.Table; +import io.trino.plugin.hive.metastore.TableInfo; import io.trino.plugin.hive.procedure.OptimizeTableProcedure; import io.trino.plugin.hive.security.AccessControlMetadata; import io.trino.plugin.hive.statistics.HiveStatisticsProvider; -import io.trino.plugin.hive.util.HiveBucketing; import io.trino.plugin.hive.util.HiveUtil; import io.trino.plugin.hive.util.SerdeConstants; import io.trino.spi.ErrorType; @@ -69,7 +72,6 @@ import io.trino.spi.connector.ColumnMetadata; import io.trino.spi.connector.ConnectorAnalyzeMetadata; import io.trino.spi.connector.ConnectorInsertTableHandle; -import io.trino.spi.connector.ConnectorMaterializedViewDefinition; import io.trino.spi.connector.ConnectorMergeTableHandle; import io.trino.spi.connector.ConnectorOutputMetadata; import io.trino.spi.connector.ConnectorOutputTableHandle; @@ -86,9 +88,9 @@ import io.trino.spi.connector.ConstraintApplicationResult; import io.trino.spi.connector.DiscretePredicates; import io.trino.spi.connector.LocalProperty; -import io.trino.spi.connector.MaterializedViewFreshness; import io.trino.spi.connector.MetadataProvider; import io.trino.spi.connector.ProjectionApplicationResult; +import io.trino.spi.connector.RelationType; import io.trino.spi.connector.RetryMode; import io.trino.spi.connector.RowChangeParadigm; import io.trino.spi.connector.SchemaNotFoundException; @@ -98,7 +100,6 @@ import io.trino.spi.connector.SystemTable; import io.trino.spi.connector.TableColumnsMetadata; import io.trino.spi.connector.TableNotFoundException; -import io.trino.spi.connector.TableScanRedirectApplicationResult; import io.trino.spi.connector.ViewNotFoundException; import io.trino.spi.expression.ConnectorExpression; import io.trino.spi.expression.Variable; @@ -123,15 +124,9 @@ import io.trino.spi.type.TypeManager; import org.apache.avro.Schema; import org.apache.avro.SchemaParseException; -import org.apache.hadoop.fs.FileSystem; -import org.apache.hadoop.fs.LocatedFileStatus; -import org.apache.hadoop.fs.Path; -import org.apache.hadoop.fs.RemoteIterator; -import java.io.File; +import java.io.FileNotFoundException; import java.io.IOException; -import java.net.MalformedURLException; -import java.net.URL; import java.util.ArrayList; import java.util.Collection; import java.util.Collections; @@ -141,12 +136,14 @@ import java.util.List; import java.util.Map; import java.util.NoSuchElementException; +import java.util.Objects; import java.util.Optional; import java.util.OptionalInt; import java.util.OptionalLong; -import java.util.Properties; import java.util.Set; -import java.util.concurrent.CompletableFuture; +import java.util.concurrent.Callable; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.Executor; import java.util.function.Function; import java.util.function.Supplier; import java.util.regex.Pattern; @@ -161,8 +158,10 @@ import static com.google.common.collect.ImmutableSet.toImmutableSet; import static com.google.common.collect.Iterables.concat; import static com.google.common.collect.Iterables.getOnlyElement; +import static io.trino.parquet.writer.ParquetWriter.SUPPORTED_BLOOM_FILTER_TYPES; import static io.trino.plugin.base.projection.ApplyProjectionUtil.extractSupportedProjectedColumns; import static io.trino.plugin.base.projection.ApplyProjectionUtil.replaceWithNewVariables; +import static io.trino.plugin.base.util.ExecutorUtil.processWithAdditionalThreads; import static io.trino.plugin.hive.HiveAnalyzeProperties.getColumnNames; import static io.trino.plugin.hive.HiveAnalyzeProperties.getPartitionList; import static io.trino.plugin.hive.HiveApplyProjectionUtil.find; @@ -215,6 +214,7 @@ import static io.trino.plugin.hive.HiveTableProperties.NULL_FORMAT_PROPERTY; import static io.trino.plugin.hive.HiveTableProperties.ORC_BLOOM_FILTER_COLUMNS; import static io.trino.plugin.hive.HiveTableProperties.ORC_BLOOM_FILTER_FPP; +import static io.trino.plugin.hive.HiveTableProperties.PARQUET_BLOOM_FILTER_COLUMNS; import static io.trino.plugin.hive.HiveTableProperties.PARTITIONED_BY_PROPERTY; import static io.trino.plugin.hive.HiveTableProperties.REGEX_CASE_INSENSITIVE; import static io.trino.plugin.hive.HiveTableProperties.REGEX_PATTERN; @@ -226,7 +226,7 @@ import static io.trino.plugin.hive.HiveTableProperties.TEXTFILE_FIELD_SEPARATOR_ESCAPE; import static io.trino.plugin.hive.HiveTableProperties.getAvroSchemaLiteral; import static io.trino.plugin.hive.HiveTableProperties.getAvroSchemaUrl; -import static io.trino.plugin.hive.HiveTableProperties.getBucketProperty; +import static io.trino.plugin.hive.HiveTableProperties.getBucketInfo; import static io.trino.plugin.hive.HiveTableProperties.getExternalLocation; import static io.trino.plugin.hive.HiveTableProperties.getExtraProperties; import static io.trino.plugin.hive.HiveTableProperties.getFooterSkipCount; @@ -235,6 +235,7 @@ import static io.trino.plugin.hive.HiveTableProperties.getNullFormat; import static io.trino.plugin.hive.HiveTableProperties.getOrcBloomFilterColumns; import static io.trino.plugin.hive.HiveTableProperties.getOrcBloomFilterFpp; +import static io.trino.plugin.hive.HiveTableProperties.getParquetBloomFilterColumns; import static io.trino.plugin.hive.HiveTableProperties.getPartitionedBy; import static io.trino.plugin.hive.HiveTableProperties.getRegexPattern; import static io.trino.plugin.hive.HiveTableProperties.getSingleCharacterProperty; @@ -243,6 +244,7 @@ import static io.trino.plugin.hive.HiveTimestampPrecision.NANOSECONDS; import static io.trino.plugin.hive.HiveType.HIVE_STRING; import static io.trino.plugin.hive.HiveType.toHiveType; +import static io.trino.plugin.hive.HiveWritableTableHandle.BucketInfo.createBucketInfo; import static io.trino.plugin.hive.HiveWriterFactory.computeNonTransactionalBucketedFilename; import static io.trino.plugin.hive.HiveWriterFactory.computeTransactionalBucketedFilename; import static io.trino.plugin.hive.LocationHandle.WriteMode.DIRECT_TO_TARGET_EXISTING_DIRECTORY; @@ -269,16 +271,25 @@ import static io.trino.plugin.hive.metastore.PrincipalPrivileges.fromHivePrivilegeInfos; import static io.trino.plugin.hive.metastore.SemiTransactionalHiveMetastore.PartitionUpdateInfo; import static io.trino.plugin.hive.metastore.SemiTransactionalHiveMetastore.cleanExtraOutputFiles; +import static io.trino.plugin.hive.metastore.StatisticsUpdateMode.MERGE_INCREMENTAL; import static io.trino.plugin.hive.metastore.StorageFormat.VIEW_STORAGE_FORMAT; import static io.trino.plugin.hive.metastore.StorageFormat.fromHiveStorageFormat; -import static io.trino.plugin.hive.metastore.thrift.ThriftMetastoreUtil.STATS_PROPERTIES; +import static io.trino.plugin.hive.metastore.thrift.ThriftMetastoreUtil.getSupportedColumnStatistics; +import static io.trino.plugin.hive.projection.PartitionProjectionProperties.arePartitionProjectionPropertiesSet; +import static io.trino.plugin.hive.projection.PartitionProjectionProperties.getPartitionProjectionHiveTableProperties; +import static io.trino.plugin.hive.projection.PartitionProjectionProperties.getPartitionProjectionTrinoTableProperties; import static io.trino.plugin.hive.type.Category.PRIMITIVE; import static io.trino.plugin.hive.util.AcidTables.deltaSubdir; import static io.trino.plugin.hive.util.AcidTables.isFullAcidTable; import static io.trino.plugin.hive.util.AcidTables.isTransactionalTable; import static io.trino.plugin.hive.util.AcidTables.writeAcidVersionFile; -import static io.trino.plugin.hive.util.HiveBucketing.getHiveBucketHandle; +import static io.trino.plugin.hive.util.HiveBucketing.BucketingVersion.BUCKETING_V2; +import static io.trino.plugin.hive.util.HiveBucketing.getBucketingVersion; +import static io.trino.plugin.hive.util.HiveBucketing.getHiveTablePartitioningForRead; +import static io.trino.plugin.hive.util.HiveBucketing.getHiveTablePartitioningForWrite; import static io.trino.plugin.hive.util.HiveBucketing.isSupportedBucketing; +import static io.trino.plugin.hive.util.HiveTypeUtil.getType; +import static io.trino.plugin.hive.util.HiveTypeUtil.getTypeSignature; import static io.trino.plugin.hive.util.HiveUtil.columnMetadataGetter; import static io.trino.plugin.hive.util.HiveUtil.getPartitionKeyColumnHandles; import static io.trino.plugin.hive.util.HiveUtil.getRegularColumnHandles; @@ -292,19 +303,15 @@ import static io.trino.plugin.hive.util.HiveUtil.toPartitionValues; import static io.trino.plugin.hive.util.HiveUtil.verifyPartitionTypeSupported; import static io.trino.plugin.hive.util.HiveWriteUtils.checkTableIsWritable; -import static io.trino.plugin.hive.util.HiveWriteUtils.checkedDelete; import static io.trino.plugin.hive.util.HiveWriteUtils.createPartitionValues; import static io.trino.plugin.hive.util.HiveWriteUtils.isFileCreatedByQuery; -import static io.trino.plugin.hive.util.HiveWriteUtils.isS3FileSystem; import static io.trino.plugin.hive.util.HiveWriteUtils.isWritableType; -import static io.trino.plugin.hive.util.RetryDriver.retry; -import static io.trino.plugin.hive.util.Statistics.ReduceOperator.ADD; import static io.trino.plugin.hive.util.Statistics.createComputedStatisticsToPartitionMap; import static io.trino.plugin.hive.util.Statistics.createEmptyPartitionStatistics; import static io.trino.plugin.hive.util.Statistics.fromComputedStatistics; -import static io.trino.plugin.hive.util.Statistics.reduce; import static io.trino.plugin.hive.util.SystemTables.getSourceTableNameFromSystemTable; import static io.trino.spi.StandardErrorCode.INVALID_ANALYZE_PROPERTY; +import static io.trino.spi.StandardErrorCode.INVALID_COLUMN_PROPERTY; import static io.trino.spi.StandardErrorCode.INVALID_SCHEMA_PROPERTY; import static io.trino.spi.StandardErrorCode.INVALID_TABLE_PROPERTY; import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED; @@ -330,14 +337,13 @@ public class HiveMetadata { private static final Logger log = Logger.get(HiveMetadata.class); - public static final String PRESTO_VERSION_NAME = "presto_version"; + public static final String TRINO_VERSION_NAME = "trino_version"; public static final String TRINO_CREATED_BY = "trino_created_by"; - public static final String PRESTO_QUERY_ID_NAME = "presto_query_id"; + public static final String TRINO_QUERY_ID_NAME = "trino_query_id"; public static final String BUCKETING_VERSION = "bucketing_version"; public static final String TABLE_COMMENT = "comment"; public static final String STORAGE_TABLE = "storage_table"; public static final String TRANSACTIONAL = "transactional"; - public static final String PRESTO_VIEW_COMMENT = "Presto View"; public static final String PRESTO_VIEW_EXPANDED_TEXT_MARKER = "/* Presto View */"; public static final String ORC_BLOOM_FILTER_COLUMNS_KEY = "orc.bloom.filter.columns"; @@ -357,6 +363,8 @@ public class HiveMetadata private static final String CSV_QUOTE_KEY = "quoteChar"; private static final String CSV_ESCAPE_KEY = "escapeChar"; + public static final String PARQUET_BLOOM_FILTER_COLUMNS_KEY = "parquet.bloom.filter.columns"; + private static final String REGEX_KEY = "input.regex"; private static final String REGEX_CASE_SENSITIVE_KEY = "input.regex.case.insensitive"; @@ -364,12 +372,17 @@ public class HiveMetadata public static final String MODIFYING_NON_TRANSACTIONAL_TABLE_MESSAGE = "Modifying Hive table rows is only supported for transactional tables"; + private static final RetryPolicy DELETE_RETRY_POLICY = RetryPolicy.builder() + .withDelay(java.time.Duration.ofSeconds(1)) + .withMaxDuration(java.time.Duration.ofSeconds(30)) + .withMaxAttempts(10) + .build(); + private final CatalogName catalogName; private final SemiTransactionalHiveMetastore metastore; private final boolean autoCommit; private final Set fileWriterFactories; private final TrinoFileSystemFactory fileSystemFactory; - private final HdfsEnvironment hdfsEnvironment; private final HivePartitionManager partitionManager; private final TypeManager typeManager; private final MetadataProvider metadataProvider; @@ -380,17 +393,17 @@ public class HiveMetadata private final boolean translateHiveViews; private final boolean hiveViewsRunAsInvoker; private final boolean hideDeltaLakeTables; - private final String prestoVersion; + private final String trinoVersion; private final HiveStatisticsProvider hiveStatisticsProvider; - private final HiveRedirectionsProvider hiveRedirectionsProvider; private final Set systemTableProviders; - private final HiveMaterializedViewMetadata hiveMaterializedViewMetadata; private final AccessControlMetadata accessControlMetadata; private final DirectoryLister directoryLister; - private final PartitionProjectionService partitionProjectionService; + private final boolean usingSystemSecurity; + private final boolean partitionProjectionEnabled; private final boolean allowTableRename; private final long maxPartitionDropsPerQuery; private final HiveTimestampPrecision hiveViewsTimestampPrecision; + private final Executor metadataFetchingExecutor; public HiveMetadata( CatalogName catalogName, @@ -398,7 +411,6 @@ public HiveMetadata( boolean autoCommit, Set fileWriterFactories, TrinoFileSystemFactory fileSystemFactory, - HdfsEnvironment hdfsEnvironment, HivePartitionManager partitionManager, boolean writesToNonManagedTablesEnabled, boolean createsOfNonManagedTablesEnabled, @@ -411,22 +423,21 @@ public HiveMetadata( JsonCodec partitionUpdateCodec, String trinoVersion, HiveStatisticsProvider hiveStatisticsProvider, - HiveRedirectionsProvider hiveRedirectionsProvider, Set systemTableProviders, - HiveMaterializedViewMetadata hiveMaterializedViewMetadata, AccessControlMetadata accessControlMetadata, DirectoryLister directoryLister, - PartitionProjectionService partitionProjectionService, + boolean usingSystemSecurity, + boolean partitionProjectionEnabled, boolean allowTableRename, long maxPartitionDropsPerQuery, - HiveTimestampPrecision hiveViewsTimestampPrecision) + HiveTimestampPrecision hiveViewsTimestampPrecision, + Executor metadataFetchingExecutor) { this.catalogName = requireNonNull(catalogName, "catalogName is null"); this.metastore = requireNonNull(metastore, "metastore is null"); this.autoCommit = autoCommit; this.fileWriterFactories = ImmutableSet.copyOf(requireNonNull(fileWriterFactories, "fileWriterFactories is null")); this.fileSystemFactory = requireNonNull(fileSystemFactory, "fileSystemFactory is null"); - this.hdfsEnvironment = requireNonNull(hdfsEnvironment, "hdfsEnvironment is null"); this.partitionManager = requireNonNull(partitionManager, "partitionManager is null"); this.typeManager = requireNonNull(typeManager, "typeManager is null"); this.metadataProvider = requireNonNull(metadataProvider, "metadataProvider is null"); @@ -437,17 +448,17 @@ public HiveMetadata( this.translateHiveViews = translateHiveViews; this.hiveViewsRunAsInvoker = hiveViewsRunAsInvoker; this.hideDeltaLakeTables = hideDeltaLakeTables; - this.prestoVersion = requireNonNull(trinoVersion, "trinoVersion is null"); + this.trinoVersion = requireNonNull(trinoVersion, "trinoVersion is null"); this.hiveStatisticsProvider = requireNonNull(hiveStatisticsProvider, "hiveStatisticsProvider is null"); - this.hiveRedirectionsProvider = requireNonNull(hiveRedirectionsProvider, "hiveRedirectionsProvider is null"); this.systemTableProviders = requireNonNull(systemTableProviders, "systemTableProviders is null"); - this.hiveMaterializedViewMetadata = requireNonNull(hiveMaterializedViewMetadata, "hiveMaterializedViewMetadata is null"); this.accessControlMetadata = requireNonNull(accessControlMetadata, "accessControlMetadata is null"); this.directoryLister = requireNonNull(directoryLister, "directoryLister is null"); - this.partitionProjectionService = requireNonNull(partitionProjectionService, "partitionProjectionService is null"); + this.usingSystemSecurity = usingSystemSecurity; + this.partitionProjectionEnabled = partitionProjectionEnabled; this.allowTableRename = allowTableRename; this.maxPartitionDropsPerQuery = maxPartitionDropsPerQuery; this.hiveViewsTimestampPrecision = requireNonNull(hiveViewsTimestampPrecision, "hiveViewsTimestampPrecision is null"); + this.metadataFetchingExecutor = requireNonNull(metadataFetchingExecutor, "metadataFetchingExecutor is null"); } @Override @@ -508,7 +519,7 @@ public HiveTableHandle getTableHandle(ConnectorSession session, SchemaTableName table.getParameters(), getPartitionKeyColumnHandles(table, typeManager), getRegularColumnHandles(table, typeManager, getTimestampPrecision(session)), - getHiveBucketHandle(session, table, typeManager)); + getHiveTablePartitioningForRead(session, table, typeManager)); } @Override @@ -667,7 +678,7 @@ else if (isTrinoView || isTrinoMaterializedView) { // Bucket properties table.getStorage().getBucketProperty().ifPresent(property -> { - properties.put(BUCKETING_VERSION, property.getBucketingVersion().getVersion()); + properties.put(BUCKETING_VERSION, getBucketingVersion(table.getParameters()).getVersion()); properties.put(BUCKET_COUNT_PROPERTY, property.getBucketCount()); properties.put(BUCKETED_BY_PROPERTY, property.getBucketedBy()); properties.put(SORTED_BY_PROPERTY, property.getSortedBy()); @@ -729,7 +740,7 @@ else if (isTrinoView || isTrinoMaterializedView) { getSerdeProperty(table, REGEX_CASE_SENSITIVE_KEY) .ifPresent(regexCaseInsensitive -> properties.put(REGEX_CASE_INSENSITIVE, parseBoolean(regexCaseInsensitive))); - Optional comment = Optional.ofNullable(table.getParameters().get(TABLE_COMMENT)); + Optional comment = Optional.ofNullable(table.getParameters().get(Table.TABLE_COMMENT)); String autoPurgeProperty = table.getParameters().get(AUTO_PURGE_KEY); if (parseBoolean(autoPurgeProperty)) { @@ -737,7 +748,7 @@ else if (isTrinoView || isTrinoMaterializedView) { } // Partition Projection specific properties - properties.putAll(partitionProjectionService.getPartitionProjectionTrinoTableProperties(table)); + properties.putAll(getPartitionProjectionTrinoTableProperties(table)); return new ConnectorTableMetadata(tableName, columns.build(), properties.buildOrThrow(), comment); } @@ -784,26 +795,23 @@ public Optional getInfo(ConnectorTableHandle tableHandle) @Override public List listTables(ConnectorSession session, Optional optionalSchemaName) { - if (optionalSchemaName.isEmpty()) { - Optional> allTables = metastore.getAllTables(); - if (allTables.isPresent()) { - return ImmutableList.builder() - .addAll(allTables.get().stream() - .filter(table -> !isHiveSystemSchema(table.getSchemaName())) - .collect(toImmutableList())) - .addAll(listMaterializedViews(session, optionalSchemaName)) - .build(); - } + return streamTables(session, optionalSchemaName) + .map(TableInfo::tableName) + .collect(toImmutableList()); + } + + private Stream streamTables(ConnectorSession session, Optional optionalSchemaName) + { + List>> tasks = listSchemas(session, optionalSchemaName).stream() + .map(schemaName -> (Callable>) () -> metastore.getTables(schemaName)) + .collect(toImmutableList()); + try { + return processWithAdditionalThreads(tasks, metadataFetchingExecutor).stream() + .flatMap(Collection::stream); } - ImmutableList.Builder tableNames = ImmutableList.builder(); - for (String schemaName : listSchemas(session, optionalSchemaName)) { - for (String tableName : metastore.getAllTables(schemaName)) { - tableNames.add(new SchemaTableName(schemaName, tableName)); - } + catch (ExecutionException e) { + throw new RuntimeException(e.getCause()); } - - tableNames.addAll(listMaterializedViews(session, optionalSchemaName)); - return tableNames.build(); } private List listSchemas(ConnectorSession session, Optional schemaName) @@ -934,7 +942,7 @@ public void createSchema(ConnectorSession session, String schemaName, Map location = HiveSchemaProperties.getLocation(properties).map(locationUri -> { try { - hdfsEnvironment.getFileSystem(new HdfsContext(session), new Path(locationUri)); + fileSystemFactory.create(session).directoryExists(Location.of(locationUri)); } catch (IOException e) { throw new TrinoException(INVALID_SCHEMA_PROPERTY, "Invalid location URI: " + locationUri, e); @@ -947,7 +955,7 @@ public void createSchema(ConnectorSession session, String schemaName, Map partitionedBy = getPartitionedBy(tableMetadata.getProperties()); - Optional bucketProperty = getBucketProperty(tableMetadata.getProperties()); + Optional bucketInfo = getBucketInfo(tableMetadata.getProperties()); boolean isTransactional = isTransactional(tableMetadata.getProperties()).orElse(false); - if (bucketProperty.isPresent() && getAvroSchemaUrl(tableMetadata.getProperties()) != null) { + if (bucketInfo.isPresent() && getAvroSchemaUrl(tableMetadata.getProperties()) != null) { throw new TrinoException(NOT_SUPPORTED, "Bucketing columns not supported when Avro schema url is set"); } - if (bucketProperty.isPresent() && getAvroSchemaLiteral(tableMetadata.getProperties()) != null) { + if (bucketInfo.isPresent() && getAvroSchemaLiteral(tableMetadata.getProperties()) != null) { throw new TrinoException(NOT_SUPPORTED, "Bucketing/Partitioning columns not supported when Avro schema literal is set"); } @@ -996,7 +1004,7 @@ public void createTable(ConnectorSession session, ConnectorTableMetadata tableMe validateTimestampColumns(tableMetadata.getColumns(), getTimestampPrecision(session)); List columnHandles = getColumnHandles(tableMetadata, ImmutableSet.copyOf(partitionedBy)); HiveStorageFormat hiveStorageFormat = getHiveStorageFormat(tableMetadata.getProperties()); - Map tableProperties = getEmptyTableProperties(tableMetadata, bucketProperty, new HdfsContext(session)); + Map tableProperties = getEmptyTableProperties(tableMetadata, bucketInfo, session); hiveStorageFormat.validateColumns(columnHandles); @@ -1016,7 +1024,7 @@ public void createTable(ConnectorSession session, ConnectorTableMetadata tableMe external = true; targetPath = Optional.of(getValidatedExternalLocation(externalLocation)); - checkExternalPath(new HdfsContext(session), new Path(targetPath.get().toString())); + checkExternalPathAndCreateIfNotExists(session, targetPath.get()); } else { external = false; @@ -1036,11 +1044,11 @@ public void createTable(ConnectorSession session, ConnectorTableMetadata tableMe columnHandles, hiveStorageFormat, partitionedBy, - bucketProperty, + bucketInfo, tableProperties, targetPath, external, - prestoVersion, + trinoVersion, accessControlMetadata.isUsingSystemSecurity()); PrincipalPrivileges principalPrivileges = accessControlMetadata.isUsingSystemSecurity() ? NO_PRIVILEGES : buildInitialPrivilegeSet(session.getUser()); HiveBasicStatistics basicStatistics = (!external && table.getPartitionColumns().isEmpty()) ? createZeroStatistics() : createEmptyStatistics(); @@ -1055,38 +1063,45 @@ public void createTable(ConnectorSession session, ConnectorTableMetadata tableMe false); } - private Map getEmptyTableProperties(ConnectorTableMetadata tableMetadata, Optional bucketProperty, HdfsContext hdfsContext) + private Map getEmptyTableProperties(ConnectorTableMetadata tableMetadata, Optional bucketInfo, ConnectorSession session) { HiveStorageFormat hiveStorageFormat = getHiveStorageFormat(tableMetadata.getProperties()); ImmutableMap.Builder tableProperties = ImmutableMap.builder(); // When metastore is configured with metastore.create.as.acid=true, it will also change Trino-created tables // behind the scenes. In particular, this won't work with CTAS. - boolean transactional = HiveTableProperties.isTransactional(tableMetadata.getProperties()).orElse(false); + boolean transactional = isTransactional(tableMetadata.getProperties()).orElse(false); tableProperties.put(TRANSACTIONAL, String.valueOf(transactional)); boolean autoPurgeEnabled = HiveTableProperties.isAutoPurge(tableMetadata.getProperties()).orElse(false); tableProperties.put(AUTO_PURGE_KEY, String.valueOf(autoPurgeEnabled)); - bucketProperty.ifPresent(hiveBucketProperty -> - tableProperties.put(BUCKETING_VERSION, Integer.toString(hiveBucketProperty.getBucketingVersion().getVersion()))); + bucketInfo.ifPresent(info -> tableProperties.put(BUCKETING_VERSION, String.valueOf(info.bucketingVersion().getVersion()))); // ORC format specific properties - List columns = getOrcBloomFilterColumns(tableMetadata.getProperties()); - if (columns != null && !columns.isEmpty()) { + List orcBloomFilterColumns = getOrcBloomFilterColumns(tableMetadata.getProperties()); + if (orcBloomFilterColumns != null && !orcBloomFilterColumns.isEmpty()) { checkFormatForProperty(hiveStorageFormat, HiveStorageFormat.ORC, ORC_BLOOM_FILTER_COLUMNS); - validateOrcBloomFilterColumns(tableMetadata, columns); - tableProperties.put(ORC_BLOOM_FILTER_COLUMNS_KEY, Joiner.on(",").join(columns)); + validateOrcBloomFilterColumns(tableMetadata, orcBloomFilterColumns); + tableProperties.put(ORC_BLOOM_FILTER_COLUMNS_KEY, Joiner.on(",").join(orcBloomFilterColumns)); tableProperties.put(ORC_BLOOM_FILTER_FPP_KEY, String.valueOf(getOrcBloomFilterFpp(tableMetadata.getProperties()))); } + List parquetBloomFilterColumns = getParquetBloomFilterColumns(tableMetadata.getProperties()); + if (parquetBloomFilterColumns != null && !parquetBloomFilterColumns.isEmpty()) { + checkFormatForProperty(hiveStorageFormat, HiveStorageFormat.PARQUET, PARQUET_BLOOM_FILTER_COLUMNS); + validateParquetBloomFilterColumns(tableMetadata, parquetBloomFilterColumns); + tableProperties.put(PARQUET_BLOOM_FILTER_COLUMNS_KEY, Joiner.on(",").join(parquetBloomFilterColumns)); + // TODO: Enable specifying FPP + } + // Avro specific properties String avroSchemaUrl = getAvroSchemaUrl(tableMetadata.getProperties()); String avroSchemaLiteral = getAvroSchemaLiteral(tableMetadata.getProperties()); checkAvroSchemaProperties(avroSchemaUrl, avroSchemaLiteral); if (avroSchemaUrl != null) { checkFormatForProperty(hiveStorageFormat, HiveStorageFormat.AVRO, AVRO_SCHEMA_URL); - tableProperties.put(AVRO_SCHEMA_URL_KEY, validateAndNormalizeAvroSchemaUrl(avroSchemaUrl, hdfsContext)); + tableProperties.put(AVRO_SCHEMA_URL_KEY, validateAvroSchemaUrl(session, avroSchemaUrl)); } else if (avroSchemaLiteral != null) { checkFormatForProperty(hiveStorageFormat, HiveStorageFormat.AVRO, AVRO_SCHEMA_LITERAL); @@ -1179,16 +1194,23 @@ else if (avroSchemaLiteral != null) { }); // Set bogus table stats to prevent Hive 2.x from gathering these stats at table creation. - // These stats are not useful by themselves and can take very long time to collect when creating an - // external table over large data set. + // These stats are not useful by themselves and can take a very long time to collect when creating an + // external table over a large data set. tableProperties.put("numFiles", "-1"); tableProperties.put("totalSize", "-1"); // Table comment property - tableMetadata.getComment().ifPresent(value -> tableProperties.put(TABLE_COMMENT, value)); + tableMetadata.getComment().ifPresent(value -> tableProperties.put(Table.TABLE_COMMENT, value)); // Partition Projection specific properties - tableProperties.putAll(partitionProjectionService.getPartitionProjectionHiveTableProperties(tableMetadata)); + if (partitionProjectionEnabled) { + tableProperties.putAll(getPartitionProjectionHiveTableProperties(tableMetadata)); + } + else if (arePartitionProjectionPropertiesSet(tableMetadata)) { + throw new TrinoException( + INVALID_COLUMN_PROPERTY, + "Partition projection is disabled. Enable it in configuration by setting " + "hive.partition-projection-enabled" + "=true"); + } Map baseProperties = tableProperties.buildOrThrow(); @@ -1198,7 +1220,7 @@ else if (avroSchemaLiteral != null) { Set illegalExtraProperties = Sets.intersection( ImmutableSet.builder() .addAll(baseProperties.keySet()) - .addAll(STATS_PROPERTIES) + .addAll(MetastoreUtil.STATS_PROPERTIES) .build(), extraProperties.keySet()); if (!illegalExtraProperties.isEmpty()) { @@ -1237,28 +1259,32 @@ private void validateOrcBloomFilterColumns(ConnectorTableMetadata tableMetadata, } } - private String validateAndNormalizeAvroSchemaUrl(String url, HdfsContext context) + private static void validateParquetBloomFilterColumns(ConnectorTableMetadata tableMetadata, List parquetBloomFilterColumns) { - try { - new URL(url).openStream().close(); - return url; - } - catch (MalformedURLException e) { - // try locally - if (new File(url).exists()) { - // hive needs url to have a protocol - return new File(url).toURI().toString(); + Map columnTypes = tableMetadata.getColumns().stream() + .collect(toImmutableMap(ColumnMetadata::getName, ColumnMetadata::getType)); + for (String column : parquetBloomFilterColumns) { + Type type = columnTypes.get(column); + if (type == null) { + throw new TrinoException(INVALID_TABLE_PROPERTY, format("Parquet Bloom filter column %s not present in schema", column)); } - // try hdfs - try { - if (!hdfsEnvironment.getFileSystem(context, new Path(url)).exists(new Path(url))) { - throw new TrinoException(INVALID_TABLE_PROPERTY, "Cannot locate Avro schema file: " + url); - } - return url; + if (!SUPPORTED_BLOOM_FILTER_TYPES.contains(type)) { + throw new TrinoException(INVALID_TABLE_PROPERTY, format("Parquet Bloom filter column %s has unsupported type %s", column, type.getDisplayName())); } - catch (IOException ex) { - throw new TrinoException(INVALID_TABLE_PROPERTY, "Avro schema file is not a valid file system URI: " + url, ex); + } + } + + private String validateAvroSchemaUrl(ConnectorSession session, String url) + { + try { + Location location = Location.of(url); + if (!fileSystemFactory.create(session).newInputFile(location).exists()) { + throw new TrinoException(INVALID_TABLE_PROPERTY, "Cannot locate Avro schema file: " + url); } + return location.toString(); + } + catch (IllegalArgumentException e) { + throw new TrinoException(INVALID_TABLE_PROPERTY, "Avro schema file is not a valid file system URI: " + url, e); } catch (IOException e) { throw new TrinoException(INVALID_TABLE_PROPERTY, "Cannot open Avro schema file: " + url, e); @@ -1305,17 +1331,30 @@ private static Location getValidatedExternalLocation(String location) return validated; } - private void checkExternalPath(HdfsContext context, Path path) + private void checkExternalPathAndCreateIfNotExists(ConnectorSession session, Location location) { try { - if (!isS3FileSystem(context, hdfsEnvironment, path)) { - if (!hdfsEnvironment.getFileSystem(context, path).isDirectory(path)) { - throw new TrinoException(INVALID_TABLE_PROPERTY, "External location must be a directory: " + path); + if (!fileSystemFactory.create(session).directoryExists(location).orElse(true)) { + if (writesToNonManagedTablesEnabled) { + createDirectory(session, location); + } + else { + throw new TrinoException(INVALID_TABLE_PROPERTY, "External location must be a directory: " + location); } } } + catch (IOException | IllegalArgumentException e) { + throw new TrinoException(INVALID_TABLE_PROPERTY, "External location is not a valid file system URI: " + location, e); + } + } + + private void createDirectory(ConnectorSession session, Location location) + { + try { + fileSystemFactory.create(session).createDirectory(location); + } catch (IOException e) { - throw new TrinoException(INVALID_TABLE_PROPERTY, "External location is not a valid file system URI: " + path, e); + throw new TrinoException(INVALID_TABLE_PROPERTY, e.getMessage()); } } @@ -1334,16 +1373,17 @@ private static Table buildTableObject( List columnHandles, HiveStorageFormat hiveStorageFormat, List partitionedBy, - Optional bucketProperty, + Optional bucketInfo, Map additionalTableParameters, Optional targetPath, boolean external, - String prestoVersion, + String trinoVersion, boolean usingSystemSecurity) { Map columnHandlesByName = Maps.uniqueIndex(columnHandles, HiveColumnHandle::getName); List partitionColumns = partitionedBy.stream() .map(columnHandlesByName::get) + .map(Objects::requireNonNull) .map(HiveColumnHandle::toMetastoreColumn) .collect(toImmutableList()); @@ -1355,7 +1395,7 @@ private static Table buildTableObject( HiveType type = columnHandle.getHiveType(); if (!partitionColumnNames.contains(name)) { verify(!columnHandle.isPartitionKey(), "Column handles are not consistent with partitioned by property"); - columns.add(new Column(name, type, columnHandle.getComment())); + columns.add(new Column(name, type, columnHandle.getComment(), ImmutableMap.of())); } else { verify(columnHandle.isPartitionKey(), "Column handles are not consistent with partitioned by property"); @@ -1363,8 +1403,8 @@ private static Table buildTableObject( } ImmutableMap.Builder tableParameters = ImmutableMap.builder() - .put(PRESTO_VERSION_NAME, prestoVersion) - .put(PRESTO_QUERY_ID_NAME, queryId) + .put(TRINO_VERSION_NAME, trinoVersion) + .put(TRINO_QUERY_ID_NAME, queryId) .putAll(additionalTableParameters); if (external) { @@ -1381,8 +1421,8 @@ private static Table buildTableObject( .setParameters(tableParameters.buildOrThrow()); tableBuilder.getStorageBuilder() - .setStorageFormat(fromHiveStorageFormat(hiveStorageFormat)) - .setBucketProperty(bucketProperty) + .setStorageFormat(hiveStorageFormat.toStorageFormat()) + .setBucketProperty(bucketInfo.map(info -> new HiveBucketProperty(info.bucketedBy(), info.bucketCount(), info.sortedBy()))) .setLocation(targetPath.map(Object::toString)); return tableBuilder.build(); @@ -1492,12 +1532,6 @@ public void setViewColumnComment(ConnectorSession session, SchemaTableName viewN replaceView(session, viewName, view, newDefinition); } - @Override - public void setMaterializedViewColumnComment(ConnectorSession session, SchemaTableName viewName, String columnName, Optional comment) - { - hiveMaterializedViewMetadata.setMaterializedViewColumnComment(session, viewName, columnName, comment); - } - private Table getView(SchemaTableName viewName) { Table view = metastore.getTable(viewName.getSchemaName(), viewName.getTableName()) @@ -1549,6 +1583,8 @@ public ConnectorTableHandle beginStatisticsCollection(ConnectorSession session, @Override public void finishStatisticsCollection(ConnectorSession session, ConnectorTableHandle tableHandle, Collection computedStatistics) { + verify(isStatisticsEnabled(session), "statistics not enabled"); + HiveTableHandle handle = (HiveTableHandle) tableHandle; SchemaTableName tableName = handle.getSchemaTableName(); Table table = metastore.getTable(tableName.getSchemaName(), tableName.getTableName()) @@ -1562,12 +1598,12 @@ public void finishStatisticsCollection(ConnectorSession session, ConnectorTableH List hiveColumnHandles = hiveColumnHandles(table, typeManager, timestampPrecision); Map columnTypes = hiveColumnHandles.stream() .filter(columnHandle -> !columnHandle.isHidden()) - .collect(toImmutableMap(HiveColumnHandle::getName, column -> column.getHiveType().getType(typeManager, timestampPrecision))); + .collect(toImmutableMap(HiveColumnHandle::getName, column -> getType(column.getHiveType(), typeManager, timestampPrecision))); Map, ComputedStatistics> computedStatisticsMap = createComputedStatisticsToPartitionMap(computedStatistics, partitionColumnNames, columnTypes); if (partitionColumns.isEmpty()) { - // commit analyze to unpartitioned table + // commit the analysis result to an unpartitioned table metastore.setTableStatistics(table, createPartitionStatistics(columnTypes, computedStatisticsMap.get(ImmutableList.of()))); } else { @@ -1593,7 +1629,7 @@ public void finishStatisticsCollection(ConnectorSession session, ConnectorTableH Map> columnStatisticTypes = hiveColumnHandles.stream() .filter(columnHandle -> !partitionColumnNames.contains(columnHandle.getName())) .filter(column -> !column.isHidden()) - .collect(toImmutableMap(HiveColumnHandle::getName, column -> ImmutableSet.copyOf(metastore.getSupportedColumnStatistics(column.getType())))); + .collect(toImmutableMap(HiveColumnHandle::getName, column -> getSupportedColumnStatistics(column.getType()))); Supplier emptyPartitionStatistics = Suppliers.memoize(() -> createEmptyPartitionStatistics(columnTypes, columnStatisticTypes)); List partitionTypes = handle.getPartitionColumns().stream() @@ -1678,14 +1714,14 @@ public HiveOutputTableHandle beginCreateTable(ConnectorSession session, Connecto HiveStorageFormat tableStorageFormat = getHiveStorageFormat(tableMetadata.getProperties()); List partitionedBy = getPartitionedBy(tableMetadata.getProperties()); - Optional bucketProperty = getBucketProperty(tableMetadata.getProperties()); + Optional bucketProperty = getBucketInfo(tableMetadata.getProperties()); // get the root directory for the database SchemaTableName schemaTableName = tableMetadata.getTable(); String schemaName = schemaTableName.getSchemaName(); String tableName = schemaTableName.getTableName(); - Map tableProperties = getEmptyTableProperties(tableMetadata, bucketProperty, new HdfsContext(session)); + Map tableProperties = getEmptyTableProperties(tableMetadata, bucketProperty, session); List columnHandles = getColumnHandles(tableMetadata, ImmutableSet.copyOf(partitionedBy)); HiveStorageFormat partitionStorageFormat = isRespectTableFormat(session) ? tableStorageFormat : getHiveStorageFormat(session); @@ -1744,17 +1780,17 @@ public Optional finishCreateTable(ConnectorSession sess handle.getInputColumns(), handle.getTableStorageFormat(), handle.getPartitionedBy(), - handle.getBucketProperty(), + handle.getBucketInfo(), handle.getAdditionalTableParameters(), Optional.of(writeInfo.targetPath()), handle.isExternal(), - prestoVersion, - accessControlMetadata.isUsingSystemSecurity()); - PrincipalPrivileges principalPrivileges = accessControlMetadata.isUsingSystemSecurity() ? NO_PRIVILEGES : buildInitialPrivilegeSet(handle.getTableOwner()); + trinoVersion, + usingSystemSecurity); + PrincipalPrivileges principalPrivileges = usingSystemSecurity ? NO_PRIVILEGES : buildInitialPrivilegeSet(handle.getTableOwner()); partitionUpdates = PartitionUpdate.mergePartitionUpdates(partitionUpdates); - if (handle.getBucketProperty().isPresent() && isCreateEmptyBucketFiles(session)) { + if (handle.getBucketInfo().isPresent() && isCreateEmptyBucketFiles(session)) { List partitionUpdatesForMissingBuckets = computePartitionUpdatesForMissingBuckets(session, handle, true, partitionUpdates); // replace partitionUpdates before creating the empty files so that those files will be cleaned up if we end up rollback partitionUpdates = PartitionUpdate.mergePartitionUpdates(concat(partitionUpdates, partitionUpdatesForMissingBuckets)); @@ -1777,14 +1813,16 @@ public Optional finishCreateTable(ConnectorSession sess } Map columnTypes = handle.getInputColumns().stream() - .collect(toImmutableMap(HiveColumnHandle::getName, column -> column.getHiveType().getType(typeManager))); + .collect(toImmutableMap(HiveColumnHandle::getName, column -> typeManager.getType(getTypeSignature(column.getHiveType())))); Map, ComputedStatistics> partitionComputedStatistics = createComputedStatisticsToPartitionMap(computedStatistics, handle.getPartitionedBy(), columnTypes); PartitionStatistics tableStatistics; if (table.getPartitionColumns().isEmpty()) { HiveBasicStatistics basicStatistics = partitionUpdates.stream() .map(PartitionUpdate::getStatistics) - .reduce((first, second) -> reduce(first, second, ADD)) + .map(hiveBasicStatistics -> new PartitionStatistics(hiveBasicStatistics, ImmutableMap.of())) + .reduce(MERGE_INCREMENTAL::updatePartitionStatistics) + .map(PartitionStatistics::getBasicStatistics) .orElse(createZeroStatistics()); tableStatistics = createPartitionStatistics(basicStatistics, columnTypes, getColumnStatistics(partitionComputedStatistics, ImmutableList.of())); } @@ -1792,7 +1830,7 @@ public Optional finishCreateTable(ConnectorSession sess tableStatistics = new PartitionStatistics(createEmptyStatistics(), ImmutableMap.of()); } - Optional writePath = Optional.of(new Path(writeInfo.writePath().toString())); + Optional writePath = Optional.of(writeInfo.writePath()); if (handle.getPartitionedBy().isEmpty()) { List fileNames; if (partitionUpdates.isEmpty()) { @@ -1845,7 +1883,7 @@ private List computePartitionUpdatesForMissingBuckets( { ImmutableList.Builder partitionUpdatesForMissingBucketsBuilder = ImmutableList.builder(); for (PartitionUpdate partitionUpdate : partitionUpdates) { - int bucketCount = handle.getBucketProperty().get().getBucketCount(); + int bucketCount = handle.getBucketInfo().orElseThrow().bucketCount(); List fileNamesForMissingBuckets = computeFileNamesForMissingBuckets( session, @@ -1897,7 +1935,7 @@ private List computeFileNamesForMissingBuckets( private void createEmptyFiles(ConnectorSession session, Location path, Table table, Optional partition, List fileNames) { - Properties schema; + Map schema; StorageFormat format; if (partition.isPresent()) { schema = getHiveSchema(partition.get(), table); @@ -2053,7 +2091,7 @@ else if (isTransactional) { handles, metastore.generatePageSinkMetadata(tableName), locationHandle, - table.getStorage().getBucketProperty(), + createBucketInfo(table), tableStorageFormat, isRespectTableFormat(session) ? tableStorageFormat : getHiveStorageFormat(session), transaction, @@ -2112,7 +2150,7 @@ private Table finishChangingTable(AcidOperation acidOperation, String changeDesc throw new TrinoException(HIVE_CONCURRENT_MODIFICATION_DETECTED, "Table format changed during " + changeDescription); } - if (handle.getBucketProperty().isPresent() && isCreateEmptyBucketFiles(session)) { + if (handle.getBucketInfo().isPresent() && isCreateEmptyBucketFiles(session)) { List partitionUpdatesForMissingBuckets = computePartitionUpdatesForMissingBuckets(session, handle, false, partitionUpdates); // replace partitionUpdates before creating the empty files so that those files will be cleaned up if we end up rollback partitionUpdates = PartitionUpdate.mergePartitionUpdates(concat(partitionUpdates, partitionUpdatesForMissingBuckets)); @@ -2167,7 +2205,7 @@ private Table finishChangingTable(AcidOperation acidOperation, String changeDesc session, table, principalPrivileges, - Optional.of(new Path(partitionUpdate.getWritePath().toString())), + Optional.of(partitionUpdate.getWritePath()), Optional.of(partitionUpdate.getFileNames()), false, partitionStatistics, @@ -2227,8 +2265,8 @@ else if (partitionUpdate.getUpdateMode() == NEW || partitionUpdate.getUpdateMode if (handle.getLocationHandle().getWriteMode() == DIRECT_TO_TARGET_EXISTING_DIRECTORY) { removeNonCurrentQueryFiles(session, partitionUpdate.getTargetPath()); if (handle.isRetriesEnabled()) { - HdfsContext hdfsContext = new HdfsContext(session); - cleanExtraOutputFiles(hdfsEnvironment, hdfsContext, session.getQueryId(), partitionUpdate.getTargetPath(), ImmutableSet.copyOf(partitionUpdate.getFileNames())); + TrinoFileSystem fileSystem = fileSystemFactory.create(session); + cleanExtraOutputFiles(fileSystem, session.getQueryId(), partitionUpdate.getTargetPath(), ImmutableSet.copyOf(partitionUpdate.getFileNames())); } } else { @@ -2260,14 +2298,13 @@ else if (partitionUpdate.getUpdateMode() == NEW || partitionUpdate.getUpdateMode private void removeNonCurrentQueryFiles(ConnectorSession session, Location partitionLocation) { String queryId = session.getQueryId(); + TrinoFileSystem fileSystem = fileSystemFactory.create(session); try { - Path partitionPath = new Path(partitionLocation.toString()); - FileSystem fileSystem = hdfsEnvironment.getFileSystem(new HdfsContext(session), partitionPath); - RemoteIterator iterator = fileSystem.listFiles(partitionPath, false); + FileIterator iterator = fileSystem.listFiles(partitionLocation); while (iterator.hasNext()) { - Path file = iterator.next().getPath(); - if (!isFileCreatedByQuery(file.getName(), queryId)) { - checkedDelete(fileSystem, file, false); + Location location = iterator.next().location(); + if (!isFileCreatedByQuery(location.fileName(), queryId)) { + fileSystem.deleteFile(location); } } } @@ -2297,8 +2334,8 @@ private Partition buildPartitionObject(ConnectorSession session, Table table, Pa .setColumns(table.getDataColumns()) .setValues(extractPartitionValues(partitionUpdate.getName())) .setParameters(ImmutableMap.builder() - .put(PRESTO_VERSION_NAME, prestoVersion) - .put(PRESTO_QUERY_ID_NAME, session.getQueryId()) + .put(TRINO_VERSION_NAME, trinoVersion) + .put(TRINO_QUERY_ID_NAME, session.getQueryId()) .buildOrThrow()) .withStorage(storage -> storage .setStorageFormat(isRespectTableFormat(session) ? @@ -2427,7 +2464,7 @@ private Optional getTableHandleForOptimize(Connecto columns, metastore.generatePageSinkMetadata(tableName), locationHandle, - table.getStorage().getBucketProperty(), + createBucketInfo(table), tableStorageFormat, // TODO: test with multiple partitions using different storage format tableStorageFormat, @@ -2496,7 +2533,7 @@ private void finishOptimize(ConnectorSession session, ConnectorTableExecuteHandl } // Support for bucketed tables disabled mostly so we do not need to think about grouped execution in an initial version. Possibly no change apart from testing required. - verify(handle.getBucketProperty().isEmpty(), "bucketed table not supported"); + verify(handle.getBucketInfo().isEmpty(), "bucketed table not supported"); ImmutableList.Builder partitionUpdateInfosBuilder = ImmutableList.builder(); for (PartitionUpdate partitionUpdate : partitionUpdates) { @@ -2544,40 +2581,36 @@ private void finishOptimize(ConnectorSession session, ConnectorTableExecuteHandl handle.isRetriesEnabled()); } - // get filesystem - FileSystem fs; - try { - fs = hdfsEnvironment.getFileSystem(new HdfsContext(session), new Path(table.getStorage().getLocation())); - } - catch (IOException e) { - throw new TrinoException(HIVE_FILESYSTEM_ERROR, e); - } - // path to be deleted - Set scannedPaths = splitSourceInfo.stream() - .map(file -> new Path((String) file)) + Set scannedPaths = splitSourceInfo.stream() + .map(file -> Location.of((String) file)) .collect(toImmutableSet()); // track remaining files to be delted for error reporting - Set remainingFilesToDelete = new HashSet<>(scannedPaths); + Set remainingFilesToDelete = new HashSet<>(scannedPaths); // delete loop + TrinoFileSystem fileSystem = fileSystemFactory.create(session); boolean someDeleted = false; - Optional firstScannedPath = Optional.empty(); + Optional firstScannedPath = Optional.empty(); try { - for (Path scannedPath : scannedPaths) { + for (Location scannedPath : scannedPaths) { if (firstScannedPath.isEmpty()) { firstScannedPath = Optional.of(scannedPath); } - retry().run("delete " + scannedPath, () -> { - checkedDelete(fs, scannedPath, false); - return null; + Failsafe.with(DELETE_RETRY_POLICY).run(() -> { + try { + fileSystem.deleteFile(scannedPath); + } + catch (FileNotFoundException e) { + // ignore missing files + } }); someDeleted = true; remainingFilesToDelete.remove(scannedPath); } } catch (Exception e) { - if (!someDeleted && (firstScannedPath.isEmpty() || exists(fs, firstScannedPath.get()))) { + if (!someDeleted && (firstScannedPath.isEmpty() || exists(fileSystem, firstScannedPath.get()))) { // we are good - we did not delete any source files so we can just throw error and allow rollback to happend // if someDeleted flag is false we do extra checkig if first file we tried to delete is still there. There is a chance that // fs.delete above could throw exception but file was actually deleted. @@ -2594,10 +2627,10 @@ private void finishOptimize(ConnectorSession session, ConnectorTableExecuteHandl } } - private boolean exists(FileSystem fs, Path path) + private boolean exists(TrinoFileSystem fs, Location location) { try { - return fs.exists(path); + return fs.newInputFile(location).exists(); } catch (IOException e) { // on failure pessimistically assume file does not exist @@ -2613,14 +2646,14 @@ public void createView(ConnectorSession session, SchemaTableName viewName, Conne } Map properties = ImmutableMap.builder() - .put(TABLE_COMMENT, PRESTO_VIEW_COMMENT) + .put(Table.TABLE_COMMENT, TableInfo.PRESTO_VIEW_COMMENT) .put(PRESTO_VIEW_FLAG, "true") .put(TRINO_CREATED_BY, "Trino Hive connector") - .put(PRESTO_VERSION_NAME, prestoVersion) - .put(PRESTO_QUERY_ID_NAME, session.getQueryId()) + .put(TRINO_VERSION_NAME, trinoVersion) + .put(TRINO_QUERY_ID_NAME, session.getQueryId()) .buildOrThrow(); - Column dummyColumn = new Column("dummy", HIVE_STRING, Optional.empty()); + Column dummyColumn = new Column("dummy", HIVE_STRING, Optional.empty(), ImmutableMap.of()); Table.Builder tableBuilder = Table.builder() .setDatabaseName(viewName.getSchemaName()) @@ -2689,26 +2722,10 @@ public void dropView(ConnectorSession session, SchemaTableName viewName) @Override public List listViews(ConnectorSession session, Optional optionalSchemaName) { - Set materializedViews = ImmutableSet.copyOf(listMaterializedViews(session, optionalSchemaName)); - if (optionalSchemaName.isEmpty()) { - Optional> allViews = metastore.getAllViews(); - if (allViews.isPresent()) { - return allViews.get().stream() - .filter(view -> !isHiveSystemSchema(view.getSchemaName())) - .filter(view -> !materializedViews.contains(view)) - .collect(toImmutableList()); - } - } - ImmutableList.Builder tableNames = ImmutableList.builder(); - for (String schemaName : listSchemas(session, optionalSchemaName)) { - for (String tableName : metastore.getAllViews(schemaName)) { - SchemaTableName schemaTableName = new SchemaTableName(schemaName, tableName); - if (!materializedViews.contains(schemaTableName)) { - tableNames.add(schemaTableName); - } - } - } - return tableNames.build(); + return streamTables(session, optionalSchemaName) + .filter(tableInfo -> tableInfo.extendedRelationType().toRelationType() == RelationType.VIEW) + .map(TableInfo::tableName) + .collect(toImmutableList()); } @Override @@ -2808,7 +2825,6 @@ public Optional getUpdateLayout(ConnectorSession se handle.getBucketingVersion(), handle.getBucketCount(), handle.getHiveTypes(), - handle.getMaxCompatibleBucketCount(), handle.isUsePartitionedBucketing())); } @@ -2892,32 +2908,27 @@ public ConnectorTableProperties getTableProperties(ConnectorSession session, Con } } - Optional tablePartitioning = Optional.empty(); - List> sortingProperties = ImmutableList.of(); - if (hiveTable.getBucketHandle().isPresent()) { - if (isPropagateTableScanSortingProperties(session) && !hiveTable.getBucketHandle().get().getSortedBy().isEmpty()) { + Optional newTablePartitioning = Optional.empty(); + List> newSortingProperties = ImmutableList.of(); + Optional hiveTablePartitioning = hiveTable.getTablePartitioning(); + if (hiveTablePartitioning.isPresent()) { + HiveTablePartitioning tablePartitioning = hiveTablePartitioning.get(); + if (isPropagateTableScanSortingProperties(session) && !tablePartitioning.sortedBy().isEmpty()) { // Populating SortingProperty guarantees to the engine that it is reading pre-sorted input. // We detect compatibility between table and partition level sorted_by properties // and fail the query if there is a mismatch in HiveSplitManager#getPartitionMetadata. // This can lead to incorrect results if a sorted_by property is defined over unsorted files. Map columnHandles = getColumnHandles(session, table); - sortingProperties = hiveTable.getBucketHandle().get().getSortedBy().stream() + newSortingProperties = tablePartitioning.sortedBy().stream() .map(sortingColumn -> new SortingProperty<>( columnHandles.get(sortingColumn.getColumnName()), sortingColumn.getOrder().getSortOrder())) .collect(toImmutableList()); } - if (isBucketExecutionEnabled(session)) { - tablePartitioning = hiveTable.getBucketHandle().map(bucketing -> new ConnectorTablePartitioning( - new HivePartitioningHandle( - bucketing.getBucketingVersion(), - bucketing.getReadBucketCount(), - bucketing.getColumns().stream() - .map(HiveColumnHandle::getHiveType) - .collect(toImmutableList()), - OptionalInt.empty(), - false), - bucketing.getColumns().stream() + if (isBucketExecutionEnabled(session) && tablePartitioning.active()) { + newTablePartitioning = Optional.of(new ConnectorTablePartitioning( + tablePartitioning.partitioningHandle(), + tablePartitioning.columns().stream() .map(ColumnHandle.class::cast) .collect(toImmutableList()))); } @@ -2925,9 +2936,9 @@ public ConnectorTableProperties getTableProperties(ConnectorSession session, Con return new ConnectorTableProperties( predicate, - tablePartitioning, + newTablePartitioning, discretePredicates, - sortingProperties); + newSortingProperties); } @Override @@ -3096,60 +3107,58 @@ private HiveColumnHandle createProjectedColumnHandle(HiveColumnHandle column, Li column.getComment()); } - @Override - public Optional applyTableScanRedirect(ConnectorSession session, ConnectorTableHandle tableHandle) - { - return hiveRedirectionsProvider.getTableScanRedirection(session, (HiveTableHandle) tableHandle); - } - @Override public Optional getCommonPartitioningHandle(ConnectorSession session, ConnectorPartitioningHandle left, ConnectorPartitioningHandle right) { HivePartitioningHandle leftHandle = (HivePartitioningHandle) left; HivePartitioningHandle rightHandle = (HivePartitioningHandle) right; - if (leftHandle.isUsePartitionedBucketing() != rightHandle.isUsePartitionedBucketing()) { + // If either side is using partitioned bucketing, we cannot merge the partitioning handles + if (leftHandle.isUsePartitionedBucketing() || rightHandle.isUsePartitionedBucketing()) { return Optional.empty(); } - if (!leftHandle.getHiveTypes().equals(rightHandle.getHiveTypes())) { + + // If bucketing version or types are different, we cannot merge the partitioning handles + if (leftHandle.getBucketingVersion() != rightHandle.getBucketingVersion()) { return Optional.empty(); } - if (leftHandle.getBucketingVersion() != rightHandle.getBucketingVersion()) { + if (!leftHandle.getHiveTypes().equals(rightHandle.getHiveTypes())) { return Optional.empty(); } + + // If the bucket count is the same, the partitioning handles are the same if (leftHandle.getBucketCount() == rightHandle.getBucketCount()) { return Optional.of(leftHandle); } - if (!isOptimizedMismatchedBucketCount(session)) { - return Optional.empty(); - } - int largerBucketCount = Math.max(leftHandle.getBucketCount(), rightHandle.getBucketCount()); - int smallerBucketCount = Math.min(leftHandle.getBucketCount(), rightHandle.getBucketCount()); - if (largerBucketCount % smallerBucketCount != 0) { - // must be evenly divisible - return Optional.empty(); - } - if (Integer.bitCount(largerBucketCount / smallerBucketCount) != 1) { - // ratio must be power of two + // At this point the handles only differ in bucket count + if (!isOptimizedMismatchedBucketCount(session)) { return Optional.empty(); } - OptionalInt maxCompatibleBucketCount = min(leftHandle.getMaxCompatibleBucketCount(), rightHandle.getMaxCompatibleBucketCount()); - if (maxCompatibleBucketCount.isPresent() && maxCompatibleBucketCount.getAsInt() < smallerBucketCount) { - // maxCompatibleBucketCount must be larger than or equal to smallerBucketCount - // because the current code uses the smallerBucketCount as the common partitioning handle. + OptionalInt commonBucketCount = getCommonBucketCount(leftHandle.getBucketCount(), rightHandle.getBucketCount()); + if (commonBucketCount.isEmpty()) { return Optional.empty(); } return Optional.of(new HivePartitioningHandle( - leftHandle.getBucketingVersion(), // same as rightHandle.getBucketingVersion() - smallerBucketCount, + leftHandle.getBucketingVersion(), + commonBucketCount.getAsInt(), leftHandle.getHiveTypes(), - maxCompatibleBucketCount, false)); } + private static OptionalInt getCommonBucketCount(int left, int right) + { + int largerBucketCount = Math.max(left, right); + int smallerBucketCount = Math.min(left, right); + + if (largerBucketCount % smallerBucketCount != 0 || Integer.bitCount(largerBucketCount / smallerBucketCount) != 1) { + return OptionalInt.empty(); + } + return OptionalInt.of(smallerBucketCount); + } + private static OptionalInt min(OptionalInt left, OptionalInt right) { if (left.isEmpty()) { @@ -3167,16 +3176,16 @@ public ConnectorTableHandle makeCompatiblePartitioning(ConnectorSession session, HiveTableHandle hiveTable = (HiveTableHandle) tableHandle; HivePartitioningHandle hivePartitioningHandle = (HivePartitioningHandle) partitioningHandle; - checkArgument(hiveTable.getBucketHandle().isPresent(), "Hive connector only provides alternative layout for bucketed table"); - HiveBucketHandle bucketHandle = hiveTable.getBucketHandle().get(); - ImmutableList bucketTypes = bucketHandle.getColumns().stream().map(HiveColumnHandle::getHiveType).collect(toImmutableList()); + checkArgument(hiveTable.getTablePartitioning().isPresent(), "Hive connector only provides alternative layout for bucketed table"); + HiveTablePartitioning bucketHandle = hiveTable.getTablePartitioning().get(); + ImmutableList bucketTypes = bucketHandle.columns().stream().map(HiveColumnHandle::getHiveType).collect(toImmutableList()); checkArgument( hivePartitioningHandle.getHiveTypes().equals(bucketTypes), "Types from the new PartitioningHandle (%s) does not match the TableHandle (%s)", hivePartitioningHandle.getHiveTypes(), bucketTypes); - int largerBucketCount = Math.max(bucketHandle.getTableBucketCount(), hivePartitioningHandle.getBucketCount()); - int smallerBucketCount = Math.min(bucketHandle.getTableBucketCount(), hivePartitioningHandle.getBucketCount()); + int largerBucketCount = Math.max(bucketHandle.tableBucketCount(), hivePartitioningHandle.getBucketCount()); + int smallerBucketCount = Math.min(bucketHandle.tableBucketCount(), hivePartitioningHandle.getBucketCount()); checkArgument( largerBucketCount % smallerBucketCount == 0 && Integer.bitCount(largerBucketCount / smallerBucketCount) == 1, "The requested partitioning is not a valid alternative for the table layout"); @@ -3191,12 +3200,7 @@ public ConnectorTableHandle makeCompatiblePartitioning(ConnectorSession session, hiveTable.getPartitions(), hiveTable.getCompactEffectivePredicate(), hiveTable.getEnforcedConstraint(), - Optional.of(new HiveBucketHandle( - bucketHandle.getColumns(), - bucketHandle.getBucketingVersion(), - bucketHandle.getTableBucketCount(), - hivePartitioningHandle.getBucketCount(), - bucketHandle.getSortedBy())), + hiveTable.getTablePartitioning(), hiveTable.getBucketFilter(), hiveTable.getAnalyzePartitionValues(), ImmutableSet.of(), @@ -3279,14 +3283,14 @@ public Optional getInsertLayout(ConnectorSession session, // Note: we cannot use hiveTableHandle.isInAcidTransaction() here as transaction is not yet set in HiveTableHandle when getInsertLayout is called else if (isFullAcidTable(table.getParameters())) { table = Table.builder(table) - .withStorage(storage -> storage.setBucketProperty(Optional.of( - new HiveBucketProperty(ImmutableList.of(), HiveBucketing.BucketingVersion.BUCKETING_V2, 1, ImmutableList.of())))) + .setParameter(BUCKETING_VERSION, String.valueOf(BUCKETING_V2.getVersion())) + .withStorage(storage -> storage.setBucketProperty(Optional.of(new HiveBucketProperty(ImmutableList.of(), 1, ImmutableList.of())))) .build(); } - Optional hiveBucketHandle = getHiveBucketHandle(session, table, typeManager); + Optional tablePartitioning = getHiveTablePartitioningForWrite(session, table, typeManager); List partitionColumns = table.getPartitionColumns(); - if (hiveBucketHandle.isEmpty()) { + if (tablePartitioning.isEmpty()) { // return preferred layout which is partitioned by partition columns if (partitionColumns.isEmpty()) { return Optional.empty(); @@ -3304,7 +3308,7 @@ else if (isFullAcidTable(table.getParameters())) { } ImmutableList.Builder partitioningColumns = ImmutableList.builder(); - hiveBucketHandle.get().getColumns().stream() + tablePartitioning.get().columns().stream() .map(HiveColumnHandle::getName) .forEach(partitioningColumns::add); partitionColumns.stream() @@ -3315,15 +3319,7 @@ else if (isFullAcidTable(table.getParameters())) { // per partition. boolean multipleWritersPerPartitionSupported = !isTransactionalTable(table.getParameters()); - HivePartitioningHandle partitioningHandle = new HivePartitioningHandle( - hiveBucketHandle.get().getBucketingVersion(), - hiveBucketHandle.get().getTableBucketCount(), - hiveBucketHandle.get().getColumns().stream() - .map(HiveColumnHandle::getHiveType) - .collect(toImmutableList()), - OptionalInt.of(hiveBucketHandle.get().getTableBucketCount()), - !partitionColumns.isEmpty() && isParallelPartitionedBucketedWrites(session)); - return Optional.of(new ConnectorTableLayout(partitioningHandle, partitioningColumns.build(), multipleWritersPerPartitionSupported)); + return Optional.of(new ConnectorTableLayout(tablePartitioning.get().partitioningHandle(), partitioningColumns.build(), multipleWritersPerPartitionSupported)); } @Override @@ -3333,9 +3329,9 @@ public Optional getNewTableLayout(ConnectorSession session validatePartitionColumns(tableMetadata); validateBucketColumns(tableMetadata); validateColumns(tableMetadata); - Optional bucketProperty = getBucketProperty(tableMetadata.getProperties()); + Optional bucketInfo = getBucketInfo(tableMetadata.getProperties()); List partitionedBy = getPartitionedBy(tableMetadata.getProperties()); - if (bucketProperty.isEmpty()) { + if (bucketInfo.isEmpty()) { // return preferred layout which is partitioned by partition columns if (partitionedBy.isEmpty()) { return Optional.empty(); @@ -3343,7 +3339,7 @@ public Optional getNewTableLayout(ConnectorSession session return Optional.of(new ConnectorTableLayout(partitionedBy)); } - if (!bucketProperty.get().getSortedBy().isEmpty() && !isSortedWritingEnabled(session)) { + if (!bucketInfo.get().sortedBy().isEmpty() && !isSortedWritingEnabled(session)) { throw new TrinoException(NOT_SUPPORTED, "Writing to bucketed sorted Hive tables is disabled"); } @@ -3351,17 +3347,16 @@ public Optional getNewTableLayout(ConnectorSession session // per partition. boolean multipleWritersPerPartitionSupported = !isTransactional(tableMetadata.getProperties()).orElse(false); - List bucketedBy = bucketProperty.get().getBucketedBy(); + List bucketedBy = bucketInfo.get().bucketedBy(); Map hiveTypeMap = tableMetadata.getColumns().stream() .collect(toMap(ColumnMetadata::getName, column -> toHiveType(column.getType()))); return Optional.of(new ConnectorTableLayout( new HivePartitioningHandle( - bucketProperty.get().getBucketingVersion(), - bucketProperty.get().getBucketCount(), + bucketInfo.get().bucketingVersion(), + bucketInfo.get().bucketCount(), bucketedBy.stream() .map(hiveTypeMap::get) .collect(toImmutableList()), - OptionalInt.of(bucketProperty.get().getBucketCount()), !partitionedBy.isEmpty() && isParallelPartitionedBucketedWrites(session)), ImmutableList.builder() .addAll(bucketedBy) @@ -3416,7 +3411,7 @@ private TableStatisticsMetadata getStatisticsCollectionMetadata(List !partitionedBy.contains(column.getName())) .filter(column -> !column.isHidden()) .filter(column -> analyzeColumns.isEmpty() || analyzeColumns.get().contains(column.getName())) - .map(this::getColumnStatisticMetadata) + .map(HiveMetadata::getColumnStatisticMetadata) .flatMap(List::stream) .collect(toImmutableSet()); @@ -3424,14 +3419,10 @@ private TableStatisticsMetadata getStatisticsCollectionMetadata(List getColumnStatisticMetadata(ColumnMetadata columnMetadata) - { - return getColumnStatisticMetadata(columnMetadata.getName(), metastore.getSupportedColumnStatistics(columnMetadata.getType())); - } - - private List getColumnStatisticMetadata(String columnName, Set statisticTypes) + private static List getColumnStatisticMetadata(ColumnMetadata columnMetadata) { - return statisticTypes.stream() + String columnName = columnMetadata.getName(); + return getSupportedColumnStatistics(columnMetadata.getType()).stream() .map(type -> type.createColumnStatisticMetadata(columnName)) .collect(toImmutableList()); } @@ -3540,15 +3531,15 @@ public static HiveStorageFormat extractHiveStorageFormat(StorageFormat storageFo private static void validateBucketColumns(ConnectorTableMetadata tableMetadata) { - Optional bucketProperty = getBucketProperty(tableMetadata.getProperties()); - if (bucketProperty.isEmpty()) { + Optional bucketInfo = getBucketInfo(tableMetadata.getProperties()); + if (bucketInfo.isEmpty()) { return; } Set allColumns = tableMetadata.getColumns().stream() .map(ColumnMetadata::getName) .collect(toImmutableSet()); - List bucketedBy = bucketProperty.get().getBucketedBy(); + List bucketedBy = bucketInfo.get().bucketedBy(); if (!allColumns.containsAll(bucketedBy)) { throw new TrinoException(INVALID_TABLE_PROPERTY, format("Bucketing columns %s not present in schema", Sets.difference(ImmutableSet.copyOf(bucketedBy), allColumns))); } @@ -3558,7 +3549,7 @@ private static void validateBucketColumns(ConnectorTableMetadata tableMetadata) throw new TrinoException(INVALID_TABLE_PROPERTY, format("Bucketing columns %s are also used as partitioning columns", Sets.intersection(ImmutableSet.copyOf(bucketedBy), partitionColumns))); } - List sortedBy = bucketProperty.get().getSortedBy().stream() + List sortedBy = bucketInfo.get().sortedBy().stream() .map(SortingColumn::getColumnName) .collect(toImmutableList()); if (!allColumns.containsAll(sortedBy)) { @@ -3570,9 +3561,9 @@ private static void validateBucketColumns(ConnectorTableMetadata tableMetadata) } List dataColumns = tableMetadata.getColumns().stream() - .map(columnMetadata -> new Column(columnMetadata.getName(), toHiveType(columnMetadata.getType()), Optional.ofNullable(columnMetadata.getComment()))) + .map(columnMetadata -> new Column(columnMetadata.getName(), toHiveType(columnMetadata.getType()), Optional.ofNullable(columnMetadata.getComment()), ImmutableMap.of())) .collect(toImmutableList()); - if (!isSupportedBucketing(bucketProperty.get(), dataColumns, tableMetadata.getTable().getTableName())) { + if (!isSupportedBucketing(bucketInfo.get().bucketedBy(), dataColumns, tableMetadata.getTable().getTableName())) { throw new TrinoException(NOT_SUPPORTED, "Cannot create a table bucketed on an unsupported type"); } } @@ -3724,66 +3715,6 @@ public void cleanupQuery(ConnectorSession session) metastore.cleanupQuery(session); } - @Override - public void createMaterializedView(ConnectorSession session, SchemaTableName viewName, ConnectorMaterializedViewDefinition definition, boolean replace, boolean ignoreExisting) - { - hiveMaterializedViewMetadata.createMaterializedView(session, viewName, definition, replace, ignoreExisting); - } - - @Override - public void dropMaterializedView(ConnectorSession session, SchemaTableName viewName) - { - hiveMaterializedViewMetadata.dropMaterializedView(session, viewName); - } - - @Override - public List listMaterializedViews(ConnectorSession session, Optional schemaName) - { - return hiveMaterializedViewMetadata.listMaterializedViews(session, schemaName); - } - - @Override - public Map getMaterializedViews(ConnectorSession session, Optional schemaName) - { - return hiveMaterializedViewMetadata.getMaterializedViews(session, schemaName); - } - - @Override - public Optional getMaterializedView(ConnectorSession session, SchemaTableName viewName) - { - return hiveMaterializedViewMetadata.getMaterializedView(session, viewName); - } - - @Override - public MaterializedViewFreshness getMaterializedViewFreshness(ConnectorSession session, SchemaTableName name) - { - return hiveMaterializedViewMetadata.getMaterializedViewFreshness(session, name); - } - - @Override - public void renameMaterializedView(ConnectorSession session, SchemaTableName source, SchemaTableName target) - { - hiveMaterializedViewMetadata.renameMaterializedView(session, source, target); - } - - @Override - public boolean delegateMaterializedViewRefreshToConnector(ConnectorSession session, SchemaTableName viewName) - { - return hiveMaterializedViewMetadata.delegateMaterializedViewRefreshToConnector(session, viewName); - } - - @Override - public CompletableFuture refreshMaterializedView(ConnectorSession session, SchemaTableName name) - { - return hiveMaterializedViewMetadata.refreshMaterializedView(session, name); - } - - @Override - public void setMaterializedViewProperties(ConnectorSession session, SchemaTableName viewName, Map> properties) - { - hiveMaterializedViewMetadata.setMaterializedViewProperties(session, viewName, properties); - } - @Override public Optional redirectTable(ConnectorSession session, SchemaTableName tableName) { diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveMetadataFactory.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveMetadataFactory.java index 991fdd36310b..e3cf63c6d917 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveMetadataFactory.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveMetadataFactory.java @@ -19,15 +19,14 @@ import io.airlift.json.JsonCodec; import io.airlift.units.Duration; import io.trino.filesystem.TrinoFileSystemFactory; -import io.trino.hdfs.HdfsEnvironment; import io.trino.plugin.base.CatalogName; -import io.trino.plugin.hive.aws.athena.PartitionProjectionService; import io.trino.plugin.hive.fs.DirectoryLister; import io.trino.plugin.hive.fs.TransactionScopeCachingDirectoryListerFactory; -import io.trino.plugin.hive.metastore.HiveMetastoreConfig; +import io.trino.plugin.hive.metastore.HiveMetastore; import io.trino.plugin.hive.metastore.HiveMetastoreFactory; import io.trino.plugin.hive.metastore.SemiTransactionalHiveMetastore; import io.trino.plugin.hive.security.AccessControlMetadataFactory; +import io.trino.plugin.hive.security.UsingSystemSecurity; import io.trino.plugin.hive.statistics.MetastoreHiveStatisticsProvider; import io.trino.spi.connector.MetadataProvider; import io.trino.spi.security.ConnectorIdentity; @@ -40,7 +39,7 @@ import java.util.concurrent.ScheduledExecutorService; import static com.google.common.util.concurrent.MoreExecutors.directExecutor; -import static io.trino.plugin.hive.metastore.cache.CachingHiveMetastore.memoizeMetastore; +import static io.trino.plugin.hive.metastore.cache.CachingHiveMetastore.createPerTransactionCache; import static java.util.Objects.requireNonNull; public class HiveMetadataFactory @@ -59,7 +58,6 @@ public class HiveMetadataFactory private final HiveMetastoreFactory metastoreFactory; private final Set fileWriterFactories; private final TrinoFileSystemFactory fileSystemFactory; - private final HdfsEnvironment hdfsEnvironment; private final HivePartitionManager partitionManager; private final TypeManager typeManager; private final MetadataProvider metadataProvider; @@ -70,27 +68,26 @@ public class HiveMetadataFactory private final Executor updateExecutor; private final long maxPartitionDropsPerQuery; private final String trinoVersion; - private final HiveRedirectionsProvider hiveRedirectionsProvider; private final Set systemTableProviders; - private final HiveMaterializedViewMetadataFactory hiveMaterializedViewMetadataFactory; private final AccessControlMetadataFactory accessControlMetadataFactory; private final Optional hiveTransactionHeartbeatInterval; private final ScheduledExecutorService heartbeatService; private final DirectoryLister directoryLister; private final TransactionScopeCachingDirectoryListerFactory transactionScopeCachingDirectoryListerFactory; - private final PartitionProjectionService partitionProjectionService; + private final boolean usingSystemSecurity; + private final boolean partitionProjectionEnabled; private final boolean allowTableRename; private final HiveTimestampPrecision hiveViewsTimestampPrecision; + private final Executor metadataFetchingExecutor; @Inject public HiveMetadataFactory( CatalogName catalogName, HiveConfig hiveConfig, - HiveMetastoreConfig hiveMetastoreConfig, + @HideDeltaLakeTables boolean hideDeltaLakeTables, HiveMetastoreFactory metastoreFactory, Set fileWriterFactories, TrinoFileSystemFactory fileSystemFactory, - HdfsEnvironment hdfsEnvironment, HivePartitionManager partitionManager, ExecutorService executorService, @ForHiveTransactionHeartbeats ScheduledExecutorService heartbeatService, @@ -99,13 +96,11 @@ public HiveMetadataFactory( LocationService locationService, JsonCodec partitionUpdateCodec, NodeVersion nodeVersion, - HiveRedirectionsProvider hiveRedirectionsProvider, Set systemTableProviders, - HiveMaterializedViewMetadataFactory hiveMaterializedViewMetadataFactory, AccessControlMetadataFactory accessControlMetadataFactory, DirectoryLister directoryLister, TransactionScopeCachingDirectoryListerFactory transactionScopeCachingDirectoryListerFactory, - PartitionProjectionService partitionProjectionService, + @UsingSystemSecurity boolean usingSystemSecurity, @AllowHiveTableRename boolean allowTableRename) { this( @@ -113,7 +108,6 @@ public HiveMetadataFactory( metastoreFactory, fileWriterFactories, fileSystemFactory, - hdfsEnvironment, partitionManager, hiveConfig.getMaxConcurrentFileSystemOperations(), hiveConfig.getMaxConcurrentMetastoreDrops(), @@ -128,7 +122,7 @@ public HiveMetadataFactory( hiveConfig.isHiveViewsRunAsInvoker(), hiveConfig.getPerTransactionMetastoreCacheMaximumSize(), hiveConfig.getHiveTransactionHeartbeatInterval(), - hiveMetastoreConfig.isHideDeltaLakeTables(), + hideDeltaLakeTables, typeManager, metadataProvider, locationService, @@ -136,15 +130,15 @@ public HiveMetadataFactory( executorService, heartbeatService, nodeVersion.toString(), - hiveRedirectionsProvider, systemTableProviders, - hiveMaterializedViewMetadataFactory, accessControlMetadataFactory, directoryLister, transactionScopeCachingDirectoryListerFactory, - partitionProjectionService, + usingSystemSecurity, + hiveConfig.isPartitionProjectionEnabled(), allowTableRename, - hiveConfig.getTimestampPrecision()); + hiveConfig.getTimestampPrecision(), + hiveConfig.getMetadataParallelism()); } public HiveMetadataFactory( @@ -152,7 +146,6 @@ public HiveMetadataFactory( HiveMetastoreFactory metastoreFactory, Set fileWriterFactories, TrinoFileSystemFactory fileSystemFactory, - HdfsEnvironment hdfsEnvironment, HivePartitionManager partitionManager, int maxConcurrentFileSystemOperations, int maxConcurrentMetastoreDrops, @@ -175,15 +168,15 @@ public HiveMetadataFactory( ExecutorService executorService, ScheduledExecutorService heartbeatService, String trinoVersion, - HiveRedirectionsProvider hiveRedirectionsProvider, Set systemTableProviders, - HiveMaterializedViewMetadataFactory hiveMaterializedViewMetadataFactory, AccessControlMetadataFactory accessControlMetadataFactory, DirectoryLister directoryLister, TransactionScopeCachingDirectoryListerFactory transactionScopeCachingDirectoryListerFactory, - PartitionProjectionService partitionProjectionService, + boolean usingSystemSecurity, + boolean partitionProjectionEnabled, boolean allowTableRename, - HiveTimestampPrecision hiveViewsTimestampPrecision) + HiveTimestampPrecision hiveViewsTimestampPrecision, + int metadataParallelism) { this.catalogName = requireNonNull(catalogName, "catalogName is null"); this.skipDeletionForAlter = skipDeletionForAlter; @@ -199,16 +192,13 @@ public HiveMetadataFactory( this.metastoreFactory = requireNonNull(metastoreFactory, "metastoreFactory is null"); this.fileWriterFactories = ImmutableSet.copyOf(requireNonNull(fileWriterFactories, "fileWriterFactories is null")); this.fileSystemFactory = requireNonNull(fileSystemFactory, "fileSystemFactory is null"); - this.hdfsEnvironment = requireNonNull(hdfsEnvironment, "hdfsEnvironment is null"); this.partitionManager = requireNonNull(partitionManager, "partitionManager is null"); this.typeManager = requireNonNull(typeManager, "typeManager is null"); this.metadataProvider = requireNonNull(metadataProvider, "metadataProvider is null"); this.locationService = requireNonNull(locationService, "locationService is null"); this.partitionUpdateCodec = requireNonNull(partitionUpdateCodec, "partitionUpdateCodec is null"); this.trinoVersion = requireNonNull(trinoVersion, "trinoVersion is null"); - this.hiveRedirectionsProvider = requireNonNull(hiveRedirectionsProvider, "hiveRedirectionsProvider is null"); this.systemTableProviders = requireNonNull(systemTableProviders, "systemTableProviders is null"); - this.hiveMaterializedViewMetadataFactory = requireNonNull(hiveMaterializedViewMetadataFactory, "hiveMaterializedViewMetadataFactory is null"); this.accessControlMetadataFactory = requireNonNull(accessControlMetadataFactory, "accessControlMetadataFactory is null"); this.hiveTransactionHeartbeatInterval = requireNonNull(hiveTransactionHeartbeatInterval, "hiveTransactionHeartbeatInterval is null"); @@ -225,21 +215,29 @@ public HiveMetadataFactory( this.heartbeatService = requireNonNull(heartbeatService, "heartbeatService is null"); this.directoryLister = requireNonNull(directoryLister, "directoryLister is null"); this.transactionScopeCachingDirectoryListerFactory = requireNonNull(transactionScopeCachingDirectoryListerFactory, "transactionScopeCachingDirectoryListerFactory is null"); - this.partitionProjectionService = requireNonNull(partitionProjectionService, "partitionProjectionService is null"); + this.usingSystemSecurity = usingSystemSecurity; + this.partitionProjectionEnabled = partitionProjectionEnabled; this.allowTableRename = allowTableRename; this.hiveViewsTimestampPrecision = requireNonNull(hiveViewsTimestampPrecision, "hiveViewsTimestampPrecision is null"); + if (metadataParallelism == 1) { + this.metadataFetchingExecutor = directExecutor(); + } + else { + this.metadataFetchingExecutor = new BoundedExecutor(executorService, metadataParallelism); + } } @Override public TransactionalMetadata create(ConnectorIdentity identity, boolean autoCommit) { - HiveMetastoreClosure hiveMetastoreClosure = new HiveMetastoreClosure( - memoizeMetastore(metastoreFactory.createMetastore(Optional.of(identity)), perTransactionCacheMaximumSize)); // per-transaction cache + HiveMetastore hiveMetastore = createPerTransactionCache(metastoreFactory.createMetastore(Optional.of(identity)), perTransactionCacheMaximumSize); DirectoryLister directoryLister = transactionScopeCachingDirectoryListerFactory.get(this.directoryLister); SemiTransactionalHiveMetastore metastore = new SemiTransactionalHiveMetastore( - hdfsEnvironment, - hiveMetastoreClosure, + typeManager, + partitionProjectionEnabled, + fileSystemFactory, + hiveMetastore, fileSystemExecutor, dropExecutor, updateExecutor, @@ -256,7 +254,6 @@ public TransactionalMetadata create(ConnectorIdentity identity, boolean autoComm autoCommit, fileWriterFactories, fileSystemFactory, - hdfsEnvironment, partitionManager, writesToNonManagedTablesEnabled, createsOfNonManagedTablesEnabled, @@ -269,14 +266,14 @@ public TransactionalMetadata create(ConnectorIdentity identity, boolean autoComm partitionUpdateCodec, trinoVersion, new MetastoreHiveStatisticsProvider(metastore), - hiveRedirectionsProvider, systemTableProviders, - hiveMaterializedViewMetadataFactory.create(hiveMetastoreClosure), accessControlMetadataFactory.create(metastore), directoryLister, - partitionProjectionService, + usingSystemSecurity, + partitionProjectionEnabled, allowTableRename, maxPartitionDropsPerQuery, - hiveViewsTimestampPrecision); + hiveViewsTimestampPrecision, + metadataFetchingExecutor); } } diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveMetastoreClosure.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveMetastoreClosure.java deleted file mode 100644 index e73a616b7ee0..000000000000 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveMetastoreClosure.java +++ /dev/null @@ -1,417 +0,0 @@ -/* - * 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.hive; - -import com.google.common.collect.ImmutableList; -import io.trino.hive.thrift.metastore.DataOperationType; -import io.trino.plugin.hive.acid.AcidOperation; -import io.trino.plugin.hive.acid.AcidTransaction; -import io.trino.plugin.hive.metastore.AcidTransactionOwner; -import io.trino.plugin.hive.metastore.Column; -import io.trino.plugin.hive.metastore.Database; -import io.trino.plugin.hive.metastore.HiveMetastore; -import io.trino.plugin.hive.metastore.HivePrincipal; -import io.trino.plugin.hive.metastore.HivePrivilegeInfo; -import io.trino.plugin.hive.metastore.HivePrivilegeInfo.HivePrivilege; -import io.trino.plugin.hive.metastore.Partition; -import io.trino.plugin.hive.metastore.PartitionWithStatistics; -import io.trino.plugin.hive.metastore.PrincipalPrivileges; -import io.trino.plugin.hive.metastore.Table; -import io.trino.spi.connector.SchemaTableName; -import io.trino.spi.connector.TableNotFoundException; -import io.trino.spi.predicate.TupleDomain; -import io.trino.spi.security.RoleGrant; -import io.trino.spi.type.Type; - -import java.util.List; -import java.util.Map; -import java.util.Optional; -import java.util.OptionalLong; -import java.util.Set; -import java.util.function.Function; - -import static com.google.common.collect.ImmutableList.toImmutableList; -import static com.google.common.collect.ImmutableMap.toImmutableMap; -import static com.google.common.collect.Maps.immutableEntry; -import static io.trino.plugin.hive.HivePartitionManager.extractPartitionValues; -import static java.util.Objects.requireNonNull; - -public class HiveMetastoreClosure -{ - private final HiveMetastore delegate; - - /** - * Do not use this directly. Instead, the closure should be fetched from the current SemiTransactionalHiveMetastore, - * which can be fetched from the current HiveMetadata. - */ - public HiveMetastoreClosure(HiveMetastore delegate) - { - this.delegate = requireNonNull(delegate, "delegate is null"); - } - - public Optional getDatabase(String databaseName) - { - return delegate.getDatabase(databaseName); - } - - public List getAllDatabases() - { - return delegate.getAllDatabases(); - } - - private Table getExistingTable(String databaseName, String tableName) - { - return delegate.getTable(databaseName, tableName) - .orElseThrow(() -> new TableNotFoundException(new SchemaTableName(databaseName, tableName))); - } - - public Optional
getTable(String databaseName, String tableName) - { - return delegate.getTable(databaseName, tableName); - } - - public Set getSupportedColumnStatistics(Type type) - { - return delegate.getSupportedColumnStatistics(type); - } - - public PartitionStatistics getTableStatistics(String databaseName, String tableName, Optional> columns) - { - Table table = getExistingTable(databaseName, tableName); - if (columns.isPresent()) { - Set requestedColumnNames = columns.get(); - List requestedColumns = table.getDataColumns().stream() - .filter(column -> requestedColumnNames.contains(column.getName())) - .collect(toImmutableList()); - table = Table.builder(table).setDataColumns(requestedColumns).build(); - } - return delegate.getTableStatistics(table); - } - - public Map getPartitionStatistics(String databaseName, String tableName, Set partitionNames) - { - return getPartitionStatistics(databaseName, tableName, partitionNames, Optional.empty()); - } - - public Map getPartitionStatistics(String databaseName, String tableName, Set partitionNames, Optional> columns) - { - Table table = getExistingTable(databaseName, tableName); - List partitions = getExistingPartitionsByNames(table, ImmutableList.copyOf(partitionNames)); - if (columns.isPresent()) { - Set requestedColumnNames = columns.get(); - List requestedColumns = table.getDataColumns().stream() - .filter(column -> requestedColumnNames.contains(column.getName())) - .collect(toImmutableList()); - table = Table.builder(table).setDataColumns(requestedColumns).build(); - partitions = partitions.stream().map(partition -> Partition.builder(partition).setColumns(requestedColumns).build()).collect(toImmutableList()); - } - return delegate.getPartitionStatistics(table, partitions); - } - - public void updateTableStatistics(String databaseName, - String tableName, - AcidTransaction transaction, - Function update) - { - delegate.updateTableStatistics(databaseName, tableName, transaction, update); - } - - public void updatePartitionStatistics(String databaseName, - String tableName, - String partitionName, - Function update) - { - Table table = getExistingTable(databaseName, tableName); - delegate.updatePartitionStatistics(table, partitionName, update); - } - - public void updatePartitionStatistics(String databaseName, String tableName, Map> updates) - { - Table table = getExistingTable(databaseName, tableName); - delegate.updatePartitionStatistics(table, updates); - } - - public List getAllTables(String databaseName) - { - return delegate.getAllTables(databaseName); - } - - public Optional> getAllTables() - { - return delegate.getAllTables(); - } - - public List getAllViews(String databaseName) - { - return delegate.getAllViews(databaseName); - } - - public Optional> getAllViews() - { - return delegate.getAllViews(); - } - - public void createDatabase(Database database) - { - delegate.createDatabase(database); - } - - public void dropDatabase(String databaseName, boolean deleteData) - { - delegate.dropDatabase(databaseName, deleteData); - } - - public void renameDatabase(String databaseName, String newDatabaseName) - { - delegate.renameDatabase(databaseName, newDatabaseName); - } - - public void setDatabaseOwner(String databaseName, HivePrincipal principal) - { - delegate.setDatabaseOwner(databaseName, principal); - } - - public void setTableOwner(String databaseName, String tableName, HivePrincipal principal) - { - delegate.setTableOwner(databaseName, tableName, principal); - } - - public void createTable(Table table, PrincipalPrivileges principalPrivileges) - { - delegate.createTable(table, principalPrivileges); - } - - public void dropTable(String databaseName, String tableName, boolean deleteData) - { - delegate.dropTable(databaseName, tableName, deleteData); - } - - public void replaceTable(String databaseName, String tableName, Table newTable, PrincipalPrivileges principalPrivileges) - { - delegate.replaceTable(databaseName, tableName, newTable, principalPrivileges); - } - - public void renameTable(String databaseName, String tableName, String newDatabaseName, String newTableName) - { - delegate.renameTable(databaseName, tableName, newDatabaseName, newTableName); - } - - public void commentTable(String databaseName, String tableName, Optional comment) - { - delegate.commentTable(databaseName, tableName, comment); - } - - public void commentColumn(String databaseName, String tableName, String columnName, Optional comment) - { - delegate.commentColumn(databaseName, tableName, columnName, comment); - } - - public void addColumn(String databaseName, String tableName, String columnName, HiveType columnType, String columnComment) - { - delegate.addColumn(databaseName, tableName, columnName, columnType, columnComment); - } - - public void renameColumn(String databaseName, String tableName, String oldColumnName, String newColumnName) - { - delegate.renameColumn(databaseName, tableName, oldColumnName, newColumnName); - } - - public void dropColumn(String databaseName, String tableName, String columnName) - { - delegate.dropColumn(databaseName, tableName, columnName); - } - - public Optional getPartition(String databaseName, String tableName, List partitionValues) - { - return delegate.getTable(databaseName, tableName) - .flatMap(table -> delegate.getPartition(table, partitionValues)); - } - - public Optional> getPartitionNamesByFilter( - String databaseName, - String tableName, - List columnNames, - TupleDomain partitionKeysFilter) - { - return delegate.getPartitionNamesByFilter(databaseName, tableName, columnNames, partitionKeysFilter); - } - - private List getExistingPartitionsByNames(Table table, List partitionNames) - { - Map partitions = delegate.getPartitionsByNames(table, partitionNames).entrySet().stream() - .map(entry -> immutableEntry(entry.getKey(), entry.getValue().orElseThrow(() -> - new PartitionNotFoundException(table.getSchemaTableName(), extractPartitionValues(entry.getKey()))))) - .collect(toImmutableMap(Map.Entry::getKey, Map.Entry::getValue)); - - return partitionNames.stream() - .map(partitions::get) - .collect(toImmutableList()); - } - - public Map> getPartitionsByNames(String databaseName, String tableName, List partitionNames) - { - return delegate.getTable(databaseName, tableName) - .map(table -> delegate.getPartitionsByNames(table, partitionNames)) - .orElseGet(() -> partitionNames.stream() - .collect(toImmutableMap(name -> name, name -> Optional.empty()))); - } - - public void addPartitions(String databaseName, String tableName, List partitions) - { - delegate.addPartitions(databaseName, tableName, partitions); - } - - public void dropPartition(String databaseName, String tableName, List parts, boolean deleteData) - { - delegate.dropPartition(databaseName, tableName, parts, deleteData); - } - - public void alterPartition(String databaseName, String tableName, PartitionWithStatistics partition) - { - delegate.alterPartition(databaseName, tableName, partition); - } - - public void createRole(String role, String grantor) - { - delegate.createRole(role, grantor); - } - - public void dropRole(String role) - { - delegate.dropRole(role); - } - - public Set listRoles() - { - return delegate.listRoles(); - } - - public void grantRoles(Set roles, Set grantees, boolean adminOption, HivePrincipal grantor) - { - delegate.grantRoles(roles, grantees, adminOption, grantor); - } - - public void revokeRoles(Set roles, Set grantees, boolean adminOption, HivePrincipal grantor) - { - delegate.revokeRoles(roles, grantees, adminOption, grantor); - } - - public Set listGrantedPrincipals(String role) - { - return delegate.listGrantedPrincipals(role); - } - - public Set listRoleGrants(HivePrincipal principal) - { - return delegate.listRoleGrants(principal); - } - - public void grantTablePrivileges(String databaseName, String tableName, String tableOwner, HivePrincipal grantee, HivePrincipal grantor, Set privileges, boolean grantOption) - { - delegate.grantTablePrivileges(databaseName, tableName, tableOwner, grantee, grantor, privileges, grantOption); - } - - public void revokeTablePrivileges(String databaseName, String tableName, String tableOwner, HivePrincipal grantee, HivePrincipal grantor, Set privileges, boolean grantOption) - { - delegate.revokeTablePrivileges(databaseName, tableName, tableOwner, grantee, grantor, privileges, grantOption); - } - - public Set listTablePrivileges(String databaseName, String tableName, Optional tableOwner, Optional principal) - { - return delegate.listTablePrivileges(databaseName, tableName, tableOwner, principal); - } - - public void checkSupportsTransactions() - { - delegate.checkSupportsTransactions(); - } - - public long openTransaction(AcidTransactionOwner transactionOwner) - { - checkSupportsTransactions(); - return delegate.openTransaction(transactionOwner); - } - - public void commitTransaction(long transactionId) - { - delegate.commitTransaction(transactionId); - } - - public void abortTransaction(long transactionId) - { - delegate.abortTransaction(transactionId); - } - - public void sendTransactionHeartbeat(long transactionId) - { - delegate.sendTransactionHeartbeat(transactionId); - } - - public void acquireSharedReadLock( - AcidTransactionOwner transactionOwner, - String queryId, - long transactionId, - List fullTables, - List partitions) - { - delegate.acquireSharedReadLock(transactionOwner, queryId, transactionId, fullTables, partitions); - } - - public String getValidWriteIds(List tables, long currentTransactionId) - { - return delegate.getValidWriteIds(tables, currentTransactionId); - } - - public Optional getConfigValue(String name) - { - return delegate.getConfigValue(name); - } - - public long allocateWriteId(String dbName, String tableName, long transactionId) - { - return delegate.allocateWriteId(dbName, tableName, transactionId); - } - - public void acquireTableWriteLock( - AcidTransactionOwner transactionOwner, - String queryId, - long transactionId, - String dbName, - String tableName, - DataOperationType operation, - boolean isPartitioned) - { - delegate.acquireTableWriteLock(transactionOwner, queryId, transactionId, dbName, tableName, operation, isPartitioned); - } - - public void updateTableWriteId(String dbName, String tableName, long transactionId, long writeId, OptionalLong rowCountChange) - { - delegate.updateTableWriteId(dbName, tableName, transactionId, writeId, rowCountChange); - } - - public void alterPartitions(String dbName, String tableName, List partitions, long writeId) - { - delegate.alterPartitions(dbName, tableName, partitions, writeId); - } - - public void addDynamicPartitions(String dbName, String tableName, List partitionNames, long transactionId, long writeId, AcidOperation operation) - { - delegate.addDynamicPartitions(dbName, tableName, partitionNames, transactionId, writeId, operation); - } - - public void alterTransactionalTable(Table table, long transactionId, long writeId, PrincipalPrivileges principalPrivileges) - { - delegate.alterTransactionalTable(table, transactionId, writeId, principalPrivileges); - } -} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveModule.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveModule.java index bb06e634a63e..3c241c550027 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveModule.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveModule.java @@ -13,21 +13,17 @@ */ package io.trino.plugin.hive; -import com.google.common.collect.ImmutableList; import com.google.inject.Binder; import com.google.inject.Module; import com.google.inject.Provides; import com.google.inject.Scopes; import com.google.inject.Singleton; import com.google.inject.multibindings.Multibinder; -import io.airlift.event.client.EventClient; -import io.trino.hdfs.HdfsNamenodeStats; -import io.trino.hdfs.TrinoFileSystemCache; -import io.trino.hdfs.TrinoFileSystemCacheStats; import io.trino.plugin.base.CatalogName; import io.trino.plugin.hive.avro.AvroFileWriterFactory; import io.trino.plugin.hive.avro.AvroPageSourceFactory; import io.trino.plugin.hive.fs.CachingDirectoryLister; +import io.trino.plugin.hive.fs.DirectoryLister; import io.trino.plugin.hive.fs.TransactionScopeCachingDirectoryListerFactory; import io.trino.plugin.hive.line.CsvFileWriterFactory; import io.trino.plugin.hive.line.CsvPageSourceFactory; @@ -52,8 +48,6 @@ import io.trino.plugin.hive.parquet.ParquetReaderConfig; import io.trino.plugin.hive.parquet.ParquetWriterConfig; import io.trino.plugin.hive.rcfile.RcFilePageSourceFactory; -import io.trino.plugin.hive.s3select.S3SelectRecordCursorProvider; -import io.trino.plugin.hive.s3select.TrinoS3ClientFactory; import io.trino.spi.connector.ConnectorNodePartitioningProvider; import io.trino.spi.connector.ConnectorPageSinkProvider; import io.trino.spi.connector.ConnectorPageSourceProvider; @@ -85,34 +79,25 @@ public void configure(Binder binder) binder.bind(HiveTableProperties.class).in(Scopes.SINGLETON); binder.bind(HiveColumnProperties.class).in(Scopes.SINGLETON); binder.bind(HiveAnalyzeProperties.class).in(Scopes.SINGLETON); - newOptionalBinder(binder, HiveMaterializedViewPropertiesProvider.class) - .setDefault().toInstance(ImmutableList::of); - - binder.bind(TrinoS3ClientFactory.class).in(Scopes.SINGLETON); binder.bind(CachingDirectoryLister.class).in(Scopes.SINGLETON); newExporter(binder).export(CachingDirectoryLister.class).withGeneratedName(); + binder.bind(DirectoryLister.class).to(CachingDirectoryLister.class).in(Scopes.SINGLETON); binder.bind(HiveWriterStats.class).in(Scopes.SINGLETON); newExporter(binder).export(HiveWriterStats.class).withGeneratedName(); - - newSetBinder(binder, EventClient.class).addBinding().to(HiveEventClient.class).in(Scopes.SINGLETON); binder.bind(HivePartitionManager.class).in(Scopes.SINGLETON); binder.bind(LocationService.class).to(HiveLocationService.class).in(Scopes.SINGLETON); Multibinder systemTableProviders = newSetBinder(binder, SystemTableProvider.class); systemTableProviders.addBinding().to(PartitionsSystemTableProvider.class).in(Scopes.SINGLETON); systemTableProviders.addBinding().to(PropertiesSystemTableProvider.class).in(Scopes.SINGLETON); - newOptionalBinder(binder, HiveRedirectionsProvider.class) - .setDefault().to(NoneHiveRedirectionsProvider.class).in(Scopes.SINGLETON); - newOptionalBinder(binder, HiveMaterializedViewMetadataFactory.class) - .setDefault().to(DefaultHiveMaterializedViewMetadataFactory.class).in(Scopes.SINGLETON); newOptionalBinder(binder, TransactionalMetadataFactory.class) .setDefault().to(HiveMetadataFactory.class).in(Scopes.SINGLETON); binder.bind(TransactionScopeCachingDirectoryListerFactory.class).in(Scopes.SINGLETON); binder.bind(HiveTransactionManager.class).in(Scopes.SINGLETON); binder.bind(ConnectorSplitManager.class).to(HiveSplitManager.class).in(Scopes.SINGLETON); newExporter(binder).export(ConnectorSplitManager.class).as(generator -> generator.generatedNameOf(HiveSplitManager.class)); - newOptionalBinder(binder, ConnectorPageSourceProvider.class).setDefault().to(HivePageSourceProvider.class).in(Scopes.SINGLETON); + binder.bind(ConnectorPageSourceProvider.class).to(HivePageSourceProvider.class).in(Scopes.SINGLETON); binder.bind(ConnectorPageSinkProvider.class).to(HivePageSinkProvider.class).in(Scopes.SINGLETON); binder.bind(ConnectorNodePartitioningProvider.class).to(HiveNodePartitioningProvider.class).in(Scopes.SINGLETON); @@ -121,16 +106,6 @@ public void configure(Binder binder) binder.bind(FileFormatDataSourceStats.class).in(Scopes.SINGLETON); newExporter(binder).export(FileFormatDataSourceStats.class).withGeneratedName(); - binder.bind(TrinoFileSystemCacheStats.class).toInstance(TrinoFileSystemCache.INSTANCE.getFileSystemCacheStats()); - newExporter(binder).export(TrinoFileSystemCacheStats.class) - .as(generator -> generator.generatedNameOf(io.trino.plugin.hive.fs.TrinoFileSystemCache.class)); - - binder.bind(HdfsNamenodeStats.class).in(Scopes.SINGLETON); - newExporter(binder).export(HdfsNamenodeStats.class) - .as(generator -> generator.generatedNameOf(NamenodeStats.class)); - - configBinder(binder).bindConfig(HiveFormatsConfig.class); - Multibinder pageSourceFactoryBinder = newSetBinder(binder, HivePageSourceFactory.class); pageSourceFactoryBinder.addBinding().to(CsvPageSourceFactory.class).in(Scopes.SINGLETON); pageSourceFactoryBinder.addBinding().to(JsonPageSourceFactory.class).in(Scopes.SINGLETON); @@ -143,11 +118,6 @@ public void configure(Binder binder) pageSourceFactoryBinder.addBinding().to(RcFilePageSourceFactory.class).in(Scopes.SINGLETON); pageSourceFactoryBinder.addBinding().to(AvroPageSourceFactory.class).in(Scopes.SINGLETON); - Multibinder recordCursorProviderBinder = newSetBinder(binder, HiveRecordCursorProvider.class); - recordCursorProviderBinder.addBinding().to(S3SelectRecordCursorProvider.class).in(Scopes.SINGLETON); - - binder.bind(GenericHiveRecordCursorProvider.class).in(Scopes.SINGLETON); - Multibinder fileWriterFactoryBinder = newSetBinder(binder, HiveFileWriterFactory.class); binder.bind(OrcFileWriterFactory.class).in(Scopes.SINGLETON); newExporter(binder).export(OrcFileWriterFactory.class).withGeneratedName(); @@ -192,4 +162,12 @@ public boolean translateHiveViews(HiveConfig hiveConfig) { return hiveConfig.isTranslateHiveViews(); } + + @Provides + @Singleton + @HideDeltaLakeTables + public boolean hideDeltaLakeTables(HiveMetastoreConfig config) + { + return config.isHideDeltaLakeTables(); + } } diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveOutputTableHandle.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveOutputTableHandle.java index 043c2395544a..29190dcee938 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveOutputTableHandle.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveOutputTableHandle.java @@ -46,7 +46,7 @@ public HiveOutputTableHandle( @JsonProperty("tableStorageFormat") HiveStorageFormat tableStorageFormat, @JsonProperty("partitionStorageFormat") HiveStorageFormat partitionStorageFormat, @JsonProperty("partitionedBy") List partitionedBy, - @JsonProperty("bucketProperty") Optional bucketProperty, + @JsonProperty("bucketInfo") Optional bucketInfo, @JsonProperty("tableOwner") String tableOwner, @JsonProperty("additionalTableParameters") Map additionalTableParameters, @JsonProperty("transaction") AcidTransaction transaction, @@ -59,7 +59,7 @@ public HiveOutputTableHandle( inputColumns, pageSinkMetadata, locationHandle, - bucketProperty, + bucketInfo, tableStorageFormat, partitionStorageFormat, transaction, diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HivePageSink.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HivePageSink.java index 26e84921435f..4629e2aeeb2e 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HivePageSink.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HivePageSink.java @@ -19,10 +19,10 @@ import com.google.common.util.concurrent.Futures; import com.google.common.util.concurrent.ListenableFuture; import com.google.common.util.concurrent.ListeningExecutorService; -import io.airlift.concurrent.MoreFutures; import io.airlift.json.JsonCodec; import io.airlift.slice.Slice; -import io.trino.hdfs.HdfsEnvironment; +import io.trino.plugin.hive.HiveWritableTableHandle.BucketInfo; +import io.trino.plugin.hive.acid.AcidTransaction; import io.trino.plugin.hive.util.HiveBucketing.BucketingVersion; import io.trino.spi.Page; import io.trino.spi.PageIndexer; @@ -54,6 +54,7 @@ import static com.google.common.base.Verify.verify; import static com.google.common.collect.ImmutableList.toImmutableList; import static com.google.common.util.concurrent.MoreExecutors.directExecutor; +import static io.airlift.concurrent.MoreFutures.toCompletableFuture; import static io.airlift.slice.Slices.wrappedBuffer; import static io.trino.plugin.hive.HiveErrorCode.HIVE_TOO_MANY_OPEN_PARTITIONS; import static io.trino.plugin.hive.HiveErrorCode.HIVE_WRITER_CLOSE_ERROR; @@ -78,7 +79,6 @@ public class HivePageSink private final HiveBucketFunction bucketFunction; private final HiveWriterPagePartitioner pagePartitioner; - private final HdfsEnvironment hdfsEnvironment; private final int maxOpenWriters; private final ListeningExecutorService writeVerificationExecutor; @@ -87,8 +87,6 @@ public class HivePageSink private final List writers = new ArrayList<>(); - private final ConnectorSession session; - private final long targetMaxFileSize; private final List closedWriterRollbackActions = new ArrayList<>(); private final List partitionUpdates = new ArrayList<>(); @@ -100,13 +98,11 @@ public class HivePageSink private long validationCpuNanos; public HivePageSink( - HiveWritableTableHandle tableHandle, HiveWriterFactory writerFactory, List inputColumns, - boolean isTransactional, - Optional bucketProperty, + AcidTransaction acidTransaction, + Optional bucketInfo, PageIndexerFactory pageIndexerFactory, - HdfsEnvironment hdfsEnvironment, int maxOpenWriters, ListeningExecutorService writeVerificationExecutor, JsonCodec partitionUpdateCodec, @@ -118,17 +114,16 @@ public HivePageSink( requireNonNull(pageIndexerFactory, "pageIndexerFactory is null"); - this.isTransactional = isTransactional; - this.hdfsEnvironment = requireNonNull(hdfsEnvironment, "hdfsEnvironment is null"); + this.isTransactional = acidTransaction.isTransactional(); this.maxOpenWriters = maxOpenWriters; this.writeVerificationExecutor = requireNonNull(writeVerificationExecutor, "writeVerificationExecutor is null"); this.partitionUpdateCodec = requireNonNull(partitionUpdateCodec, "partitionUpdateCodec is null"); - this.isMergeSink = tableHandle.getTransaction().isMerge(); - requireNonNull(bucketProperty, "bucketProperty is null"); + this.isMergeSink = acidTransaction.isMerge(); + requireNonNull(bucketInfo, "bucketInfo is null"); this.pagePartitioner = new HiveWriterPagePartitioner( inputColumns, - bucketProperty.isPresent(), + bucketInfo.isPresent(), pageIndexerFactory); // determine the input index of the partition columns and data columns @@ -151,13 +146,13 @@ public HivePageSink( this.partitionColumnsInputIndex = Ints.toArray(partitionColumns.build()); this.dataColumnInputIndex = Ints.toArray(dataColumnsInputIndex.build()); - if (bucketProperty.isPresent()) { - BucketingVersion bucketingVersion = bucketProperty.get().getBucketingVersion(); - int bucketCount = bucketProperty.get().getBucketCount(); - bucketColumns = bucketProperty.get().getBucketedBy().stream() + if (bucketInfo.isPresent()) { + BucketingVersion bucketingVersion = bucketInfo.get().bucketingVersion(); + int bucketCount = bucketInfo.get().bucketCount(); + bucketColumns = bucketInfo.get().bucketedBy().stream() .mapToInt(dataColumnNameToIdMap::get) .toArray(); - List bucketColumnTypes = bucketProperty.get().getBucketedBy().stream() + List bucketColumnTypes = bucketInfo.get().bucketedBy().stream() .map(dataColumnNameToTypeMap::get) .collect(toList()); bucketFunction = new HiveBucketFunction(bucketingVersion, bucketCount, bucketColumnTypes); @@ -167,7 +162,6 @@ public HivePageSink( bucketFunction = null; } - this.session = requireNonNull(session, "session is null"); this.targetMaxFileSize = HiveSessionProperties.getTargetMaxFileSize(session).toBytes(); } @@ -192,13 +186,7 @@ public long getValidationCpuNanos() @Override public CompletableFuture> finish() { - // Must be wrapped in doAs entirely - // Implicit FileSystem initializations are possible in HiveRecordWriter#commit -> RecordWriter#close - ListenableFuture> result = hdfsEnvironment.doAs( - session.getIdentity(), - isMergeSink ? this::doMergeSinkFinish : this::doInsertSinkFinish); - - return MoreFutures.toCompletableFuture(result); + return toCompletableFuture(isMergeSink ? doMergeSinkFinish() : doInsertSinkFinish()); } private ListenableFuture> doMergeSinkFinish() @@ -244,13 +232,6 @@ private ListenableFuture> doInsertSinkFinish() @Override public void abort() - { - // Must be wrapped in doAs entirely - // Implicit FileSystem initializations are possible in HiveRecordWriter#rollback -> RecordWriter#close - hdfsEnvironment.doAs(session.getIdentity(), this::doAbort); - } - - private void doAbort() { List rollbackActions = Streams.concat( writers.stream() @@ -278,17 +259,6 @@ private void doAbort() @Override public CompletableFuture appendPage(Page page) - { - if (page.getPositionCount() > 0) { - // Must be wrapped in doAs entirely - // Implicit FileSystem initializations are possible in HiveRecordWriter#addRow or #createWriter - hdfsEnvironment.doAs(session.getIdentity(), () -> doAppend(page)); - } - - return NOT_BLOCKED; - } - - private void doAppend(Page page) { int writeOffset = 0; while (writeOffset < page.getPositionCount()) { @@ -296,6 +266,7 @@ private void doAppend(Page page) writeOffset += chunk.getPositionCount(); writePage(chunk); } + return NOT_BLOCKED; } private void writePage(Page page) diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HivePageSinkProvider.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HivePageSinkProvider.java index cd0f2032e4ca..79488183794e 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HivePageSinkProvider.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HivePageSinkProvider.java @@ -18,15 +18,13 @@ import com.google.common.collect.ImmutableSet; import com.google.common.util.concurrent.ListeningExecutorService; import com.google.inject.Inject; -import io.airlift.event.client.EventClient; import io.airlift.json.JsonCodec; import io.airlift.units.DataSize; import io.trino.filesystem.TrinoFileSystemFactory; -import io.trino.hdfs.HdfsEnvironment; import io.trino.plugin.hive.metastore.HiveMetastoreFactory; import io.trino.plugin.hive.metastore.HivePageSinkMetadataProvider; import io.trino.plugin.hive.metastore.SortingColumn; -import io.trino.spi.NodeManager; +import io.trino.plugin.hive.metastore.cache.CachingHiveMetastore; import io.trino.spi.PageIndexerFactory; import io.trino.spi.PageSorter; import io.trino.spi.connector.ConnectorInsertTableHandle; @@ -51,7 +49,7 @@ import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.util.concurrent.MoreExecutors.listeningDecorator; import static io.airlift.concurrent.Threads.daemonThreadsNamed; -import static io.trino.plugin.hive.metastore.cache.CachingHiveMetastore.memoizeMetastore; +import static io.trino.plugin.hive.metastore.cache.CachingHiveMetastore.createPerTransactionCache; import static java.util.Objects.requireNonNull; import static java.util.concurrent.Executors.newFixedThreadPool; @@ -60,7 +58,6 @@ public class HivePageSinkProvider { private final Set fileWriterFactories; private final TrinoFileSystemFactory fileSystemFactory; - private final HdfsEnvironment hdfsEnvironment; private final PageSorter pageSorter; private final HiveMetastoreFactory metastoreFactory; private final PageIndexerFactory pageIndexerFactory; @@ -71,9 +68,6 @@ public class HivePageSinkProvider private final LocationService locationService; private final ListeningExecutorService writeVerificationExecutor; private final JsonCodec partitionUpdateCodec; - private final NodeManager nodeManager; - private final EventClient eventClient; - private final HiveSessionProperties hiveSessionProperties; private final HiveWriterStats hiveWriterStats; private final long perTransactionMetastoreCacheMaximumSize; private final DateTimeZone parquetTimeZone; @@ -84,7 +78,6 @@ public class HivePageSinkProvider public HivePageSinkProvider( Set fileWriterFactories, TrinoFileSystemFactory fileSystemFactory, - HdfsEnvironment hdfsEnvironment, PageSorter pageSorter, HiveMetastoreFactory metastoreFactory, PageIndexerFactory pageIndexerFactory, @@ -93,14 +86,10 @@ public HivePageSinkProvider( SortingFileWriterConfig sortingFileWriterConfig, LocationService locationService, JsonCodec partitionUpdateCodec, - NodeManager nodeManager, - EventClient eventClient, - HiveSessionProperties hiveSessionProperties, HiveWriterStats hiveWriterStats) { this.fileWriterFactories = ImmutableSet.copyOf(requireNonNull(fileWriterFactories, "fileWriterFactories is null")); this.fileSystemFactory = requireNonNull(fileSystemFactory, "fileSystemFactory is null"); - this.hdfsEnvironment = requireNonNull(hdfsEnvironment, "hdfsEnvironment is null"); this.pageSorter = requireNonNull(pageSorter, "pageSorter is null"); this.metastoreFactory = requireNonNull(metastoreFactory, "metastoreFactory is null"); this.pageIndexerFactory = requireNonNull(pageIndexerFactory, "pageIndexerFactory is null"); @@ -111,9 +100,6 @@ public HivePageSinkProvider( this.locationService = requireNonNull(locationService, "locationService is null"); this.writeVerificationExecutor = listeningDecorator(newFixedThreadPool(config.getWriteValidationThreads(), daemonThreadsNamed("hive-write-validation-%s"))); this.partitionUpdateCodec = requireNonNull(partitionUpdateCodec, "partitionUpdateCodec is null"); - this.nodeManager = requireNonNull(nodeManager, "nodeManager is null"); - this.eventClient = requireNonNull(eventClient, "eventClient is null"); - this.hiveSessionProperties = requireNonNull(hiveSessionProperties, "hiveSessionProperties is null"); this.hiveWriterStats = requireNonNull(hiveWriterStats, "hiveWriterStats is null"); this.perTransactionMetastoreCacheMaximumSize = config.getPerTransactionMetastoreCacheMaximumSize(); this.parquetTimeZone = config.getParquetDateTimeZone(); @@ -156,11 +142,12 @@ private HivePageSink createPageSink(HiveWritableTableHandle handle, boolean isCr OptionalInt bucketCount = OptionalInt.empty(); List sortedBy = ImmutableList.of(); - if (handle.getBucketProperty().isPresent()) { - bucketCount = OptionalInt.of(handle.getBucketProperty().get().getBucketCount()); - sortedBy = handle.getBucketProperty().get().getSortedBy(); + if (handle.getBucketInfo().isPresent()) { + bucketCount = OptionalInt.of(handle.getBucketInfo().get().bucketCount()); + sortedBy = handle.getBucketInfo().get().sortedBy(); } + CachingHiveMetastore cachingHiveMetastore = createPerTransactionCache(metastoreFactory.createMetastore(Optional.of(session.getIdentity())), perTransactionMetastoreCacheMaximumSize); HiveWriterFactory writerFactory = new HiveWriterFactory( fileWriterFactories, fileSystemFactory, @@ -177,31 +164,22 @@ private HivePageSink createPageSink(HiveWritableTableHandle handle, boolean isCr handle.getLocationHandle(), locationService, session.getQueryId(), - new HivePageSinkMetadataProvider( - handle.getPageSinkMetadata(), - new HiveMetastoreClosure(memoizeMetastore(metastoreFactory.createMetastore(Optional.of(session.getIdentity())), perTransactionMetastoreCacheMaximumSize))), + new HivePageSinkMetadataProvider(handle.getPageSinkMetadata(), cachingHiveMetastore), typeManager, - hdfsEnvironment, pageSorter, writerSortBufferSize, maxOpenSortFiles, - parquetTimeZone, session, - nodeManager, - eventClient, - hiveSessionProperties, hiveWriterStats, temporaryStagingDirectoryDirectoryEnabled, temporaryStagingDirectoryPath); return new HivePageSink( - handle, writerFactory, handle.getInputColumns(), - handle.isTransactional(), - handle.getBucketProperty(), + handle.getTransaction(), + handle.getBucketInfo(), pageIndexerFactory, - hdfsEnvironment, maxOpenPartitions, writeVerificationExecutor, partitionUpdateCodec, diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HivePageSourceFactory.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HivePageSourceFactory.java index 308fcb58b915..43979da2c704 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HivePageSourceFactory.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HivePageSourceFactory.java @@ -21,7 +21,6 @@ import java.util.List; import java.util.Optional; import java.util.OptionalInt; -import java.util.Properties; public interface HivePageSourceFactory { @@ -31,7 +30,8 @@ Optional createPageSource( long start, long length, long estimatedFileSize, - Properties schema, + long fileModifiedTime, + Schema schema, List columns, TupleDomain effectivePredicate, Optional acidInfo, diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HivePageSourceProvider.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HivePageSourceProvider.java index f7871ebbd5fa..970dadfe0d98 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HivePageSourceProvider.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HivePageSourceProvider.java @@ -19,10 +19,7 @@ import com.google.common.collect.ImmutableSet; import com.google.inject.Inject; import io.trino.filesystem.Location; -import io.trino.hdfs.HdfsContext; -import io.trino.hdfs.HdfsEnvironment; import io.trino.plugin.hive.HivePageSource.BucketValidator; -import io.trino.plugin.hive.HiveRecordCursorProvider.ReaderRecordCursorWithProjections; import io.trino.plugin.hive.HiveSplit.BucketConversion; import io.trino.plugin.hive.HiveSplit.BucketValidation; import io.trino.plugin.hive.acid.AcidTransaction; @@ -37,15 +34,10 @@ import io.trino.spi.connector.ConnectorTransactionHandle; import io.trino.spi.connector.DynamicFilter; import io.trino.spi.connector.EmptyPageSource; -import io.trino.spi.connector.RecordCursor; -import io.trino.spi.connector.RecordPageSource; import io.trino.spi.predicate.Domain; import io.trino.spi.predicate.NullableValue; import io.trino.spi.predicate.TupleDomain; -import io.trino.spi.type.Type; import io.trino.spi.type.TypeManager; -import org.apache.hadoop.conf.Configuration; -import org.apache.hadoop.fs.Path; import java.util.ArrayList; import java.util.HashMap; @@ -55,7 +47,6 @@ import java.util.Objects; import java.util.Optional; import java.util.OptionalInt; -import java.util.Properties; import java.util.Set; import java.util.regex.Pattern; @@ -86,28 +77,18 @@ public class HivePageSourceProvider private static final Pattern ORIGINAL_FILE_PATH_MATCHER = Pattern.compile("(?s)(?.*)/(?(?\\d+)_(?.*)?)$"); private final TypeManager typeManager; - private final HdfsEnvironment hdfsEnvironment; private final int domainCompactionThreshold; private final Set pageSourceFactories; - private final Set cursorProviders; @Inject public HivePageSourceProvider( TypeManager typeManager, - HdfsEnvironment hdfsEnvironment, HiveConfig hiveConfig, - Set pageSourceFactories, - Set cursorProviders, - GenericHiveRecordCursorProvider genericCursorProvider) + Set pageSourceFactories) { this.typeManager = requireNonNull(typeManager, "typeManager is null"); - this.hdfsEnvironment = requireNonNull(hdfsEnvironment, "hdfsEnvironment is null"); this.domainCompactionThreshold = hiveConfig.getDomainCompactionThreshold(); this.pageSourceFactories = ImmutableSet.copyOf(requireNonNull(pageSourceFactories, "pageSourceFactories is null")); - this.cursorProviders = ImmutableSet.builder() - .addAll(requireNonNull(cursorProviders, "cursorProviders is null")) - .add(genericCursorProvider) // generic should be last, as a fallback option - .build(); } @Override @@ -137,7 +118,7 @@ public ConnectorPageSource createPageSource( hiveSplit.getPartitionKeys(), hiveColumns, hiveSplit.getBucketConversion().map(BucketConversion::bucketColumnHandles).orElse(ImmutableList.of()), - hiveSplit.getTableToPartitionMapping(), + hiveSplit.getHiveColumnCoercions(), hiveSplit.getPath(), hiveSplit.getTableBucketNumber(), hiveSplit.getEstimatedFileSize(), @@ -149,27 +130,22 @@ public ConnectorPageSource createPageSource( return new EmptyPageSource(); } - Configuration configuration = hdfsEnvironment.getConfiguration(new HdfsContext(session), new Path(hiveSplit.getPath())); - Optional pageSource = createHivePageSource( pageSourceFactories, - cursorProviders, - configuration, session, Location.of(hiveSplit.getPath()), hiveSplit.getTableBucketNumber(), hiveSplit.getStart(), hiveSplit.getLength(), hiveSplit.getEstimatedFileSize(), + hiveSplit.getFileModifiedTime(), hiveSplit.getSchema(), hiveTable.getCompactEffectivePredicate().intersect( dynamicFilter.getCurrentPredicate().transformKeys(HiveColumnHandle.class::cast)) .simplify(domainCompactionThreshold), - hiveColumns, typeManager, hiveSplit.getBucketConversion(), hiveSplit.getBucketValidation(), - hiveSplit.isS3SelectPushdownEnabled(), hiveSplit.getAcidInfo(), originalFile, hiveTable.getTransaction(), @@ -183,21 +159,18 @@ public ConnectorPageSource createPageSource( public static Optional createHivePageSource( Set pageSourceFactories, - Set cursorProviders, - Configuration configuration, ConnectorSession session, Location path, OptionalInt tableBucketNumber, long start, long length, long estimatedFileSize, - Properties schema, + long fileModifiedTime, + Schema schema, TupleDomain effectivePredicate, - List columns, TypeManager typeManager, Optional bucketConversion, Optional bucketValidation, - boolean s3SelectPushdownEnabled, Optional acidInfo, boolean originalFile, AcidTransaction transaction, @@ -223,6 +196,7 @@ public static Optional createHivePageSource( start, length, estimatedFileSize, + fileModifiedTime, schema, desiredColumns, effectivePredicate, @@ -251,58 +225,6 @@ public static Optional createHivePageSource( } } - for (HiveRecordCursorProvider provider : cursorProviders) { - List desiredColumns = toColumnHandles(regularAndInterimColumnMappings, false, typeManager, timestampPrecision); - Optional readerWithProjections = provider.createRecordCursor( - configuration, - session, - path, - start, - length, - estimatedFileSize, - schema, - desiredColumns, - effectivePredicate, - typeManager, - s3SelectPushdownEnabled); - - if (readerWithProjections.isPresent()) { - RecordCursor delegate = readerWithProjections.get().getRecordCursor(); - Optional projections = readerWithProjections.get().getProjectedReaderColumns(); - - if (projections.isPresent()) { - ReaderProjectionsAdapter projectionsAdapter = hiveProjectionsAdapter(desiredColumns, projections.get()); - delegate = new HiveReaderProjectionsAdaptingRecordCursor(delegate, projectionsAdapter); - } - - checkArgument(acidInfo.isEmpty(), "Acid is not supported"); - - if (bucketAdaptation.isPresent()) { - delegate = new HiveBucketAdapterRecordCursor( - bucketAdaptation.get().getBucketColumnIndices(), - bucketAdaptation.get().getBucketColumnHiveTypes(), - bucketAdaptation.get().getBucketingVersion(), - bucketAdaptation.get().getTableBucketCount(), - bucketAdaptation.get().getPartitionBucketCount(), - bucketAdaptation.get().getBucketToKeep(), - typeManager, - delegate); - } - - // bucket adaptation already validates that data is in the right bucket - if (bucketAdaptation.isEmpty() && bucketValidator.isPresent()) { - delegate = bucketValidator.get().wrapRecordCursor(delegate, typeManager); - } - - HiveRecordCursor hiveRecordCursor = new HiveRecordCursor(columnMappings, delegate); - List columnTypes = columns.stream() - .map(HiveColumnHandle::getType) - .collect(toList()); - - return Optional.of(new RecordPageSource(columnTypes, hiveRecordCursor)); - } - } - return Optional.empty(); } @@ -453,7 +375,7 @@ public static List buildColumnMappings( List partitionKeys, List columns, List requiredInterimColumns, - TableToPartitionMapping tableToPartitionMapping, + Map hiveColumnCoercions, String path, OptionalInt bucketNumber, long estimatedFileSize, @@ -469,7 +391,7 @@ public static List buildColumnMappings( int regularIndex = 0; for (HiveColumnHandle column : columns) { - Optional baseTypeCoercionFrom = tableToPartitionMapping.getCoercion(column.getBaseHiveColumnIndex()); + Optional baseTypeCoercionFrom = Optional.ofNullable(hiveColumnCoercions.get(column.getBaseHiveColumnIndex())).map(HiveTypeName::toHiveType); if (column.getColumnType() == REGULAR) { if (column.isBaseColumn()) { baseColumnHiveIndices.add(column.getBaseHiveColumnIndex()); @@ -520,7 +442,8 @@ else if (isRowIdColumnHandle(column)) { } if (projectionsForColumn.containsKey(column.getBaseHiveColumnIndex())) { - columnMappings.add(interim(column, regularIndex, tableToPartitionMapping.getCoercion(column.getBaseHiveColumnIndex()))); + Optional baseTypeCoercionFrom = Optional.ofNullable(hiveColumnCoercions.get(column.getBaseHiveColumnIndex())).map(HiveTypeName::toHiveType); + columnMappings.add(interim(column, regularIndex, baseTypeCoercionFrom)); } else { // If coercion does not affect bucket number calculation, coercion doesn't need to be applied here. diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HivePartitionManager.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HivePartitionManager.java index b105971a8aac..0467990bb4ab 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HivePartitionManager.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HivePartitionManager.java @@ -74,11 +74,11 @@ public HivePartitionResult getPartitions(SemiTransactionalHiveMetastore metastor .intersect(hiveTableHandle.getEnforcedConstraint()); SchemaTableName tableName = hiveTableHandle.getSchemaTableName(); - Optional hiveBucketHandle = hiveTableHandle.getBucketHandle(); + Optional tablePartitioning = hiveTableHandle.getTablePartitioning(); List partitionColumns = hiveTableHandle.getPartitionColumns(); if (effectivePredicate.isNone()) { - return new HivePartitionResult(partitionColumns, Optional.empty(), ImmutableList.of(), TupleDomain.none(), TupleDomain.none(), hiveBucketHandle, Optional.empty()); + return new HivePartitionResult(partitionColumns, Optional.empty(), ImmutableList.of(), TupleDomain.none(), TupleDomain.none(), tablePartitioning, Optional.empty()); } Optional bucketFilter = getHiveBucketFilter(hiveTableHandle, effectivePredicate); @@ -93,7 +93,7 @@ public HivePartitionResult getPartitions(SemiTransactionalHiveMetastore metastor ImmutableList.of(new HivePartition(tableName)), effectivePredicate, compactEffectivePredicate, - hiveBucketHandle, + tablePartitioning, bucketFilter); } @@ -117,15 +117,15 @@ public HivePartitionResult getPartitions(SemiTransactionalHiveMetastore metastor partitionNames = Optional.of(partitionNamesList); } - return new HivePartitionResult(partitionColumns, partitionNames, partitionsIterable, effectivePredicate, compactEffectivePredicate, hiveBucketHandle, bucketFilter); + return new HivePartitionResult(partitionColumns, partitionNames, partitionsIterable, effectivePredicate, compactEffectivePredicate, tablePartitioning, bucketFilter); } public HivePartitionResult getPartitions(ConnectorTableHandle tableHandle, List> partitionValuesList) { HiveTableHandle hiveTableHandle = (HiveTableHandle) tableHandle; SchemaTableName tableName = hiveTableHandle.getSchemaTableName(); + Optional tablePartitioning = hiveTableHandle.getTablePartitioning(); List partitionColumns = hiveTableHandle.getPartitionColumns(); - Optional bucketHandle = hiveTableHandle.getBucketHandle(); List partitionColumnNames = partitionColumns.stream() .map(HiveColumnHandle::getName) @@ -137,7 +137,7 @@ public HivePartitionResult getPartitions(ConnectorTableHandle tableHandle, List< .map(partition -> partition.orElseThrow(() -> new VerifyException("partition must exist"))) .collect(toImmutableList()); - return new HivePartitionResult(partitionColumns, Optional.empty(), partitionList, TupleDomain.all(), TupleDomain.all(), bucketHandle, Optional.empty()); + return new HivePartitionResult(partitionColumns, Optional.empty(), partitionList, TupleDomain.all(), TupleDomain.all(), tablePartitioning, Optional.empty()); } public HiveTableHandle applyPartitionResult(HiveTableHandle handle, HivePartitionResult partitions, Constraint constraint) @@ -166,7 +166,7 @@ public HiveTableHandle applyPartitionResult(HiveTableHandle handle, HivePartitio partitionList, partitions.getCompactEffectivePredicate(), enforcedConstraint, - partitions.getBucketHandle(), + partitions.getTablePartitioning(), partitions.getBucketFilter(), handle.getAnalyzePartitionValues(), Sets.union(handle.getConstraintColumns(), constraint.getPredicateColumns().orElseGet(ImmutableSet::of)), diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HivePartitionMetadata.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HivePartitionMetadata.java index f484f42868c9..15efa2929ad9 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HivePartitionMetadata.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HivePartitionMetadata.java @@ -13,8 +13,10 @@ */ package io.trino.plugin.hive; +import com.google.common.collect.ImmutableMap; import io.trino.plugin.hive.metastore.Partition; +import java.util.Map; import java.util.Optional; import static java.util.Objects.requireNonNull; @@ -23,16 +25,16 @@ public class HivePartitionMetadata { private final Optional partition; private final HivePartition hivePartition; - private final TableToPartitionMapping tableToPartitionMapping; + private final Map hiveColumnCoercions; HivePartitionMetadata( HivePartition hivePartition, Optional partition, - TableToPartitionMapping tableToPartitionMapping) + Map hiveColumnCoercions) { this.partition = requireNonNull(partition, "partition is null"); this.hivePartition = requireNonNull(hivePartition, "hivePartition is null"); - this.tableToPartitionMapping = requireNonNull(tableToPartitionMapping, "tableToPartitionMapping is null"); + this.hiveColumnCoercions = ImmutableMap.copyOf(requireNonNull(hiveColumnCoercions, "hiveColumnCoercions is null")); } public HivePartition getHivePartition() @@ -48,8 +50,8 @@ public Optional getPartition() return partition; } - public TableToPartitionMapping getTableToPartitionMapping() + public Map getHiveColumnCoercions() { - return tableToPartitionMapping; + return hiveColumnCoercions; } } diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HivePartitionResult.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HivePartitionResult.java index cf6fd3943f0b..31fd716ed2a9 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HivePartitionResult.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HivePartitionResult.java @@ -38,7 +38,7 @@ public class HivePartitionResult private final Iterable partitions; private final TupleDomain effectivePredicate; private final TupleDomain compactEffectivePredicate; - private final Optional bucketHandle; + private final Optional tablePartitioning; private final Optional bucketFilter; private final Optional> partitionNames; @@ -48,7 +48,7 @@ public HivePartitionResult( Iterable partitions, TupleDomain effectivePredicate, TupleDomain compactEffectivePredicate, - Optional bucketHandle, + Optional tablePartitioning, Optional bucketFilter) { this.partitionColumns = requireNonNull(partitionColumns, "partitionColumns is null"); @@ -56,7 +56,7 @@ public HivePartitionResult( this.partitions = requireNonNull(partitions, "partitions is null"); this.effectivePredicate = requireNonNull(effectivePredicate, "effectivePredicate is null"); this.compactEffectivePredicate = requireNonNull(compactEffectivePredicate, "compactEffectivePredicate is null"); - this.bucketHandle = requireNonNull(bucketHandle, "bucketHandle is null"); + this.tablePartitioning = requireNonNull(tablePartitioning, "bucketHandle is null"); this.bucketFilter = requireNonNull(bucketFilter, "bucketFilter is null"); } @@ -85,9 +85,9 @@ public TupleDomain getCompactEffectivePredicate() return compactEffectivePredicate; } - public Optional getBucketHandle() + public Optional getTablePartitioning() { - return bucketHandle; + return tablePartitioning; } public Optional getBucketFilter() diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HivePartitioningHandle.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HivePartitioningHandle.java index 51195f2c5189..3be110b5b162 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HivePartitioningHandle.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HivePartitioningHandle.java @@ -21,7 +21,6 @@ import java.util.List; import java.util.Objects; -import java.util.OptionalInt; import static com.google.common.base.MoreObjects.toStringHelper; import static java.util.Objects.requireNonNull; @@ -32,7 +31,6 @@ public class HivePartitioningHandle private final BucketingVersion bucketingVersion; private final int bucketCount; private final List hiveTypes; - private final OptionalInt maxCompatibleBucketCount; private final boolean usePartitionedBucketing; @JsonCreator @@ -40,13 +38,11 @@ public HivePartitioningHandle( @JsonProperty("bucketingVersion") BucketingVersion bucketingVersion, @JsonProperty("bucketCount") int bucketCount, @JsonProperty("hiveBucketTypes") List hiveTypes, - @JsonProperty("maxCompatibleBucketCount") OptionalInt maxCompatibleBucketCount, @JsonProperty("usePartitionedBucketing") boolean usePartitionedBucketing) { this.bucketingVersion = requireNonNull(bucketingVersion, "bucketingVersion is null"); this.bucketCount = bucketCount; this.hiveTypes = requireNonNull(hiveTypes, "hiveTypes is null"); - this.maxCompatibleBucketCount = maxCompatibleBucketCount; this.usePartitionedBucketing = usePartitionedBucketing; } @@ -68,12 +64,6 @@ public List getHiveTypes() return hiveTypes; } - @JsonProperty - public OptionalInt getMaxCompatibleBucketCount() - { - return maxCompatibleBucketCount; - } - @JsonProperty public boolean isUsePartitionedBucketing() { diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveRecordCursorProvider.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveRecordCursorProvider.java index aad75d221610..960c63c0e743 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveRecordCursorProvider.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveRecordCursorProvider.java @@ -21,8 +21,8 @@ import org.apache.hadoop.conf.Configuration; import java.util.List; +import java.util.Map; import java.util.Optional; -import java.util.Properties; import static java.util.Objects.requireNonNull; @@ -35,7 +35,7 @@ Optional createRecordCursor( long start, long length, long fileSize, - Properties schema, + Map schema, List columns, TupleDomain effectivePredicate, TypeManager typeManager, diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveSessionProperties.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveSessionProperties.java index 19e70dc17e1c..8be402bc86dc 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveSessionProperties.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveSessionProperties.java @@ -41,6 +41,7 @@ import static io.trino.plugin.base.session.PropertyMetadataUtil.durationProperty; import static io.trino.plugin.base.session.PropertyMetadataUtil.validateMaxDataSize; import static io.trino.plugin.base.session.PropertyMetadataUtil.validateMinDataSize; +import static io.trino.plugin.hive.parquet.ParquetReaderConfig.PARQUET_READER_MAX_SMALL_FILE_THRESHOLD; import static io.trino.plugin.hive.parquet.ParquetWriterConfig.PARQUET_WRITER_MAX_BLOCK_SIZE; import static io.trino.plugin.hive.parquet.ParquetWriterConfig.PARQUET_WRITER_MAX_PAGE_SIZE; import static io.trino.plugin.hive.parquet.ParquetWriterConfig.PARQUET_WRITER_MIN_PAGE_SIZE; @@ -105,6 +106,7 @@ public final class HiveSessionProperties private static final String PARQUET_USE_BLOOM_FILTER = "parquet_use_bloom_filter"; private static final String PARQUET_MAX_READ_BLOCK_SIZE = "parquet_max_read_block_size"; private static final String PARQUET_MAX_READ_BLOCK_ROW_COUNT = "parquet_max_read_block_row_count"; + private static final String PARQUET_SMALL_FILE_THRESHOLD = "parquet_small_file_threshold"; private static final String PARQUET_OPTIMIZED_READER_ENABLED = "parquet_optimized_reader_enabled"; private static final String PARQUET_OPTIMIZED_NESTED_READER_ENABLED = "parquet_optimized_nested_reader_enabled"; private static final String PARQUET_WRITER_BLOCK_SIZE = "parquet_writer_block_size"; @@ -163,7 +165,6 @@ static boolean isValid(InsertExistingPartitionsBehavior value, boolean immutable @Inject public HiveSessionProperties( HiveConfig hiveConfig, - HiveFormatsConfig hiveFormatsConfig, OrcReaderConfig orcReaderConfig, OrcWriterConfig orcWriterConfig, ParquetReaderConfig parquetReaderConfig, @@ -204,71 +205,6 @@ public HiveSessionProperties( false, value -> InsertExistingPartitionsBehavior.valueOf((String) value, hiveConfig.isImmutablePartitions()), InsertExistingPartitionsBehavior::toString), - booleanProperty( - AVRO_NATIVE_READER_ENABLED, - "Use native Avro file reader", - hiveFormatsConfig.isAvroFileNativeReaderEnabled(), - false), - booleanProperty( - AVRO_NATIVE_WRITER_ENABLED, - "Use native Avro file writer", - hiveFormatsConfig.isAvroFileNativeWriterEnabled(), - false), - booleanProperty( - CSV_NATIVE_READER_ENABLED, - "Use native CSV reader", - hiveFormatsConfig.isCsvNativeReaderEnabled(), - false), - booleanProperty( - CSV_NATIVE_WRITER_ENABLED, - "Use native CSV writer", - hiveFormatsConfig.isCsvNativeWriterEnabled(), - false), - booleanProperty( - JSON_NATIVE_READER_ENABLED, - "Use native JSON reader", - hiveFormatsConfig.isJsonNativeReaderEnabled(), - false), - booleanProperty( - JSON_NATIVE_WRITER_ENABLED, - "Use native JSON writer", - hiveFormatsConfig.isJsonNativeWriterEnabled(), - false), - booleanProperty( - OPENX_JSON_NATIVE_READER_ENABLED, - "Use native OpenX JSON reader", - hiveFormatsConfig.isOpenXJsonNativeReaderEnabled(), - false), - booleanProperty( - OPENX_JSON_NATIVE_WRITER_ENABLED, - "Use native OpenX JSON writer", - hiveFormatsConfig.isOpenXJsonNativeWriterEnabled(), - false), - booleanProperty( - REGEX_NATIVE_READER_ENABLED, - "Use native REGEX reader", - hiveFormatsConfig.isRegexNativeReaderEnabled(), - false), - booleanProperty( - TEXT_FILE_NATIVE_READER_ENABLED, - "Use native text file reader", - hiveFormatsConfig.isTextFileNativeReaderEnabled(), - false), - booleanProperty( - TEXT_FILE_NATIVE_WRITER_ENABLED, - "Use native text file writer", - hiveFormatsConfig.isTextFileNativeWriterEnabled(), - false), - booleanProperty( - SEQUENCE_FILE_NATIVE_READER_ENABLED, - "Use native sequence file reader", - hiveFormatsConfig.isSequenceFileNativeReaderEnabled(), - false), - booleanProperty( - SEQUENCE_FILE_NATIVE_WRITER_ENABLED, - "Use native sequence file writer", - hiveFormatsConfig.isSequenceFileNativeWriterEnabled(), - false), booleanProperty( ORC_BLOOM_FILTERS_ENABLED, "ORC: Enable bloom filters for predicate pushdown", @@ -426,15 +362,11 @@ public HiveSessionProperties( } }, false), - booleanProperty( - PARQUET_OPTIMIZED_READER_ENABLED, - "Use optimized Parquet reader", - parquetReaderConfig.isOptimizedReaderEnabled(), - false), - booleanProperty( - PARQUET_OPTIMIZED_NESTED_READER_ENABLED, - "Use optimized Parquet reader for nested columns", - parquetReaderConfig.isOptimizedNestedReaderEnabled(), + dataSizeProperty( + PARQUET_SMALL_FILE_THRESHOLD, + "Parquet: Size below which a parquet file will be read entirely", + parquetReaderConfig.getSmallFileThreshold(), + value -> validateMaxDataSize(PARQUET_SMALL_FILE_THRESHOLD, value, DataSize.valueOf(PARQUET_READER_MAX_SMALL_FILE_THRESHOLD)), false), dataSizeProperty( PARQUET_WRITER_BLOCK_SIZE, @@ -872,6 +804,11 @@ public static int getParquetMaxReadBlockRowCount(ConnectorSession session) return session.getProperty(PARQUET_MAX_READ_BLOCK_ROW_COUNT, Integer.class); } + public static DataSize getParquetSmallFileThreshold(ConnectorSession session) + { + return session.getProperty(PARQUET_SMALL_FILE_THRESHOLD, DataSize.class); + } + public static boolean isParquetOptimizedReaderEnabled(ConnectorSession session) { return session.getProperty(PARQUET_OPTIMIZED_READER_ENABLED, Boolean.class); diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveSplit.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveSplit.java index 1167039edf67..69a0694d8b00 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveSplit.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveSplit.java @@ -23,39 +23,38 @@ import io.trino.spi.connector.ConnectorSplit; import java.util.List; +import java.util.Map; import java.util.Optional; import java.util.OptionalInt; -import java.util.Properties; import static com.google.common.base.MoreObjects.toStringHelper; import static com.google.common.base.Preconditions.checkArgument; import static io.airlift.slice.SizeOf.estimatedSizeOf; import static io.airlift.slice.SizeOf.instanceSize; import static io.airlift.slice.SizeOf.sizeOf; -import static io.trino.plugin.hive.util.HiveUtil.getDeserializerClassName; import static java.util.Objects.requireNonNull; public class HiveSplit implements ConnectorSplit { private static final int INSTANCE_SIZE = instanceSize(HiveSplit.class); + private static final int INTEGER_INSTANCE_SIZE = instanceSize(Integer.class); private final String path; private final long start; private final long length; private final long estimatedFileSize; private final long fileModifiedTime; - private final Properties schema; + private final Schema schema; private final List partitionKeys; private final List addresses; private final String partitionName; private final OptionalInt readBucketNumber; private final OptionalInt tableBucketNumber; private final boolean forceLocalScheduling; - private final TableToPartitionMapping tableToPartitionMapping; + private final Map hiveColumnCoercions; private final Optional bucketConversion; private final Optional bucketValidation; - private final boolean s3SelectPushdownEnabled; private final Optional acidInfo; private final SplitWeight splitWeight; @@ -67,18 +66,55 @@ public HiveSplit( @JsonProperty("length") long length, @JsonProperty("estimatedFileSize") long estimatedFileSize, @JsonProperty("fileModifiedTime") long fileModifiedTime, - @JsonProperty("schema") Properties schema, + @JsonProperty("schema") Schema schema, @JsonProperty("partitionKeys") List partitionKeys, - @JsonProperty("addresses") List addresses, @JsonProperty("readBucketNumber") OptionalInt readBucketNumber, @JsonProperty("tableBucketNumber") OptionalInt tableBucketNumber, @JsonProperty("forceLocalScheduling") boolean forceLocalScheduling, - @JsonProperty("tableToPartitionMapping") TableToPartitionMapping tableToPartitionMapping, + @JsonProperty("hiveColumnCoercions") Map hiveColumnCoercions, @JsonProperty("bucketConversion") Optional bucketConversion, @JsonProperty("bucketValidation") Optional bucketValidation, - @JsonProperty("s3SelectPushdownEnabled") boolean s3SelectPushdownEnabled, @JsonProperty("acidInfo") Optional acidInfo, @JsonProperty("splitWeight") SplitWeight splitWeight) + { + this( + partitionName, + path, + start, + length, + estimatedFileSize, + fileModifiedTime, + schema, + partitionKeys, + ImmutableList.of(), + readBucketNumber, + tableBucketNumber, + forceLocalScheduling, + hiveColumnCoercions, + bucketConversion, + bucketValidation, + acidInfo, + splitWeight); + } + + public HiveSplit( + String partitionName, + String path, + long start, + long length, + long estimatedFileSize, + long fileModifiedTime, + Schema schema, + List partitionKeys, + List addresses, + OptionalInt readBucketNumber, + OptionalInt tableBucketNumber, + boolean forceLocalScheduling, + Map hiveColumnCoercions, + Optional bucketConversion, + Optional bucketValidation, + Optional acidInfo, + SplitWeight splitWeight) { checkArgument(start >= 0, "start must be positive"); checkArgument(length >= 0, "length must be positive"); @@ -90,7 +126,7 @@ public HiveSplit( requireNonNull(addresses, "addresses is null"); requireNonNull(readBucketNumber, "readBucketNumber is null"); requireNonNull(tableBucketNumber, "tableBucketNumber is null"); - requireNonNull(tableToPartitionMapping, "tableToPartitionMapping is null"); + requireNonNull(hiveColumnCoercions, "hiveColumnCoercions is null"); requireNonNull(bucketConversion, "bucketConversion is null"); requireNonNull(bucketValidation, "bucketValidation is null"); requireNonNull(acidInfo, "acidInfo is null"); @@ -107,10 +143,9 @@ public HiveSplit( this.readBucketNumber = readBucketNumber; this.tableBucketNumber = tableBucketNumber; this.forceLocalScheduling = forceLocalScheduling; - this.tableToPartitionMapping = tableToPartitionMapping; + this.hiveColumnCoercions = ImmutableMap.copyOf(hiveColumnCoercions); this.bucketConversion = bucketConversion; this.bucketValidation = bucketValidation; - this.s3SelectPushdownEnabled = s3SelectPushdownEnabled; this.acidInfo = acidInfo; this.splitWeight = requireNonNull(splitWeight, "splitWeight is null"); } @@ -152,7 +187,7 @@ public long getFileModifiedTime() } @JsonProperty - public Properties getSchema() + public Schema getSchema() { return schema; } @@ -189,9 +224,9 @@ public boolean isForceLocalScheduling() } @JsonProperty - public TableToPartitionMapping getTableToPartitionMapping() + public Map getHiveColumnCoercions() { - return tableToPartitionMapping; + return hiveColumnCoercions; } @JsonProperty @@ -212,12 +247,6 @@ public boolean isRemotelyAccessible() return !forceLocalScheduling; } - @JsonProperty - public boolean isS3SelectPushdownEnabled() - { - return s3SelectPushdownEnabled; - } - @JsonProperty public Optional getAcidInfo() { @@ -236,13 +265,13 @@ public long getRetainedSizeInBytes() { return INSTANCE_SIZE + estimatedSizeOf(path) - + estimatedSizeOf(schema, key -> estimatedSizeOf((String) key), value -> estimatedSizeOf((String) value)) + + schema.getRetainedSizeInBytes() + estimatedSizeOf(partitionKeys, HivePartitionKey::getEstimatedSizeInBytes) + estimatedSizeOf(addresses, HostAddress::getRetainedSizeInBytes) + estimatedSizeOf(partitionName) + sizeOf(readBucketNumber) + sizeOf(tableBucketNumber) - + tableToPartitionMapping.getEstimatedSizeInBytes() + + estimatedSizeOf(hiveColumnCoercions, (Integer key) -> INTEGER_INSTANCE_SIZE, HiveTypeName::getEstimatedSizeInBytes) + sizeOf(bucketConversion, BucketConversion::getRetainedSizeInBytes) + sizeOf(bucketValidation, BucketValidation::getRetainedSizeInBytes) + sizeOf(acidInfo, AcidInfo::getRetainedSizeInBytes) @@ -260,8 +289,7 @@ public Object getInfo() .put("hosts", addresses) .put("forceLocalScheduling", forceLocalScheduling) .put("partitionName", partitionName) - .put("deserializerClassName", getDeserializerClassName(schema)) - .put("s3SelectPushdownEnabled", s3SelectPushdownEnabled) + .put("serializationLibraryName", schema.serializationLibraryName()) .buildOrThrow(); } diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveSplitManager.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveSplitManager.java index 077e5a52540d..334d25381219 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveSplitManager.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveSplitManager.java @@ -18,13 +18,13 @@ import com.google.common.collect.ImmutableMap; import com.google.common.collect.Lists; import com.google.common.collect.PeekingIterator; +import com.google.common.collect.Streams; import com.google.inject.Inject; import io.airlift.concurrent.BoundedExecutor; import io.airlift.stats.CounterStat; import io.airlift.units.DataSize; import io.trino.filesystem.TrinoFileSystemFactory; -import io.trino.hdfs.HdfsEnvironment; -import io.trino.hdfs.HdfsNamenodeStats; +import io.trino.filesystem.cache.CachingHostAddressProvider; import io.trino.plugin.hive.metastore.Column; import io.trino.plugin.hive.metastore.Partition; import io.trino.plugin.hive.metastore.SemiTransactionalHiveMetastore; @@ -54,6 +54,7 @@ import java.util.List; import java.util.Map; import java.util.Optional; +import java.util.Set; import java.util.concurrent.Executor; import java.util.concurrent.ExecutorService; import java.util.concurrent.RejectedExecutionException; @@ -63,6 +64,7 @@ import static com.google.common.base.Strings.isNullOrEmpty; import static com.google.common.base.Verify.verify; import static com.google.common.collect.ImmutableList.toImmutableList; +import static com.google.common.collect.ImmutableSet.toImmutableSet; import static com.google.common.collect.Iterators.peekingIterator; import static com.google.common.collect.Iterators.singletonIterator; import static com.google.common.collect.Iterators.transform; @@ -75,12 +77,10 @@ import static io.trino.plugin.hive.HiveSessionProperties.getDynamicFilteringWaitTimeout; import static io.trino.plugin.hive.HiveSessionProperties.getTimestampPrecision; import static io.trino.plugin.hive.HiveSessionProperties.isIgnoreAbsentPartitions; -import static io.trino.plugin.hive.HiveSessionProperties.isOptimizeSymlinkListing; import static io.trino.plugin.hive.HiveSessionProperties.isPropagateTableScanSortingProperties; import static io.trino.plugin.hive.HiveSessionProperties.isUseOrcColumnNames; import static io.trino.plugin.hive.HiveSessionProperties.isUseParquetColumnNames; import static io.trino.plugin.hive.HiveStorageFormat.getHiveStorageFormat; -import static io.trino.plugin.hive.TableToPartitionMapping.mapColumnsByIndex; import static io.trino.plugin.hive.metastore.MetastoreUtil.getProtectMode; import static io.trino.plugin.hive.metastore.MetastoreUtil.makePartitionName; import static io.trino.plugin.hive.metastore.MetastoreUtil.verifyOnline; @@ -104,8 +104,6 @@ public class HiveSplitManager private final HiveTransactionManager transactionManager; private final HivePartitionManager partitionManager; private final TrinoFileSystemFactory fileSystemFactory; - private final HdfsNamenodeStats hdfsNamenodeStats; - private final HdfsEnvironment hdfsEnvironment; private final Executor executor; private final int maxOutstandingSplits; private final DataSize maxOutstandingSplitsSize; @@ -117,6 +115,7 @@ public class HiveSplitManager private final boolean recursiveDfsWalkerEnabled; private final CounterStat highMemorySplitSourceCounter; private final TypeManager typeManager; + private final CachingHostAddressProvider cachingHostAddressProvider; private final int maxPartitionsPerScan; @Inject @@ -125,18 +124,15 @@ public HiveSplitManager( HiveTransactionManager transactionManager, HivePartitionManager partitionManager, TrinoFileSystemFactory fileSystemFactory, - HdfsNamenodeStats hdfsNamenodeStats, - HdfsEnvironment hdfsEnvironment, ExecutorService executorService, VersionEmbedder versionEmbedder, - TypeManager typeManager) + TypeManager typeManager, + CachingHostAddressProvider cachingHostAddressProvider) { this( transactionManager, partitionManager, fileSystemFactory, - hdfsNamenodeStats, - hdfsEnvironment, versionEmbedder.embedVersion(new BoundedExecutor(executorService, hiveConfig.getMaxSplitIteratorThreads())), new CounterStat(), hiveConfig.getMaxOutstandingSplits(), @@ -148,6 +144,7 @@ public HiveSplitManager( hiveConfig.getMaxSplitsPerSecond(), hiveConfig.getRecursiveDirWalkerEnabled(), typeManager, + cachingHostAddressProvider, hiveConfig.getMaxPartitionsPerScan()); } @@ -155,8 +152,6 @@ public HiveSplitManager( HiveTransactionManager transactionManager, HivePartitionManager partitionManager, TrinoFileSystemFactory fileSystemFactory, - HdfsNamenodeStats hdfsNamenodeStats, - HdfsEnvironment hdfsEnvironment, Executor executor, CounterStat highMemorySplitSourceCounter, int maxOutstandingSplits, @@ -168,13 +163,12 @@ public HiveSplitManager( @Nullable Integer maxSplitsPerSecond, boolean recursiveDfsWalkerEnabled, TypeManager typeManager, + CachingHostAddressProvider cachingHostAddressProvider, int maxPartitionsPerScan) { this.transactionManager = requireNonNull(transactionManager, "transactionManager is null"); this.partitionManager = requireNonNull(partitionManager, "partitionManager is null"); this.fileSystemFactory = requireNonNull(fileSystemFactory, "fileSystemFactory is null"); - this.hdfsNamenodeStats = requireNonNull(hdfsNamenodeStats, "hdfsNamenodeStats is null"); - this.hdfsEnvironment = requireNonNull(hdfsEnvironment, "hdfsEnvironment is null"); this.executor = new ErrorCodedExecutor(executor); this.highMemorySplitSourceCounter = requireNonNull(highMemorySplitSourceCounter, "highMemorySplitSourceCounter is null"); checkArgument(maxOutstandingSplits >= 1, "maxOutstandingSplits must be at least 1"); @@ -187,6 +181,7 @@ public HiveSplitManager( this.maxSplitsPerSecond = firstNonNull(maxSplitsPerSecond, Integer.MAX_VALUE); this.recursiveDfsWalkerEnabled = recursiveDfsWalkerEnabled; this.typeManager = requireNonNull(typeManager, "typeManager is null"); + this.cachingHostAddressProvider = requireNonNull(cachingHostAddressProvider, "cachingHostAddressProvider is null"); this.maxPartitionsPerScan = maxPartitionsPerScan; } @@ -223,13 +218,13 @@ public ConnectorSplitSource getSplits( Optional bucketFilter = hiveTable.getBucketFilter(); // validate bucket bucketed execution - Optional bucketHandle = hiveTable.getBucketHandle(); + Optional tablePartitioning = hiveTable.getTablePartitioning(); - bucketHandle.ifPresent(bucketing -> - verify(bucketing.getReadBucketCount() <= bucketing.getTableBucketCount(), + tablePartitioning.ifPresent(bucketing -> + verify(bucketing.partitioningHandle().getBucketCount() <= bucketing.tableBucketCount(), "readBucketCount (%s) is greater than the tableBucketCount (%s) which generally points to an issue in plan generation", - bucketing.getReadBucketCount(), - bucketing.getTableBucketCount())); + bucketing.partitioningHandle().getBucketCount(), + bucketing.tableBucketCount())); // get partitions Iterator partitions = partitionManager.getPartitions(metastore, hiveTable); @@ -242,12 +237,20 @@ public ConnectorSplitSource getSplits( return emptySplitSource(); } + Set neededColumnNames = Streams.concat(hiveTable.getProjectedColumns().stream(), hiveTable.getConstraintColumns().stream()) + .map(columnHandle -> ((HiveColumnHandle) columnHandle).getBaseColumnName()) // possible duplicates are handled by toImmutableSet at the end + .map(columnName -> columnName.toLowerCase(ENGLISH)) + .collect(toImmutableSet()); + Iterator hivePartitions = getPartitionMetadata( session, metastore, table, peekingIterator(partitions), - bucketHandle.map(HiveBucketHandle::toTableBucketProperty)); + tablePartitioning.map(HiveTablePartitioning::toTableBucketProperty), + neededColumnNames, + dynamicFilter, + hiveTable); HiveSplitLoader hiveSplitLoader = new BackgroundHiveSplitLoader( table, @@ -256,17 +259,14 @@ public ConnectorSplitSource getSplits( dynamicFilter, getDynamicFilteringWaitTimeout(session), typeManager, - createBucketSplitInfo(bucketHandle, bucketFilter), + createBucketSplitInfo(tablePartitioning, bucketFilter), session, fileSystemFactory, - hdfsEnvironment, - hdfsNamenodeStats, transactionalMetadata.getDirectoryLister(), executor, splitLoaderConcurrency, recursiveDfsWalkerEnabled, !hiveTable.getPartitionColumns().isEmpty() && isIgnoreAbsentPartitions(session), - isOptimizeSymlinkListing(session), metastore.getValidWriteIds(session, hiveTable) .map(value -> value.getTableValidWriteIdList(table.getDatabaseName() + "." + table.getTableName())), hiveTable.getMaxScannedFileSize(), @@ -283,6 +283,7 @@ public ConnectorSplitSource getSplits( hiveSplitLoader, executor, highMemorySplitSourceCounter, + cachingHostAddressProvider, hiveTable.isRecordScannedFiles()); hiveSplitLoader.start(splitSource); @@ -301,7 +302,10 @@ private Iterator getPartitionMetadata( SemiTransactionalHiveMetastore metastore, Table table, PeekingIterator hivePartitions, - Optional bucketProperty) + Optional bucketProperty, + Set neededColumnNames, + DynamicFilter dynamicFilter, + HiveTableHandle tableHandle) { if (!hivePartitions.hasNext()) { return emptyIterator(); @@ -311,7 +315,7 @@ private Iterator getPartitionMetadata( if (firstPartition.getPartitionId().equals(UNPARTITIONED_ID)) { hivePartitions.next(); checkArgument(!hivePartitions.hasNext(), "single partition is expected for unpartitioned table"); - return singletonIterator(new HivePartitionMetadata(firstPartition, Optional.empty(), TableToPartitionMapping.empty())); + return singletonIterator(new HivePartitionMetadata(firstPartition, Optional.empty(), ImmutableMap.of())); } HiveTimestampPrecision hiveTimestampPrecision = getTimestampPrecision(session); @@ -347,7 +351,8 @@ private Iterator getPartitionMetadata( table, bucketProperty, hivePartition, - partition.get())); + partition.get(), + neededColumnNames)); } return results.build(); @@ -365,7 +370,8 @@ private static HivePartitionMetadata toPartitionMetadata( Table table, Optional bucketProperty, HivePartition hivePartition, - Partition partition) + Partition partition, + Set neededColumnNames) { SchemaTableName tableName = table.getSchemaTableName(); String partName = makePartitionName(table, partition); @@ -388,7 +394,8 @@ private static HivePartitionMetadata toPartitionMetadata( if ((tableColumns == null) || (partitionColumns == null)) { throw new TrinoException(HIVE_INVALID_METADATA, format("Table '%s' or partition '%s' has null columns", tableName, partName)); } - TableToPartitionMapping tableToPartitionMapping = getTableToPartitionMapping(usePartitionColumnNames, typeManager, hiveTimestampPrecision, tableName, partName, tableColumns, partitionColumns); + + Map hiveColumnCoercions = getHiveColumnCoercions(usePartitionColumnNames, typeManager, hiveTimestampPrecision, tableName, partName, tableColumns, partitionColumns, neededColumnNames); if (bucketProperty.isPresent()) { HiveBucketProperty partitionBucketProperty = partition.getStorage().getBucketProperty() @@ -423,16 +430,28 @@ private static HivePartitionMetadata toPartitionMetadata( } } } - return new HivePartitionMetadata(hivePartition, Optional.of(partition), tableToPartitionMapping); + return new HivePartitionMetadata(hivePartition, Optional.of(partition), hiveColumnCoercions); } - private static TableToPartitionMapping getTableToPartitionMapping(boolean usePartitionColumnNames, TypeManager typeManager, HiveTimestampPrecision hiveTimestampPrecision, SchemaTableName tableName, String partName, List tableColumns, List partitionColumns) + private static Map getHiveColumnCoercions( + boolean usePartitionColumnNames, + TypeManager typeManager, + HiveTimestampPrecision hiveTimestampPrecision, + SchemaTableName tableName, + String partName, + List tableColumns, + List partitionColumns, + Set neededColumnNames) { if (usePartitionColumnNames) { - return getTableToPartitionMappingByColumnNames(typeManager, tableName, partName, tableColumns, partitionColumns, hiveTimestampPrecision); + return getHiveColumnCoercionsByColumnNames(typeManager, tableName, partName, tableColumns, partitionColumns, neededColumnNames, hiveTimestampPrecision); } ImmutableMap.Builder columnCoercions = ImmutableMap.builder(); for (int i = 0; i < min(partitionColumns.size(), tableColumns.size()); i++) { + if (!neededColumnNames.contains(tableColumns.get(i).getName().toLowerCase(ENGLISH))) { + // skip columns not used in the query + continue; + } HiveType tableType = tableColumns.get(i).getType(); HiveType partitionType = partitionColumns.get(i).getType(); if (!tableType.equals(partitionType)) { @@ -442,7 +461,7 @@ private static TableToPartitionMapping getTableToPartitionMapping(boolean usePar columnCoercions.put(i, partitionType.getHiveTypeName()); } } - return mapColumnsByIndex(columnCoercions.buildOrThrow()); + return columnCoercions.buildOrThrow(); } private static boolean isPartitionUsesColumnNames(ConnectorSession session, Optional storageFormat) @@ -458,16 +477,27 @@ private static boolean isPartitionUsesColumnNames(ConnectorSession session, Opti }; } - private static TableToPartitionMapping getTableToPartitionMappingByColumnNames(TypeManager typeManager, SchemaTableName tableName, String partName, List tableColumns, List partitionColumns, HiveTimestampPrecision hiveTimestampPrecision) + private static Map getHiveColumnCoercionsByColumnNames( + TypeManager typeManager, + SchemaTableName tableName, + String partName, + List tableColumns, + List partitionColumns, + Set neededColumnNames, + HiveTimestampPrecision hiveTimestampPrecision) { ImmutableMap.Builder partitionColumnIndexesBuilder = ImmutableMap.builderWithExpectedSize(partitionColumns.size()); for (int i = 0; i < partitionColumns.size(); i++) { - partitionColumnIndexesBuilder.put(partitionColumns.get(i).getName().toLowerCase(ENGLISH), i); + String columnName = partitionColumns.get(i).getName().toLowerCase(ENGLISH); + if (!neededColumnNames.contains(columnName)) { + // skip columns not used in the query + continue; + } + partitionColumnIndexesBuilder.put(columnName, i); } Map partitionColumnsByIndex = partitionColumnIndexesBuilder.buildOrThrow(); ImmutableMap.Builder columnCoercions = ImmutableMap.builder(); - ImmutableMap.Builder tableToPartitionColumns = ImmutableMap.builder(); for (int tableColumnIndex = 0; tableColumnIndex < tableColumns.size(); tableColumnIndex++) { Column tableColumn = tableColumns.get(tableColumnIndex); HiveType tableType = tableColumn.getType(); @@ -475,18 +505,17 @@ private static TableToPartitionMapping getTableToPartitionMappingByColumnNames(T if (partitionColumnIndex == null) { continue; } - tableToPartitionColumns.put(tableColumnIndex, partitionColumnIndex); Column partitionColumn = partitionColumns.get(partitionColumnIndex); HiveType partitionType = partitionColumn.getType(); if (!tableType.equals(partitionType)) { if (!canCoerce(typeManager, partitionType, tableType, hiveTimestampPrecision)) { throw tablePartitionColumnMismatchException(tableName, partName, tableColumn.getName(), tableType, partitionColumn.getName(), partitionType); } - columnCoercions.put(partitionColumnIndex, partitionType.getHiveTypeName()); + columnCoercions.put(tableColumnIndex, partitionType.getHiveTypeName()); } } - return new TableToPartitionMapping(Optional.of(tableToPartitionColumns.buildOrThrow()), columnCoercions.buildOrThrow()); + return columnCoercions.buildOrThrow(); } private static TrinoException tablePartitionColumnMismatchException(SchemaTableName tableName, String partName, String tableColumnName, HiveType tableType, String partitionColumnName, HiveType partitionType) diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveSplitSource.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveSplitSource.java index 1c8affff2a00..94a66bf09101 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveSplitSource.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveSplitSource.java @@ -19,6 +19,7 @@ import io.airlift.log.Logger; import io.airlift.stats.CounterStat; import io.airlift.units.DataSize; +import io.trino.filesystem.cache.CachingHostAddressProvider; import io.trino.plugin.hive.InternalHiveSplit.InternalHiveBlock; import io.trino.plugin.hive.util.AsyncQueue; import io.trino.plugin.hive.util.AsyncQueue.BorrowResult; @@ -84,6 +85,7 @@ class HiveSplitSource private final CounterStat highMemorySplitSourceCounter; private final AtomicBoolean loggedHighMemoryWarning = new AtomicBoolean(); private final HiveSplitWeightProvider splitWeightProvider; + private final CachingHostAddressProvider cachingHostAddressProvider; private final boolean recordScannedFiles; private final ImmutableList.Builder scannedFilePaths = ImmutableList.builder(); @@ -98,6 +100,7 @@ private HiveSplitSource( HiveSplitLoader splitLoader, AtomicReference stateReference, CounterStat highMemorySplitSourceCounter, + CachingHostAddressProvider cachingHostAddressProvider, boolean recordScannedFiles) { requireNonNull(session, "session is null"); @@ -114,6 +117,7 @@ private HiveSplitSource( this.maxInitialSplitSize = getMaxInitialSplitSize(session); this.remainingInitialSplits = new AtomicInteger(maxInitialSplits); this.splitWeightProvider = isSizeBasedSplitWeightsEnabled(session) ? new SizeBasedSplitWeightProvider(getMinimumAssignedSplitWeight(session), maxSplitSize) : HiveSplitWeightProvider.uniformStandardWeightProvider(); + this.cachingHostAddressProvider = requireNonNull(cachingHostAddressProvider, "cachingHostAddressProvider is null"); this.recordScannedFiles = recordScannedFiles; } @@ -128,6 +132,7 @@ public static HiveSplitSource allAtOnce( HiveSplitLoader splitLoader, Executor executor, CounterStat highMemorySplitSourceCounter, + CachingHostAddressProvider cachingHostAddressProvider, boolean recordScannedFiles) { AtomicReference stateReference = new AtomicReference<>(State.initial()); @@ -168,6 +173,7 @@ public boolean isFinished() splitLoader, stateReference, highMemorySplitSourceCounter, + cachingHostAddressProvider, recordScannedFiles); } @@ -305,14 +311,13 @@ else if (maxSplitBytes * 2 >= remainingBlockBytes) { internalSplit.getFileModifiedTime(), internalSplit.getSchema(), internalSplit.getPartitionKeys(), - block.getAddresses(), + cachingHostAddressProvider.getHosts(internalSplit.getPath(), block.getAddresses()), internalSplit.getReadBucketNumber(), internalSplit.getTableBucketNumber(), internalSplit.isForceLocalScheduling(), - internalSplit.getTableToPartitionMapping(), + internalSplit.getHiveColumnCoercions(), internalSplit.getBucketConversion(), internalSplit.getBucketValidation(), - internalSplit.isS3SelectPushdownEnabled(), internalSplit.getAcidInfo(), splitWeightProvider.weightForSplitSizeInBytes(splitBytes))); diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveStorageFormat.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveStorageFormat.java index 203212afdd42..01f154ce3b26 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveStorageFormat.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveStorageFormat.java @@ -158,6 +158,11 @@ public boolean isSplittable(String path) }; } + public StorageFormat toStorageFormat() + { + return StorageFormat.create(serde, inputFormat, outputFormat); + } + public void validateColumns(List handles) { if (this == AVRO) { @@ -214,4 +219,13 @@ public static Optional getHiveStorageFormat(StorageFormat sto return Optional.ofNullable(HIVE_STORAGE_FORMATS.get( new SerdeAndInputFormat(storageFormat.getSerde(), storageFormat.getInputFormat()))); } + + public String humanName() + { + return switch (this) { + case AVRO -> "Avro"; + case PARQUET -> "Parquet"; + default -> toString(); + }; + } } diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveTableExecuteHandle.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveTableExecuteHandle.java index 9776d4d61c4e..490ec4972ffc 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveTableExecuteHandle.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveTableExecuteHandle.java @@ -43,7 +43,7 @@ public HiveTableExecuteHandle( @JsonProperty("inputColumns") List inputColumns, @JsonProperty("pageSinkMetadata") HivePageSinkMetadata pageSinkMetadata, @JsonProperty("locationHandle") LocationHandle locationHandle, - @JsonProperty("bucketProperty") Optional bucketProperty, + @JsonProperty("bucketInfo") Optional bucketInfo, @JsonProperty("tableStorageFormat") HiveStorageFormat tableStorageFormat, @JsonProperty("partitionStorageFormat") HiveStorageFormat partitionStorageFormat, @JsonProperty("transaction") AcidTransaction transaction, @@ -55,14 +55,14 @@ public HiveTableExecuteHandle( inputColumns, pageSinkMetadata, locationHandle, - bucketProperty, + bucketInfo, tableStorageFormat, partitionStorageFormat, transaction, retriesEnabled); // todo to be added soon - verify(bucketProperty.isEmpty(), "bucketed tables not supported yet"); + verify(bucketInfo.isEmpty(), "bucketed tables not supported yet"); this.procedureName = requireNonNull(procedureName, "procedureName is null"); this.writeDeclarationId = requireNonNull(writeDeclarationId, "writeDeclarationId is null"); @@ -98,7 +98,7 @@ public HiveTableExecuteHandle withWriteDeclarationId(String writeDeclarationId) getInputColumns(), getPageSinkMetadata(), getLocationHandle(), - getBucketProperty(), + getBucketInfo(), getTableStorageFormat(), getPartitionStorageFormat(), getTransaction(), diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveTableHandle.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveTableHandle.java index 8df28e5c557a..ce7aec7785cf 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveTableHandle.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveTableHandle.java @@ -50,7 +50,7 @@ public class HiveTableHandle private final Optional> partitions; private final TupleDomain compactEffectivePredicate; private final TupleDomain enforcedConstraint; - private final Optional bucketHandle; + private final Optional tablePartitioning; private final Optional bucketFilter; private final Optional>> analyzePartitionValues; private final Set constraintColumns; @@ -67,7 +67,7 @@ public HiveTableHandle( @JsonProperty("dataColumns") List dataColumns, @JsonProperty("compactEffectivePredicate") TupleDomain compactEffectivePredicate, @JsonProperty("enforcedConstraint") TupleDomain enforcedConstraint, - @JsonProperty("bucketHandle") Optional bucketHandle, + @JsonProperty("bucketHandle") Optional tablePartitioning, @JsonProperty("bucketFilter") Optional bucketFilter, @JsonProperty("analyzePartitionValues") Optional>> analyzePartitionValues, @JsonProperty("transaction") AcidTransaction transaction) @@ -82,7 +82,7 @@ public HiveTableHandle( Optional.empty(), compactEffectivePredicate, enforcedConstraint, - bucketHandle, + tablePartitioning, bucketFilter, analyzePartitionValues, ImmutableSet.of(), @@ -97,7 +97,7 @@ public HiveTableHandle( Map tableParameters, List partitionColumns, List dataColumns, - Optional bucketHandle) + Optional tablePartitioning) { this( schemaName, @@ -109,7 +109,7 @@ public HiveTableHandle( Optional.empty(), TupleDomain.all(), TupleDomain.all(), - bucketHandle, + tablePartitioning, Optional.empty(), Optional.empty(), ImmutableSet.of(), @@ -128,7 +128,7 @@ private HiveTableHandle( Optional> partitions, TupleDomain compactEffectivePredicate, TupleDomain enforcedConstraint, - Optional bucketHandle, + Optional tablePartitioning, Optional bucketFilter, Optional>> analyzePartitionValues, Set constraintColumns, @@ -146,7 +146,7 @@ private HiveTableHandle( partitions, compactEffectivePredicate, enforcedConstraint, - bucketHandle, + tablePartitioning, bucketFilter, analyzePartitionValues, constraintColumns, @@ -166,7 +166,7 @@ public HiveTableHandle( Optional> partitions, TupleDomain compactEffectivePredicate, TupleDomain enforcedConstraint, - Optional bucketHandle, + Optional tablePartitioning, Optional bucketFilter, Optional>> analyzePartitionValues, Set constraintColumns, @@ -185,7 +185,7 @@ public HiveTableHandle( this.partitions = partitions.map(ImmutableList::copyOf); this.compactEffectivePredicate = requireNonNull(compactEffectivePredicate, "compactEffectivePredicate is null"); this.enforcedConstraint = requireNonNull(enforcedConstraint, "enforcedConstraint is null"); - this.bucketHandle = requireNonNull(bucketHandle, "bucketHandle is null"); + this.tablePartitioning = requireNonNull(tablePartitioning, "tablePartitioning is null"); this.bucketFilter = requireNonNull(bucketFilter, "bucketFilter is null"); this.analyzePartitionValues = analyzePartitionValues.map(ImmutableList::copyOf); this.constraintColumns = ImmutableSet.copyOf(requireNonNull(constraintColumns, "constraintColumns is null")); @@ -207,7 +207,7 @@ public HiveTableHandle withAnalyzePartitionValues(List> analyzePart partitions, compactEffectivePredicate, enforcedConstraint, - bucketHandle, + tablePartitioning, bucketFilter, Optional.of(analyzePartitionValues), constraintColumns, @@ -229,7 +229,7 @@ public HiveTableHandle withTransaction(AcidTransaction transaction) partitions, compactEffectivePredicate, enforcedConstraint, - bucketHandle, + tablePartitioning, bucketFilter, analyzePartitionValues, constraintColumns, @@ -251,7 +251,7 @@ public HiveTableHandle withProjectedColumns(Set projectedColumns) partitions, compactEffectivePredicate, enforcedConstraint, - bucketHandle, + tablePartitioning, bucketFilter, analyzePartitionValues, constraintColumns, @@ -273,7 +273,7 @@ public HiveTableHandle withRecordScannedFiles(boolean recordScannedFiles) partitions, compactEffectivePredicate, enforcedConstraint, - bucketHandle, + tablePartitioning, bucketFilter, analyzePartitionValues, constraintColumns, @@ -295,7 +295,7 @@ public HiveTableHandle withMaxScannedFileSize(Optional maxScannedFileSize) partitions, compactEffectivePredicate, enforcedConstraint, - bucketHandle, + tablePartitioning, bucketFilter, analyzePartitionValues, constraintColumns, @@ -372,9 +372,9 @@ public TupleDomain getEnforcedConstraint() } @JsonProperty - public Optional getBucketHandle() + public Optional getTablePartitioning() { - return bucketHandle; + return tablePartitioning; } @JsonProperty @@ -463,7 +463,7 @@ public boolean equals(Object o) Objects.equals(partitions, that.partitions) && Objects.equals(compactEffectivePredicate, that.compactEffectivePredicate) && Objects.equals(enforcedConstraint, that.enforcedConstraint) && - Objects.equals(bucketHandle, that.bucketHandle) && + Objects.equals(tablePartitioning, that.tablePartitioning) && Objects.equals(bucketFilter, that.bucketFilter) && Objects.equals(analyzePartitionValues, that.analyzePartitionValues) && Objects.equals(transaction, that.transaction) && @@ -482,7 +482,7 @@ public int hashCode() partitions, compactEffectivePredicate, enforcedConstraint, - bucketHandle, + tablePartitioning, bucketFilter, analyzePartitionValues, transaction, @@ -494,11 +494,11 @@ public String toString() { StringBuilder builder = new StringBuilder(); builder.append(schemaName).append(":").append(tableName); - bucketHandle.ifPresent(bucket -> { - builder.append(" buckets=").append(bucket.getReadBucketCount()); - if (!bucket.getSortedBy().isEmpty()) { + tablePartitioning.ifPresent(bucket -> { + builder.append(" buckets=").append(bucket.partitioningHandle().getBucketCount()); + if (!bucket.sortedBy().isEmpty()) { builder.append(" sorted_by=") - .append(bucket.getSortedBy().stream() + .append(bucket.sortedBy().stream() .map(HiveUtil::sortingColumnToString) .collect(joining(", ", "[", "]"))); } diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveTablePartitioning.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveTablePartitioning.java new file mode 100644 index 000000000000..61fd0510df6d --- /dev/null +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveTablePartitioning.java @@ -0,0 +1,95 @@ +/* + * 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.hive; + +import com.google.common.collect.ImmutableList; +import io.trino.plugin.hive.metastore.SortingColumn; +import io.trino.plugin.hive.util.HiveBucketing; + +import java.util.List; + +import static com.google.common.base.Preconditions.checkArgument; +import static com.google.common.collect.ImmutableList.toImmutableList; +import static java.lang.String.format; +import static java.util.Objects.requireNonNull; +import static java.util.stream.Collectors.toList; + +/** + * @param active is the partitioning enabled for reading + * @param tableBucketCount number of buckets in the table, as specified in table metadata + * @param forWrite if true, the partitioning is being used for writing and cannot be changed + */ +public record HiveTablePartitioning( + boolean active, + List columns, + HivePartitioningHandle partitioningHandle, + int tableBucketCount, + List sortedBy, + boolean forWrite) +{ + public HiveTablePartitioning + { + columns = ImmutableList.copyOf(requireNonNull(columns, "columns is null")); + columns.forEach(column -> checkArgument(column.isBaseColumn(), format("projected column %s is not allowed for bucketing", column))); + checkArgument(columns.stream().map(HiveColumnHandle::getHiveType).toList().equals(partitioningHandle.getHiveTypes()), "columns do not match partitioning handle"); + checkArgument(tableBucketCount > 0, "tableBucketCount must be greater than zero"); + checkArgument(tableBucketCount >= partitioningHandle.getBucketCount(), "tableBucketCount must be greater than or equal to partitioningHandle.bucketCount"); + sortedBy = ImmutableList.copyOf(requireNonNull(sortedBy, "sortedBy is null")); + checkArgument(!forWrite || active, "Partitioning must be active for write"); + checkArgument(forWrite || !partitioningHandle.isUsePartitionedBucketing(), "Partitioned bucketing is only supported for write"); + } + + public HiveTablePartitioning( + boolean active, + HiveBucketing.BucketingVersion bucketingVersion, + int bucketCount, + List columns, + boolean usePartitionedBucketing, + List sortedBy, + boolean forWrite) + { + this(active, + columns, + new HivePartitioningHandle( + bucketingVersion, + bucketCount, + columns.stream() + .map(HiveColumnHandle::getHiveType) + .collect(toImmutableList()), + usePartitionedBucketing), + bucketCount, + sortedBy, + forWrite); + } + + public HiveTablePartitioning withActivePartitioning() + { + return new HiveTablePartitioning(true, columns, partitioningHandle, tableBucketCount, sortedBy, forWrite); + } + + public HiveTablePartitioning withPartitioningHandle(HivePartitioningHandle hivePartitioningHandle) + { + return new HiveTablePartitioning(active, columns, hivePartitioningHandle, tableBucketCount, sortedBy, forWrite); + } + + public HiveBucketProperty toTableBucketProperty() + { + return new HiveBucketProperty( + columns.stream() + .map(HiveColumnHandle::getName) + .collect(toList()), + tableBucketCount, + sortedBy); + } +} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveTableProperties.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveTableProperties.java index 51a70f9bc7e6..8a8812e90aac 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveTableProperties.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveTableProperties.java @@ -30,9 +30,9 @@ import java.util.Optional; import static com.google.common.collect.ImmutableList.toImmutableList; -import static io.trino.plugin.hive.aws.athena.PartitionProjectionProperties.PARTITION_PROJECTION_ENABLED; -import static io.trino.plugin.hive.aws.athena.PartitionProjectionProperties.PARTITION_PROJECTION_IGNORE; -import static io.trino.plugin.hive.aws.athena.PartitionProjectionProperties.PARTITION_PROJECTION_LOCATION_TEMPLATE; +import static io.trino.plugin.hive.projection.PartitionProjectionProperties.PARTITION_PROJECTION_ENABLED; +import static io.trino.plugin.hive.projection.PartitionProjectionProperties.PARTITION_PROJECTION_IGNORE; +import static io.trino.plugin.hive.projection.PartitionProjectionProperties.PARTITION_PROJECTION_LOCATION_TEMPLATE; import static io.trino.plugin.hive.util.HiveBucketing.BucketingVersion.BUCKETING_V1; import static io.trino.plugin.hive.util.HiveBucketing.BucketingVersion.BUCKETING_V2; import static io.trino.spi.StandardErrorCode.INVALID_TABLE_PROPERTY; @@ -66,6 +66,7 @@ public class HiveTableProperties public static final String CSV_SEPARATOR = "csv_separator"; public static final String CSV_QUOTE = "csv_quote"; public static final String CSV_ESCAPE = "csv_escape"; + public static final String PARQUET_BLOOM_FILTER_COLUMNS = "parquet_bloom_filter_columns"; public static final String REGEX_PATTERN = "regex"; public static final String REGEX_CASE_INSENSITIVE = "regex_case_insensitive"; public static final String TRANSACTIONAL = "transactional"; @@ -244,12 +245,12 @@ public static List getPartitionedBy(Map tableProperties) return partitionedBy == null ? ImmutableList.of() : ImmutableList.copyOf(partitionedBy); } - public static Optional getBucketProperty(Map tableProperties) + public static Optional getBucketInfo(Map tableProperties) { List bucketedBy = getBucketedBy(tableProperties); List sortedBy = getSortedBy(tableProperties); int bucketCount = (Integer) tableProperties.get(BUCKET_COUNT_PROPERTY); - if ((bucketedBy.isEmpty()) && (bucketCount == 0)) { + if (bucketedBy.isEmpty() && (bucketCount == 0)) { if (!sortedBy.isEmpty()) { throw new TrinoException(INVALID_TABLE_PROPERTY, format("%s may be specified only when %s is specified", SORTED_BY_PROPERTY, BUCKETED_BY_PROPERTY)); } @@ -262,7 +263,7 @@ public static Optional getBucketProperty(Map throw new TrinoException(INVALID_TABLE_PROPERTY, format("%s and %s must be specified together", BUCKETED_BY_PROPERTY, BUCKET_COUNT_PROPERTY)); } BucketingVersion bucketingVersion = getBucketingVersion(tableProperties); - return Optional.of(new HiveBucketProperty(bucketedBy, bucketingVersion, bucketCount, sortedBy)); + return Optional.of(new HiveWritableTableHandle.BucketInfo(bucketedBy, bucketingVersion, bucketCount, sortedBy)); } public static BucketingVersion getBucketingVersion(Map tableProperties) @@ -295,6 +296,12 @@ public static List getOrcBloomFilterColumns(Map tablePro return (List) tableProperties.get(ORC_BLOOM_FILTER_COLUMNS); } + @SuppressWarnings("unchecked") + public static List getParquetBloomFilterColumns(Map tableProperties) + { + return (List) tableProperties.get(PARQUET_BLOOM_FILTER_COLUMNS); + } + public static Double getOrcBloomFilterFpp(Map tableProperties) { return (Double) tableProperties.get(ORC_BLOOM_FILTER_FPP); diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveType.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveType.java index d4c9b7e851d7..c21d8ec048b4 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveType.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveType.java @@ -308,6 +308,11 @@ else if (typeInfo instanceof UnionTypeInfo unionTypeInfo) { return dereferenceNames.build(); } + public static HiveType fromTypeInfo(TypeInfo typeInfo) + { + return new HiveType(typeInfo); + } + public long getRetainedSizeInBytes() { return INSTANCE_SIZE + hiveTypeName.getEstimatedSizeInBytes() + typeInfo.getRetainedSizeInBytes(); diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveUpdateHandle.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveUpdateHandle.java index 9d1fb06f5fbf..2f49efb30ef7 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveUpdateHandle.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveUpdateHandle.java @@ -18,7 +18,6 @@ import io.trino.plugin.hive.util.HiveBucketing.BucketingVersion; import java.util.List; -import java.util.OptionalInt; public class HiveUpdateHandle extends HivePartitioningHandle @@ -28,9 +27,8 @@ public HiveUpdateHandle( @JsonProperty("bucketingVersion") BucketingVersion bucketingVersion, @JsonProperty("bucketCount") int bucketCount, @JsonProperty("hiveBucketTypes") List hiveTypes, - @JsonProperty("maxCompatibleBucketCount") OptionalInt maxCompatibleBucketCount, @JsonProperty("usePartitionedBucketing") boolean usePartitionedBucketing) { - super(bucketingVersion, bucketCount, hiveTypes, maxCompatibleBucketCount, usePartitionedBucketing); + super(bucketingVersion, bucketCount, hiveTypes, usePartitionedBucketing); } } diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveWritableTableHandle.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveWritableTableHandle.java index b0adba504dca..84d6108b1f47 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveWritableTableHandle.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveWritableTableHandle.java @@ -18,11 +18,15 @@ import com.google.common.collect.ImmutableList; import io.trino.plugin.hive.acid.AcidTransaction; import io.trino.plugin.hive.metastore.HivePageSinkMetadata; +import io.trino.plugin.hive.metastore.SortingColumn; +import io.trino.plugin.hive.metastore.Table; +import io.trino.plugin.hive.util.HiveBucketing; import io.trino.spi.connector.SchemaTableName; import java.util.List; import java.util.Optional; +import static io.trino.plugin.hive.util.HiveBucketing.getBucketingVersion; import static java.util.Objects.requireNonNull; public class HiveWritableTableHandle @@ -32,7 +36,7 @@ public class HiveWritableTableHandle private final List inputColumns; private final HivePageSinkMetadata pageSinkMetadata; private final LocationHandle locationHandle; - private final Optional bucketProperty; + private final Optional bucketInfo; private final HiveStorageFormat tableStorageFormat; private final HiveStorageFormat partitionStorageFormat; private final AcidTransaction transaction; @@ -44,7 +48,7 @@ public HiveWritableTableHandle( List inputColumns, HivePageSinkMetadata pageSinkMetadata, LocationHandle locationHandle, - Optional bucketProperty, + Optional bucketInfo, HiveStorageFormat tableStorageFormat, HiveStorageFormat partitionStorageFormat, AcidTransaction transaction, @@ -55,7 +59,7 @@ public HiveWritableTableHandle( this.inputColumns = ImmutableList.copyOf(requireNonNull(inputColumns, "inputColumns is null")); this.pageSinkMetadata = requireNonNull(pageSinkMetadata, "pageSinkMetadata is null"); this.locationHandle = requireNonNull(locationHandle, "locationHandle is null"); - this.bucketProperty = requireNonNull(bucketProperty, "bucketProperty is null"); + this.bucketInfo = requireNonNull(bucketInfo, "bucketInfo is null"); this.tableStorageFormat = requireNonNull(tableStorageFormat, "tableStorageFormat is null"); this.partitionStorageFormat = requireNonNull(partitionStorageFormat, "partitionStorageFormat is null"); this.transaction = requireNonNull(transaction, "transaction is null"); @@ -99,9 +103,9 @@ public LocationHandle getLocationHandle() } @JsonProperty - public Optional getBucketProperty() + public Optional getBucketInfo() { - return bucketProperty; + return bucketInfo; } @JsonProperty @@ -139,4 +143,28 @@ public String toString() { return schemaName + "." + tableName; } + + public record BucketInfo(List bucketedBy, HiveBucketing.BucketingVersion bucketingVersion, int bucketCount, List sortedBy) + { + public BucketInfo + { + bucketedBy = ImmutableList.copyOf(requireNonNull(bucketedBy, "bucketedBy is null")); + requireNonNull(bucketingVersion, "bucketingVersion is null"); + sortedBy = ImmutableList.copyOf(requireNonNull(sortedBy, "sortedBy is null")); + } + + public static Optional createBucketInfo(Table table) + { + if (table.getStorage().getBucketProperty().isEmpty()) { + return Optional.empty(); + } + HiveBucketProperty bucketProperty = table.getStorage().getBucketProperty().get(); + HiveBucketing.BucketingVersion bucketingVersion = getBucketingVersion(table.getParameters()); + return Optional.of(new BucketInfo( + bucketProperty.getBucketedBy(), + bucketingVersion, + bucketProperty.getBucketCount(), + bucketProperty.getSortedBy())); + } + } } diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveWriter.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveWriter.java index 7628b45362c5..a0f8a6359fb9 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveWriter.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveWriter.java @@ -19,7 +19,6 @@ import java.io.Closeable; import java.util.Optional; -import java.util.function.Consumer; import static com.google.common.base.MoreObjects.toStringHelper; import static java.util.Objects.requireNonNull; @@ -32,7 +31,6 @@ public class HiveWriter private final String fileName; private final String writePath; private final String targetPath; - private final Consumer onCommit; private final HiveWriterStats hiveWriterStats; private long rowCount; @@ -45,7 +43,6 @@ public HiveWriter( String fileName, String writePath, String targetPath, - Consumer onCommit, HiveWriterStats hiveWriterStats) { this.fileWriter = requireNonNull(fileWriter, "fileWriter is null"); @@ -54,7 +51,6 @@ public HiveWriter( this.fileName = requireNonNull(fileName, "fileName is null"); this.writePath = requireNonNull(writePath, "writePath is null"); this.targetPath = requireNonNull(targetPath, "targetPath is null"); - this.onCommit = requireNonNull(onCommit, "onCommit is null"); this.hiveWriterStats = requireNonNull(hiveWriterStats, "hiveWriterStats is null"); } @@ -89,9 +85,7 @@ public void append(Page dataPage) public Closeable commit() { - Closeable rollbackAction = fileWriter.commit(); - onCommit.accept(this); - return rollbackAction; + return fileWriter.commit(); } long getValidationCpuNanos() diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveWriterFactory.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveWriterFactory.java index 0402d29933f4..5470f4c206a4 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveWriterFactory.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveWriterFactory.java @@ -19,13 +19,11 @@ import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; import com.google.common.collect.Sets; -import io.airlift.event.client.EventClient; import io.airlift.units.DataSize; import io.trino.filesystem.Location; import io.trino.filesystem.TrinoFileSystem; import io.trino.filesystem.TrinoFileSystemFactory; -import io.trino.hdfs.HdfsContext; -import io.trino.hdfs.HdfsEnvironment; +import io.trino.hive.formats.compression.CompressionKind; import io.trino.plugin.hive.HiveSessionProperties.InsertExistingPartitionsBehavior; import io.trino.plugin.hive.LocationService.WriteInfo; import io.trino.plugin.hive.PartitionUpdate.UpdateMode; @@ -38,7 +36,6 @@ import io.trino.plugin.hive.metastore.Table; import io.trino.plugin.hive.orc.OrcFileWriterFactory; import io.trino.plugin.hive.util.HiveWriteUtils; -import io.trino.spi.NodeManager; import io.trino.spi.Page; import io.trino.spi.PageSorter; import io.trino.spi.TrinoException; @@ -47,38 +44,23 @@ import io.trino.spi.type.RowType; import io.trino.spi.type.Type; import io.trino.spi.type.TypeManager; -import org.apache.hadoop.conf.Configuration; -import org.apache.hadoop.fs.Path; -import org.apache.hadoop.hive.conf.HiveConf; -import org.apache.hadoop.io.compress.CompressionCodec; -import org.apache.hadoop.io.compress.DefaultCodec; -import org.apache.hadoop.mapred.JobConf; -import org.apache.hadoop.util.ReflectionUtils; -import org.joda.time.DateTimeZone; import java.io.IOException; -import java.security.Principal; import java.util.ArrayList; import java.util.HashMap; import java.util.HashSet; import java.util.List; import java.util.Map; -import java.util.Map.Entry; import java.util.Optional; import java.util.OptionalInt; -import java.util.Properties; import java.util.Set; import java.util.UUID; -import java.util.function.Consumer; import java.util.regex.Matcher; import java.util.regex.Pattern; import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.collect.ImmutableList.toImmutableList; -import static com.google.common.collect.ImmutableMap.toImmutableMap; -import static com.google.common.collect.Maps.immutableEntry; import static com.google.common.collect.MoreCollectors.onlyElement; -import static io.trino.hdfs.ConfigurationUtils.toJobConf; import static io.trino.plugin.hive.HiveCompressionCodecs.selectCompressionCodec; import static io.trino.plugin.hive.HiveErrorCode.HIVE_FILESYSTEM_ERROR; import static io.trino.plugin.hive.HiveErrorCode.HIVE_INVALID_METADATA; @@ -97,7 +79,6 @@ import static io.trino.plugin.hive.util.AcidTables.deltaSubdir; import static io.trino.plugin.hive.util.AcidTables.isFullAcidTable; import static io.trino.plugin.hive.util.AcidTables.isInsertOnlyTable; -import static io.trino.plugin.hive.util.CompressionConfigUtil.configureCompression; import static io.trino.plugin.hive.util.HiveClassNames.HIVE_IGNORE_KEY_OUTPUT_FORMAT_CLASS; import static io.trino.plugin.hive.util.HiveUtil.getColumnNames; import static io.trino.plugin.hive.util.HiveUtil.getColumnTypes; @@ -116,7 +97,6 @@ import static java.util.stream.Collectors.joining; import static java.util.stream.Collectors.toList; import static java.util.stream.Collectors.toMap; -import static org.apache.hadoop.hive.conf.HiveConf.ConfVars.COMPRESSRESULT; public class HiveWriterFactory { @@ -147,7 +127,6 @@ public class HiveWriterFactory private final HivePageSinkMetadataProvider pageSinkMetadataProvider; private final TypeManager typeManager; private final PageSorter pageSorter; - private final JobConf conf; private final Table table; private final DataSize sortBufferSize; @@ -155,16 +134,11 @@ public class HiveWriterFactory private final boolean sortedWritingTempStagingPathEnabled; private final String sortedWritingTempStagingPath; private final InsertExistingPartitionsBehavior insertExistingPartitionsBehavior; - private final DateTimeZone parquetTimeZone; private final ConnectorSession session; private final OptionalInt bucketCount; private final List sortedBy; - private final NodeManager nodeManager; - private final EventClient eventClient; - private final Map sessionProperties; - private final HiveWriterStats hiveWriterStats; private final Optional rowType; private final Optional hiveRowtype; @@ -187,15 +161,10 @@ public HiveWriterFactory( String queryId, HivePageSinkMetadataProvider pageSinkMetadataProvider, TypeManager typeManager, - HdfsEnvironment hdfsEnvironment, PageSorter pageSorter, DataSize sortBufferSize, int maxOpenSortFiles, - DateTimeZone parquetTimeZone, ConnectorSession session, - NodeManager nodeManager, - EventClient eventClient, - HiveSessionProperties hiveSessionProperties, HiveWriterStats hiveWriterStats, boolean sortedWritingTempStagingPathEnabled, String sortedWritingTempStagingPath) @@ -223,7 +192,6 @@ public HiveWriterFactory( this.sortedWritingTempStagingPathEnabled = sortedWritingTempStagingPathEnabled; this.sortedWritingTempStagingPath = requireNonNull(sortedWritingTempStagingPath, "sortedWritingTempStagingPath is null"); this.insertExistingPartitionsBehavior = getInsertExistingPartitionsBehavior(session); - this.parquetTimeZone = requireNonNull(parquetTimeZone, "parquetTimeZone is null"); // divide input columns into partition and data columns ImmutableList.Builder partitionColumnNames = ImmutableList.builder(); @@ -275,31 +243,7 @@ public HiveWriterFactory( } this.sortedBy = ImmutableList.copyOf(requireNonNull(sortedBy, "sortedBy is null")); - this.session = requireNonNull(session, "session is null"); - this.nodeManager = requireNonNull(nodeManager, "nodeManager is null"); - this.eventClient = requireNonNull(eventClient, "eventClient is null"); - - requireNonNull(hiveSessionProperties, "hiveSessionProperties is null"); - this.sessionProperties = hiveSessionProperties.getSessionProperties().stream() - .map(propertyMetadata -> immutableEntry( - propertyMetadata.getName(), - session.getProperty(propertyMetadata.getName(), propertyMetadata.getJavaType()))) - // The session properties collected here are used for events only. Filter out nulls to avoid problems with downstream consumers - .filter(entry -> entry.getValue() != null) - .collect(toImmutableMap(Entry::getKey, entry -> entry.getValue().toString())); - - Configuration conf = hdfsEnvironment.getConfiguration(new HdfsContext(session), new Path(writePath.toString())); - this.conf = toJobConf(conf); - - // make sure the FileSystem is created with the correct Configuration object - try { - hdfsEnvironment.getFileSystem(session.getIdentity(), new Path(writePath.toString()), conf); - } - catch (IOException e) { - throw new TrinoException(HIVE_FILESYSTEM_ERROR, "Failed getting FileSystem: " + writePath, e); - } - this.hiveWriterStats = requireNonNull(hiveWriterStats, "hiveWriterStats is null"); } @@ -330,7 +274,7 @@ public HiveWriter createWriter(Page partitionColumns, int position, OptionalInt } UpdateMode updateMode; - Properties schema; + Map schema = new HashMap<>(); WriteInfo writeInfo; StorageFormat outputStorageFormat; HiveCompressionCodec compressionCodec; @@ -339,11 +283,10 @@ public HiveWriter createWriter(Page partitionColumns, int position, OptionalInt // Write to: a new partition in a new partitioned table, // or a new unpartitioned table. updateMode = UpdateMode.NEW; - schema = new Properties(); - schema.setProperty(LIST_COLUMNS, dataColumns.stream() + schema.put(LIST_COLUMNS, dataColumns.stream() .map(DataColumn::getName) .collect(joining(","))); - schema.setProperty(LIST_COLUMN_TYPES, dataColumns.stream() + schema.put(LIST_COLUMN_TYPES, dataColumns.stream() .map(DataColumn::getHiveType) .map(HiveType::getHiveTypeName) .map(HiveTypeName::toString) @@ -402,7 +345,7 @@ public HiveWriter createWriter(Page partitionColumns, int position, OptionalInt } } - schema = getHiveSchema(table); + schema.putAll(getHiveSchema(table)); } if (partitionName.isPresent()) { @@ -448,7 +391,7 @@ public HiveWriter createWriter(Page partitionColumns, int position, OptionalInt outputStorageFormat = partition.get().getStorage().getStorageFormat(); compressionCodec = selectCompressionCodec(session, outputStorageFormat); - schema = getHiveSchema(partition.get(), table); + schema.putAll(getHiveSchema(partition.get(), table)); writeInfo = locationService.getPartitionWriteInfo(locationHandle, partition, partitionName.get()); break; @@ -462,7 +405,7 @@ public HiveWriter createWriter(Page partitionColumns, int position, OptionalInt outputStorageFormat = fromHiveStorageFormat(partitionStorageFormat); compressionCodec = selectCompressionCodec(session, partitionStorageFormat); - schema = getHiveSchema(table); + schema.putAll(getHiveSchema(table)); writeInfo = locationService.getPartitionWriteInfo(locationHandle, Optional.empty(), partitionName.get()); break; @@ -473,10 +416,7 @@ public HiveWriter createWriter(Page partitionColumns, int position, OptionalInt } } - JobConf outputConf = new JobConf(conf); - configureCompression(outputConf, compressionCodec); - - additionalTableParameters.forEach(schema::setProperty); + schema.putAll(additionalTableParameters); validateSchema(partitionName, schema); @@ -489,7 +429,7 @@ public HiveWriter createWriter(Page partitionColumns, int position, OptionalInt path = path.appendPath(subdir).appendPath(nameFormat.formatted(bucketToUse)); } else { - path = path.appendPath(computeFileName(bucketNumber) + getFileExtension(outputConf, outputStorageFormat)); + path = path.appendPath(computeFileName(bucketNumber) + getFileExtension(compressionCodec, outputStorageFormat)); } boolean useAcidSchema = isCreateTransactionalTable || (table != null && isFullAcidTable(table.getParameters())); @@ -538,50 +478,9 @@ public HiveWriter createWriter(Page partitionColumns, int position, OptionalInt } if (hiveFileWriter == null) { - hiveFileWriter = new RecordFileWriter( - new Path(path.toString()), - dataColumns.stream() - .map(DataColumn::getName) - .collect(toList()), - outputStorageFormat, - schema, - partitionStorageFormat.getEstimatedWriterMemoryUsage(), - outputConf, - typeManager, - parquetTimeZone, - session); + throw new TrinoException(HIVE_UNSUPPORTED_FORMAT, "Writing not supported for " + outputStorageFormat); } - String writePath = path.toString(); - String writerImplementation = hiveFileWriter.getClass().getName(); - - Consumer onCommit = hiveWriter -> { - Optional size; - try { - size = Optional.of(hiveWriter.getWrittenBytes()); - } - catch (RuntimeException e) { - // Do not fail the query if file system is not available - size = Optional.empty(); - } - - eventClient.post(new WriteCompletedEvent( - session.getQueryId(), - writePath, - schemaName, - tableName, - partitionName.orElse(null), - outputStorageFormat.getOutputFormat(), - writerImplementation, - nodeManager.getCurrentNode().getVersion(), - nodeManager.getCurrentNode().getHost(), - session.getIdentity().getPrincipal().map(Principal::getName).orElse(null), - nodeManager.getEnvironment(), - sessionProperties, - size.orElse(null), - hiveWriter.getRowCount())); - }; - if (!sortedBy.isEmpty()) { Location tempFilePath; if (sortedWritingTempStagingPathEnabled) { @@ -634,7 +533,6 @@ public HiveWriter createWriter(Page partitionColumns, int position, OptionalInt path.fileName(), writeInfo.writePath().toString(), writeInfo.targetPath().toString(), - onCommit, hiveWriterStats); } @@ -668,7 +566,7 @@ public SortingFileWriter makeRowIdSortingWriter(FileWriter deleteFileWriter, Loc OrcFileWriterFactory::createOrcDataSink); } - private void validateSchema(Optional partitionName, Properties schema) + private void validateSchema(Optional partitionName, Map schema) { // existing tables may have columns in a different order List fileColumnNames = getColumnNames(schema); @@ -782,28 +680,13 @@ public static int getBucketFromFileName(String fileName) return Integer.parseInt(matcher.group(1)); } - public static String getFileExtension(JobConf conf, StorageFormat storageFormat) + public static String getFileExtension(HiveCompressionCodec compression, StorageFormat format) { // text format files must have the correct extension when compressed - if (!HiveConf.getBoolVar(conf, COMPRESSRESULT) || !HIVE_IGNORE_KEY_OUTPUT_FORMAT_CLASS.equals(storageFormat.getOutputFormat())) { - return ""; - } - - String compressionCodecClass = conf.get("mapred.output.compression.codec"); - if (compressionCodecClass == null) { - return new DefaultCodec().getDefaultExtension(); - } - - try { - Class codecClass = conf.getClassByName(compressionCodecClass).asSubclass(CompressionCodec.class); - return ReflectionUtils.newInstance(codecClass, conf).getDefaultExtension(); - } - catch (ClassNotFoundException e) { - throw new TrinoException(HIVE_UNSUPPORTED_FORMAT, "Compression codec not found: " + compressionCodecClass, e); - } - catch (RuntimeException e) { - throw new TrinoException(HIVE_UNSUPPORTED_FORMAT, "Failed to load compression codec: " + compressionCodecClass, e); - } + return compression.getHiveCompressionKind() + .filter(ignore -> format.getOutputFormat().equals(HIVE_IGNORE_KEY_OUTPUT_FORMAT_CLASS)) + .map(CompressionKind::getFileExtension) + .orElse(""); } @VisibleForTesting diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/InternalHiveConnectorFactory.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/InternalHiveConnectorFactory.java index 7d231385477b..a3aef1ca0ae3 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/InternalHiveConnectorFactory.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/InternalHiveConnectorFactory.java @@ -21,19 +21,11 @@ import com.google.inject.TypeLiteral; import io.airlift.bootstrap.Bootstrap; import io.airlift.bootstrap.LifeCycleManager; -import io.airlift.event.client.EventModule; import io.airlift.json.JsonModule; import io.opentelemetry.api.OpenTelemetry; import io.opentelemetry.api.trace.Tracer; import io.trino.filesystem.TrinoFileSystemFactory; import io.trino.filesystem.manager.FileSystemModule; -import io.trino.hdfs.HdfsModule; -import io.trino.hdfs.authentication.HdfsAuthenticationModule; -import io.trino.hdfs.azure.HiveAzureModule; -import io.trino.hdfs.cos.HiveCosModule; -import io.trino.hdfs.gcs.HiveGcsModule; -import io.trino.hdfs.rubix.RubixEnabledConfig; -import io.trino.hdfs.rubix.RubixModule; import io.trino.plugin.base.CatalogName; import io.trino.plugin.base.CatalogNameModule; import io.trino.plugin.base.TypeDeserializerModule; @@ -41,14 +33,10 @@ import io.trino.plugin.base.classloader.ClassLoaderSafeConnectorPageSinkProvider; import io.trino.plugin.base.classloader.ClassLoaderSafeConnectorPageSourceProvider; import io.trino.plugin.base.classloader.ClassLoaderSafeConnectorSplitManager; -import io.trino.plugin.base.classloader.ClassLoaderSafeEventListener; import io.trino.plugin.base.classloader.ClassLoaderSafeNodePartitioningProvider; import io.trino.plugin.base.jmx.ConnectorObjectNameGeneratorModule; import io.trino.plugin.base.jmx.MBeanServerModule; import io.trino.plugin.base.session.SessionPropertiesProvider; -import io.trino.plugin.hive.aws.athena.PartitionProjectionModule; -import io.trino.plugin.hive.fs.CachingDirectoryListerModule; -import io.trino.plugin.hive.fs.DirectoryLister; import io.trino.plugin.hive.metastore.HiveMetastore; import io.trino.plugin.hive.metastore.HiveMetastoreModule; import io.trino.plugin.hive.procedure.HiveProcedureModule; @@ -68,7 +56,6 @@ import io.trino.spi.connector.ConnectorSplitManager; import io.trino.spi.connector.MetadataProvider; import io.trino.spi.connector.TableProcedureMetadata; -import io.trino.spi.eventlistener.EventListener; import io.trino.spi.procedure.Procedure; import org.weakref.jmx.guice.MBeanModule; @@ -76,9 +63,7 @@ import java.util.Optional; import java.util.Set; -import static com.google.common.collect.ImmutableSet.toImmutableSet; import static com.google.inject.multibindings.Multibinder.newSetBinder; -import static io.airlift.configuration.ConditionalModule.conditionalModule; import static java.util.Objects.requireNonNull; public final class InternalHiveConnectorFactory @@ -87,7 +72,7 @@ private InternalHiveConnectorFactory() {} public static Connector createConnector(String catalogName, Map config, ConnectorContext context, Module module) { - return createConnector(catalogName, config, context, module, Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty()); + return createConnector(catalogName, config, context, module, Optional.empty(), Optional.empty()); } public static Connector createConnector( @@ -96,9 +81,7 @@ public static Connector createConnector( ConnectorContext context, Module module, Optional metastore, - Optional fileSystemFactory, - Optional openTelemetry, - Optional directoryLister) + Optional fileSystemFactory) { requireNonNull(config, "config is null"); @@ -106,29 +89,20 @@ public static Connector createConnector( try (ThreadContextClassLoader ignored = new ThreadContextClassLoader(classLoader)) { Bootstrap app = new Bootstrap( new CatalogNameModule(catalogName), - new EventModule(), new MBeanModule(), new ConnectorObjectNameGeneratorModule("io.trino.plugin.hive", "trino.plugin.hive"), new JsonModule(), new TypeDeserializerModule(context.getTypeManager()), new HiveModule(), - new PartitionProjectionModule(), - new CachingDirectoryListerModule(directoryLister), - new HdfsModule(), - new HiveGcsModule(), - new HiveAzureModule(), - new HiveCosModule(), - conditionalModule(RubixEnabledConfig.class, RubixEnabledConfig::isCacheEnabled, new RubixModule()), new HiveMetastoreModule(metastore), new HiveSecurityModule(), - new HdfsAuthenticationModule(), fileSystemFactory .map(factory -> (Module) binder -> binder.bind(TrinoFileSystemFactory.class).toInstance(factory)) - .orElseGet(FileSystemModule::new), + .orElseGet(() -> new FileSystemModule(catalogName, context.getNodeManager(), context.getOpenTelemetry(), false)), new HiveProcedureModule(), new MBeanServerModule(), binder -> { - binder.bind(OpenTelemetry.class).toInstance(openTelemetry.orElse(context.getOpenTelemetry())); + binder.bind(OpenTelemetry.class).toInstance(context.getOpenTelemetry()); binder.bind(Tracer.class).toInstance(context.getTracer()); binder.bind(NodeVersion.class).toInstance(new NodeVersion(context.getNodeManager().getCurrentNode().getVersion())); binder.bind(NodeManager.class).toInstance(context.getNodeManager()); @@ -138,7 +112,6 @@ public static Connector createConnector( binder.bind(PageSorter.class).toInstance(context.getPageSorter()); binder.bind(CatalogName.class).toInstance(new CatalogName(catalogName)); }, - binder -> newSetBinder(binder, EventListener.class), binder -> bindSessionPropertiesProvider(binder, HiveSessionProperties.class), module); @@ -157,13 +130,8 @@ public static Connector createConnector( HiveTableProperties hiveTableProperties = injector.getInstance(HiveTableProperties.class); HiveColumnProperties hiveColumnProperties = injector.getInstance(HiveColumnProperties.class); HiveAnalyzeProperties hiveAnalyzeProperties = injector.getInstance(HiveAnalyzeProperties.class); - HiveMaterializedViewPropertiesProvider hiveMaterializedViewPropertiesProvider = injector.getInstance(HiveMaterializedViewPropertiesProvider.class); Set procedures = injector.getInstance(Key.get(new TypeLiteral>() {})); Set tableProcedures = injector.getInstance(Key.get(new TypeLiteral>() {})); - Set eventListeners = injector.getInstance(Key.get(new TypeLiteral>() {})) - .stream() - .map(listener -> new ClassLoaderSafeEventListener(listener, classLoader)) - .collect(toImmutableSet()); Set systemTableProviders = injector.getInstance(Key.get(new TypeLiteral>() {})); Optional hiveAccessControl = injector.getInstance(Key.get(new TypeLiteral>() {})) .map(accessControl -> new SystemTableAwareAccessControl(accessControl, systemTableProviders)) @@ -179,13 +147,11 @@ public static Connector createConnector( new ClassLoaderSafeNodePartitioningProvider(connectorDistributionProvider, classLoader), procedures, tableProcedures, - eventListeners, sessionPropertiesProviders, HiveSchemaProperties.SCHEMA_PROPERTIES, hiveTableProperties.getTableProperties(), hiveColumnProperties.getColumnProperties(), hiveAnalyzeProperties.getAnalyzeProperties(), - hiveMaterializedViewPropertiesProvider.getMaterializedViewProperties(), hiveAccessControl, injector.getInstance(HiveConfig.class).isSingleStatementWritesOnly(), classLoader); diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/InternalHiveSplit.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/InternalHiveSplit.java index 40c7e7a52b6e..2e64e327d2d7 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/InternalHiveSplit.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/InternalHiveSplit.java @@ -14,12 +14,14 @@ package io.trino.plugin.hive; import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; import io.trino.annotation.NotThreadSafe; import io.trino.plugin.hive.HiveSplit.BucketConversion; import io.trino.plugin.hive.HiveSplit.BucketValidation; import io.trino.spi.HostAddress; import java.util.List; +import java.util.Map; import java.util.Optional; import java.util.OptionalInt; import java.util.Properties; @@ -38,12 +40,13 @@ public class InternalHiveSplit { private static final int INSTANCE_SIZE = instanceSize(InternalHiveSplit.class) + instanceSize(Properties.class) + instanceSize(OptionalInt.class); + private static final int INTEGER_INSTANCE_SIZE = instanceSize(Integer.class); private final String path; private final long end; private final long estimatedFileSize; private final long fileModifiedTime; - private final Properties schema; + private final Schema schema; private final List partitionKeys; private final List blocks; private final String partitionName; @@ -51,10 +54,9 @@ public class InternalHiveSplit private final OptionalInt tableBucketNumber; private final boolean splittable; private final boolean forceLocalScheduling; - private final TableToPartitionMapping tableToPartitionMapping; + private final Map hiveColumnCoercions; private final Optional bucketConversion; private final Optional bucketValidation; - private final boolean s3SelectPushdownEnabled; private final Optional acidInfo; private final BooleanSupplier partitionMatchSupplier; @@ -68,17 +70,16 @@ public InternalHiveSplit( long end, long estimatedFileSize, long fileModifiedTime, - Properties schema, + Schema schema, List partitionKeys, List blocks, OptionalInt readBucketNumber, OptionalInt tableBucketNumber, boolean splittable, boolean forceLocalScheduling, - TableToPartitionMapping tableToPartitionMapping, + Map hiveColumnCoercions, Optional bucketConversion, Optional bucketValidation, - boolean s3SelectPushdownEnabled, Optional acidInfo, BooleanSupplier partitionMatchSupplier) { @@ -92,7 +93,7 @@ public InternalHiveSplit( requireNonNull(blocks, "blocks is null"); requireNonNull(readBucketNumber, "readBucketNumber is null"); requireNonNull(tableBucketNumber, "tableBucketNumber is null"); - requireNonNull(tableToPartitionMapping, "tableToPartitionMapping is null"); + requireNonNull(hiveColumnCoercions, "hiveColumnCoercions is null"); requireNonNull(bucketConversion, "bucketConversion is null"); requireNonNull(bucketValidation, "bucketValidation is null"); requireNonNull(acidInfo, "acidInfo is null"); @@ -111,10 +112,9 @@ public InternalHiveSplit( this.tableBucketNumber = tableBucketNumber; this.splittable = splittable; this.forceLocalScheduling = forceLocalScheduling; - this.tableToPartitionMapping = tableToPartitionMapping; + this.hiveColumnCoercions = ImmutableMap.copyOf(hiveColumnCoercions); this.bucketConversion = bucketConversion; this.bucketValidation = bucketValidation; - this.s3SelectPushdownEnabled = s3SelectPushdownEnabled; this.acidInfo = acidInfo; this.partitionMatchSupplier = partitionMatchSupplier; } @@ -144,12 +144,7 @@ public long getFileModifiedTime() return fileModifiedTime; } - public boolean isS3SelectPushdownEnabled() - { - return s3SelectPushdownEnabled; - } - - public Properties getSchema() + public Schema getSchema() { return schema; } @@ -184,9 +179,9 @@ public boolean isForceLocalScheduling() return forceLocalScheduling; } - public TableToPartitionMapping getTableToPartitionMapping() + public Map getHiveColumnCoercions() { - return tableToPartitionMapping; + return hiveColumnCoercions; } public Optional getBucketConversion() @@ -229,7 +224,7 @@ public int getEstimatedSizeInBytes() estimatedSizeOf(partitionKeys, HivePartitionKey::getEstimatedSizeInBytes) + estimatedSizeOf(blocks, InternalHiveBlock::getEstimatedSizeInBytes) + estimatedSizeOf(partitionName) + - tableToPartitionMapping.getEstimatedSizeInBytes(); + estimatedSizeOf(hiveColumnCoercions, (Integer key) -> INTEGER_INSTANCE_SIZE, HiveTypeName::getEstimatedSizeInBytes); return toIntExact(result); } diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/LegacyHiveViewReader.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/LegacyHiveViewReader.java index 6da51a5c8729..fe61c48afe26 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/LegacyHiveViewReader.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/LegacyHiveViewReader.java @@ -24,8 +24,8 @@ import static com.google.common.collect.ImmutableList.toImmutableList; import static io.trino.plugin.hive.HiveErrorCode.HIVE_INVALID_METADATA; -import static io.trino.plugin.hive.HiveMetadata.TABLE_COMMENT; import static io.trino.plugin.hive.HiveToTrinoTranslator.translateHiveViewToTrino; +import static io.trino.plugin.hive.metastore.Table.TABLE_COMMENT; public class LegacyHiveViewReader implements ViewReaderUtil.ViewReader diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/MergeFileWriter.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/MergeFileWriter.java index a63435535669..11443f39e0d7 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/MergeFileWriter.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/MergeFileWriter.java @@ -32,9 +32,9 @@ import java.io.Closeable; import java.io.IOException; import java.util.List; +import java.util.Map; import java.util.Optional; import java.util.OptionalInt; -import java.util.Properties; import java.util.regex.Matcher; import java.util.regex.Pattern; @@ -82,7 +82,7 @@ public class MergeFileWriter private final RowIdSortingFileWriterMaker sortingFileWriterMaker; private final OrcFileWriterFactory orcFileWriterFactory; private final HiveCompressionCodec compressionCodec; - private final Properties hiveAcidSchema; + private final Map hiveAcidSchema; private final String bucketFilename; private Optional deleteFileWriter = Optional.empty(); private Optional insertFileWriter = Optional.empty(); @@ -245,14 +245,12 @@ private Page buildDeletePage(Block rowIds, long writeId) private FileWriter getOrCreateInsertFileWriter() { if (insertFileWriter.isEmpty()) { - Properties schemaCopy = new Properties(); - schemaCopy.putAll(hiveAcidSchema); insertFileWriter = orcFileWriterFactory.createFileWriter( deltaDirectory.appendPath(bucketFilename), ACID_COLUMN_NAMES, fromHiveStorageFormat(ORC), compressionCodec, - schemaCopy, + hiveAcidSchema, session, bucketNumber, transaction, @@ -265,15 +263,13 @@ private FileWriter getOrCreateInsertFileWriter() private FileWriter getOrCreateDeleteFileWriter() { if (deleteFileWriter.isEmpty()) { - Properties schemaCopy = new Properties(); - schemaCopy.putAll(hiveAcidSchema); Location deletePath = deleteDeltaDirectory.appendPath(bucketFilename); FileWriter writer = getWriter(orcFileWriterFactory.createFileWriter( deletePath, ACID_COLUMN_NAMES, fromHiveStorageFormat(ORC), compressionCodec, - schemaCopy, + hiveAcidSchema, session, bucketNumber, transaction, diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/NoneHiveMaterializedViewMetadata.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/NoneHiveMaterializedViewMetadata.java deleted file mode 100644 index a23d562b616d..000000000000 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/NoneHiveMaterializedViewMetadata.java +++ /dev/null @@ -1,100 +0,0 @@ -/* - * 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.hive; - -import com.google.common.collect.ImmutableList; -import com.google.common.collect.ImmutableMap; -import io.trino.spi.TrinoException; -import io.trino.spi.connector.ConnectorMaterializedViewDefinition; -import io.trino.spi.connector.ConnectorSession; -import io.trino.spi.connector.MaterializedViewFreshness; -import io.trino.spi.connector.SchemaTableName; - -import java.util.List; -import java.util.Map; -import java.util.Optional; -import java.util.concurrent.CompletableFuture; - -import static io.trino.spi.StandardErrorCode.NOT_FOUND; -import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED; - -public class NoneHiveMaterializedViewMetadata - implements HiveMaterializedViewMetadata -{ - @Override - public void createMaterializedView(ConnectorSession session, SchemaTableName viewName, ConnectorMaterializedViewDefinition definition, boolean replace, boolean ignoreExisting) - { - throw new TrinoException(NOT_SUPPORTED, "This connector does not support creating materialized views"); - } - - @Override - public void dropMaterializedView(ConnectorSession session, SchemaTableName viewName) - { - throw new TrinoException(NOT_SUPPORTED, "This connector does not support dropping materialized views"); - } - - @Override - public List listMaterializedViews(ConnectorSession session, Optional schemaName) - { - return ImmutableList.of(); - } - - @Override - public Map getMaterializedViews(ConnectorSession session, Optional schemaName) - { - return ImmutableMap.of(); - } - - @Override - public Optional getMaterializedView(ConnectorSession session, SchemaTableName viewName) - { - return Optional.empty(); - } - - @Override - public MaterializedViewFreshness getMaterializedViewFreshness(ConnectorSession session, SchemaTableName name) - { - throw new TrinoException(NOT_FOUND, "This connector does not support materialized views"); - } - - @Override - public boolean delegateMaterializedViewRefreshToConnector(ConnectorSession session, SchemaTableName viewName) - { - throw new TrinoException(NOT_SUPPORTED, "This connector does not support materialized views"); - } - - @Override - public CompletableFuture refreshMaterializedView(ConnectorSession session, SchemaTableName viewName) - { - throw new TrinoException(NOT_SUPPORTED, "This connector does not support materialized views"); - } - - @Override - public void renameMaterializedView(ConnectorSession session, SchemaTableName existingViewName, SchemaTableName newViewName) - { - throw new TrinoException(NOT_SUPPORTED, "This connector does not support renaming materialized views"); - } - - @Override - public void setMaterializedViewProperties(ConnectorSession session, SchemaTableName viewName, Map> properties) - { - throw new TrinoException(NOT_SUPPORTED, "This connector does not support setting materialized view properties"); - } - - @Override - public void setMaterializedViewColumnComment(ConnectorSession session, SchemaTableName viewName, String columnName, Optional comment) - { - throw new TrinoException(NOT_SUPPORTED, "This connector does not support setting materialized view column comment"); - } -} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/PartitionStatistics.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/PartitionStatistics.java index 3537517abe6a..d8571c3a5eb1 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/PartitionStatistics.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/PartitionStatistics.java @@ -60,6 +60,11 @@ public Map getColumnStatistics() return columnStatistics; } + public PartitionStatistics withBasicStatistics(HiveBasicStatistics basicStatistics) + { + return new PartitionStatistics(basicStatistics, columnStatistics); + } + @Override public boolean equals(Object o) { diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/PartitionsSystemTableProvider.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/PartitionsSystemTableProvider.java index c472e5010fce..9bfb86ad5ef8 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/PartitionsSystemTableProvider.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/PartitionsSystemTableProvider.java @@ -34,7 +34,7 @@ import static io.trino.plugin.hive.SystemTableHandler.PARTITIONS; import static io.trino.plugin.hive.metastore.MetastoreUtil.getProtectMode; import static io.trino.plugin.hive.metastore.MetastoreUtil.verifyOnline; -import static io.trino.plugin.hive.util.HiveBucketing.getHiveBucketHandle; +import static io.trino.plugin.hive.util.HiveBucketing.getHiveTablePartitioningForRead; import static io.trino.plugin.hive.util.HiveUtil.getPartitionKeyColumnHandles; import static io.trino.plugin.hive.util.HiveUtil.getRegularColumnHandles; import static io.trino.plugin.hive.util.HiveUtil.isDeltaLakeTable; @@ -87,7 +87,7 @@ public Optional getSystemTable(HiveMetadata metadata, ConnectorSess sourceTable.getParameters(), getPartitionKeyColumnHandles(sourceTable, typeManager), getRegularColumnHandles(sourceTable, typeManager, getTimestampPrecision(session)), - getHiveBucketHandle(session, sourceTable, typeManager)); + getHiveTablePartitioningForRead(session, sourceTable, typeManager)); List partitionColumns = sourceTableHandle.getPartitionColumns(); if (partitionColumns.isEmpty()) { diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/RcFileFileWriterFactory.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/RcFileFileWriterFactory.java index 98717d361a0c..7db3d6b9b960 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/RcFileFileWriterFactory.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/RcFileFileWriterFactory.java @@ -14,7 +14,6 @@ package io.trino.plugin.hive; import com.google.common.collect.ImmutableMap; -import com.google.common.collect.Maps; import com.google.inject.Inject; import io.trino.filesystem.Location; import io.trino.filesystem.TrinoFileSystem; @@ -36,15 +35,15 @@ import java.io.Closeable; import java.io.OutputStream; import java.util.List; +import java.util.Map; import java.util.Optional; import java.util.OptionalInt; -import java.util.Properties; import java.util.function.Supplier; import static io.trino.memory.context.AggregatedMemoryContext.newSimpleAggregatedMemoryContext; import static io.trino.plugin.hive.HiveErrorCode.HIVE_WRITER_OPEN_ERROR; -import static io.trino.plugin.hive.HiveMetadata.PRESTO_QUERY_ID_NAME; -import static io.trino.plugin.hive.HiveMetadata.PRESTO_VERSION_NAME; +import static io.trino.plugin.hive.HiveMetadata.TRINO_QUERY_ID_NAME; +import static io.trino.plugin.hive.HiveMetadata.TRINO_VERSION_NAME; import static io.trino.plugin.hive.HiveSessionProperties.getTimestampPrecision; import static io.trino.plugin.hive.HiveSessionProperties.isRcfileOptimizedWriterValidate; import static io.trino.plugin.hive.util.HiveClassNames.COLUMNAR_SERDE_CLASS; @@ -91,7 +90,7 @@ public Optional createFileWriter( List inputColumnNames, StorageFormat storageFormat, HiveCompressionCodec compressionCodec, - Properties schema, + Map schema, ConnectorSession session, OptionalInt bucketNumber, AcidTransaction transaction, @@ -107,7 +106,7 @@ public Optional createFileWriter( columnEncodingFactory = new BinaryColumnEncodingFactory(timeZone); } else if (COLUMNAR_SERDE_CLASS.equals(storageFormat.getSerde())) { - columnEncodingFactory = new TextColumnEncodingFactory(TextEncodingOptions.fromSchema(Maps.fromProperties(schema))); + columnEncodingFactory = new TextColumnEncodingFactory(TextEncodingOptions.fromSchema(schema)); } else { return Optional.empty(); @@ -145,8 +144,8 @@ else if (COLUMNAR_SERDE_CLASS.equals(storageFormat.getSerde())) { compressionCodec.getHiveCompressionKind(), fileInputColumnIndexes, ImmutableMap.builder() - .put(PRESTO_VERSION_NAME, nodeVersion.toString()) - .put(PRESTO_QUERY_ID_NAME, session.getQueryId()) + .put(TRINO_VERSION_NAME, nodeVersion.toString()) + .put(TRINO_QUERY_ID_NAME, session.getQueryId()) .buildOrThrow(), validationInputFactory)); } diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/RecordFileWriter.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/RecordFileWriter.java index 7874b0326b37..6f06fc47ee77 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/RecordFileWriter.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/RecordFileWriter.java @@ -39,6 +39,7 @@ import java.io.IOException; import java.io.UncheckedIOException; import java.util.List; +import java.util.Map; import java.util.Optional; import java.util.Properties; @@ -81,7 +82,7 @@ public RecordFileWriter( Path path, List inputColumnNames, StorageFormat storageFormat, - Properties schema, + Map schema, DataSize estimatedWriterMemoryUsage, JobConf conf, TypeManager typeManager, @@ -100,17 +101,19 @@ public RecordFileWriter( fieldCount = fileColumnNames.size(); String serde = storageFormat.getSerde(); - serializer = initializeSerializer(conf, schema, serde); + Properties properties = new Properties(); + properties.putAll(schema); + serializer = initializeSerializer(conf, properties, serde); List objectInspectors = getRowColumnInspectors(fileColumnTypes); tableInspector = getStandardStructObjectInspector(fileColumnNames, objectInspectors); if (storageFormat.getOutputFormat().equals(HIVE_IGNORE_KEY_OUTPUT_FORMAT_CLASS)) { Optional textHeaderWriter = Optional.of(new TextHeaderWriter(serializer, typeManager, session, fileColumnNames)); - recordWriter = createRecordWriter(path, conf, schema, storageFormat.getOutputFormat(), session, textHeaderWriter); + recordWriter = createRecordWriter(path, conf, properties, storageFormat.getOutputFormat(), session, textHeaderWriter); } else { - recordWriter = createRecordWriter(path, conf, schema, storageFormat.getOutputFormat(), session, Optional.empty()); + recordWriter = createRecordWriter(path, conf, properties, storageFormat.getOutputFormat(), session, Optional.empty()); } // reorder (and possibly reduce) struct fields to match input diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/S3StorageClassFilter.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/S3StorageClassFilter.java new file mode 100644 index 000000000000..c31536481291 --- /dev/null +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/S3StorageClassFilter.java @@ -0,0 +1,59 @@ +/* + * 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.hive; + +import io.trino.filesystem.FileEntry; + +import java.util.function.Predicate; + +import static com.google.common.base.Predicates.alwaysTrue; + +public enum S3StorageClassFilter { + READ_ALL, + READ_NON_GLACIER, + READ_NON_GLACIER_AND_RESTORED; + + private static final String S3_GLACIER_TAG = "s3:glacier"; + private static final String S3_GLACIER_AND_RESTORED_TAG = "s3:glacierRestored"; + + /** + * Checks if the S3 object is not an object with a storage class of glacier/deep_archive + * + * @return boolean that helps identify if FileEntry object contains tags for glacier object + */ + private static boolean isNotGlacierObject(FileEntry fileEntry) + { + return !fileEntry.tags().contains(S3_GLACIER_TAG); + } + + /** + * Only restored objects will have the restoreExpiryDate set. + * Ignore not-restored objects and in-progress restores. + * + * @return boolean that helps identify if FileEntry object contains tags for glacier or glacierRestored object + */ + private static boolean isCompletedRestoredObject(FileEntry fileEntry) + { + return isNotGlacierObject(fileEntry) || fileEntry.tags().contains(S3_GLACIER_AND_RESTORED_TAG); + } + + public Predicate toFileEntryPredicate() + { + return switch (this) { + case READ_ALL -> alwaysTrue(); + case READ_NON_GLACIER -> S3StorageClassFilter::isNotGlacierObject; + case READ_NON_GLACIER_AND_RESTORED -> S3StorageClassFilter::isCompletedRestoredObject; + }; + } +} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/Schema.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/Schema.java new file mode 100644 index 000000000000..9ed591927af5 --- /dev/null +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/Schema.java @@ -0,0 +1,43 @@ +/* + * 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.hive; + +import io.airlift.slice.SizeOf; + +import java.util.Map; + +import static io.airlift.slice.SizeOf.estimatedSizeOf; +import static io.airlift.slice.SizeOf.instanceSize; +import static java.util.Objects.requireNonNull; + +public record Schema( + String serializationLibraryName, + boolean isFullAcidTable, + Map serdeProperties) +{ + private static final int INSTANCE_SIZE = instanceSize(Schema.class); + + public Schema + { + requireNonNull(serializationLibraryName, "serializationLibraryName is null"); + requireNonNull(serdeProperties, "serdeProperties is null"); + } + + public long getRetainedSizeInBytes() + { + return INSTANCE_SIZE + + estimatedSizeOf(serializationLibraryName) + + estimatedSizeOf(serdeProperties, SizeOf::estimatedSizeOf, SizeOf::estimatedSizeOf); + } +} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/SchemaAlreadyExistsException.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/SchemaAlreadyExistsException.java index 5511366ddecc..21eebb05e3fc 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/SchemaAlreadyExistsException.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/SchemaAlreadyExistsException.java @@ -25,12 +25,22 @@ public class SchemaAlreadyExistsException public SchemaAlreadyExistsException(String schemaName) { - this(schemaName, format("Schema already exists: '%s'", schemaName)); + this(schemaName, (Throwable) null); } public SchemaAlreadyExistsException(String schemaName, String message) { - super(ALREADY_EXISTS, message); + this(schemaName, message, null); + } + + public SchemaAlreadyExistsException(String schemaName, Throwable cause) + { + this(schemaName, format("Schema already exists: '%s'", schemaName), cause); + } + + public SchemaAlreadyExistsException(String schemaName, String message, Throwable cause) + { + super(ALREADY_EXISTS, message, cause); this.schemaName = schemaName; } diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/TableAlreadyExistsException.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/TableAlreadyExistsException.java index 157e0744f8c1..0433d7968361 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/TableAlreadyExistsException.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/TableAlreadyExistsException.java @@ -26,7 +26,7 @@ public class TableAlreadyExistsException public TableAlreadyExistsException(SchemaTableName tableName) { - this(tableName, format("Table already exists: '%s'", tableName)); + this(tableName, (Throwable) null); } public TableAlreadyExistsException(SchemaTableName tableName, String message) @@ -34,6 +34,11 @@ public TableAlreadyExistsException(SchemaTableName tableName, String message) this(tableName, message, null); } + public TableAlreadyExistsException(SchemaTableName tableName, Throwable cause) + { + this(tableName, format("Table already exists: '%s'", tableName), cause); + } + public TableAlreadyExistsException(SchemaTableName tableName, String message, Throwable cause) { super(ALREADY_EXISTS, message, cause); diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/TableInvalidationCallback.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/TableInvalidationCallback.java index d444eb45eab5..ad8ae443bbd4 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/TableInvalidationCallback.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/TableInvalidationCallback.java @@ -13,12 +13,22 @@ */ package io.trino.plugin.hive; +import io.trino.filesystem.Location; import io.trino.plugin.hive.metastore.Partition; import io.trino.plugin.hive.metastore.Table; public interface TableInvalidationCallback { - void invalidate(Partition partition); + default boolean isCached(Location location) + { + return false; + } - void invalidate(Table table); + default void invalidate(Location location) {} + + default void invalidate(Partition partition) {} + + default void invalidate(Table table) {} + + default void invalidateAll() {} } diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/TransformConnectorPageSource.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/TransformConnectorPageSource.java new file mode 100644 index 000000000000..cc10926ebc7e --- /dev/null +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/TransformConnectorPageSource.java @@ -0,0 +1,324 @@ +/* + * 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.hive; + +import com.google.common.collect.ImmutableList; +import com.google.errorprone.annotations.CanIgnoreReturnValue; +import com.google.errorprone.annotations.CheckReturnValue; +import io.trino.spi.Page; +import io.trino.spi.block.Block; +import io.trino.spi.block.RunLengthEncodedBlock; +import io.trino.spi.connector.ConnectorPageSource; +import io.trino.spi.metrics.Metrics; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.List; +import java.util.Optional; +import java.util.OptionalLong; +import java.util.concurrent.CompletableFuture; +import java.util.function.Function; + +import static com.google.common.base.Preconditions.checkArgument; +import static io.trino.plugin.base.util.Closables.closeAllSuppress; +import static io.trino.spi.block.RowBlock.getRowFieldsFromBlock; +import static java.util.Objects.requireNonNull; + +public final class TransformConnectorPageSource + implements ConnectorPageSource +{ + private final ConnectorPageSource connectorPageSource; + private final Function transform; + + @CheckReturnValue + public static TransformConnectorPageSource create(ConnectorPageSource connectorPageSource, Function transform) + { + return new TransformConnectorPageSource(connectorPageSource, transform); + } + + private TransformConnectorPageSource(ConnectorPageSource connectorPageSource, Function transform) + { + this.connectorPageSource = requireNonNull(connectorPageSource, "connectorPageSource is null"); + this.transform = requireNonNull(transform, "transform is null"); + } + + @Override + public long getCompletedBytes() + { + return connectorPageSource.getCompletedBytes(); + } + + @Override + public OptionalLong getCompletedPositions() + { + return connectorPageSource.getCompletedPositions(); + } + + @Override + public long getReadTimeNanos() + { + return connectorPageSource.getReadTimeNanos(); + } + + @Override + public boolean isFinished() + { + return connectorPageSource.isFinished(); + } + + @Override + public Page getNextPage() + { + try { + Page page = connectorPageSource.getNextPage(); + if (page == null) { + return null; + } + return transform.apply(page); + } + catch (Throwable e) { + closeAllSuppress(e, connectorPageSource); + throw e; + } + } + + @Override + public long getMemoryUsage() + { + return connectorPageSource.getMemoryUsage(); + } + + @Override + public void close() + throws IOException + { + connectorPageSource.close(); + } + + @Override + public CompletableFuture isBlocked() + { + return connectorPageSource.isBlocked(); + } + + @Override + public Metrics getMetrics() + { + return connectorPageSource.getMetrics(); + } + + @CheckReturnValue + public static Builder builder() + { + return new Builder(); + } + + public static final class Builder + { + private final List> transforms = new ArrayList<>(); + private boolean requiresTransform; + + private Builder() {} + + @CanIgnoreReturnValue + public Builder constantValue(Block constantValue) + { + requiresTransform = true; + transforms.add(new ConstantValue(constantValue)); + return this; + } + + @CanIgnoreReturnValue + public Builder column(int inputField) + { + return column(inputField, Optional.empty()); + } + + @CanIgnoreReturnValue + public Builder column(int inputField, Optional> transform) + { + if (transform.isPresent()) { + return transform(inputField, transform.get()); + } + + if (inputField != transforms.size()) { + requiresTransform = true; + } + transforms.add(new InputColumn(inputField)); + return this; + } + + @CanIgnoreReturnValue + public Builder dereferenceField(List path) + { + return dereferenceField(path, Optional.empty()); + } + + @CanIgnoreReturnValue + public Builder dereferenceField(List path, Optional> transform) + { + requireNonNull(path, "path is null"); + if (path.size() == 1) { + return column(path.get(0), transform); + } + + requiresTransform = true; + transforms.add(new DereferenceFieldTransform(path, transform)); + return this; + } + + @CanIgnoreReturnValue + public Builder transform(int inputColumn, Function transform) + { + requireNonNull(transform, "transform is null"); + requiresTransform = true; + transforms.add(new TransformBlock(transform, inputColumn)); + return this; + } + + @CanIgnoreReturnValue + public Builder transform(Function transform) + { + requiresTransform = true; + transforms.add(transform); + return this; + } + + @CheckReturnValue + public ConnectorPageSource build(ConnectorPageSource pageSource) + { + if (!requiresTransform) { + return pageSource; + } + + List> functions = List.copyOf(transforms); + return new TransformConnectorPageSource(pageSource, new TransformPages(functions)); + } + } + + private record ConstantValue(Block constantValue) + implements Function + { + @Override + public Block apply(Page page) + { + return RunLengthEncodedBlock.create(constantValue, page.getPositionCount()); + } + } + + private record InputColumn(int inputField) + implements Function + { + @Override + public Block apply(Page page) + { + return page.getBlock(inputField); + } + } + + private record DereferenceFieldTransform(List path, Optional> transform) + implements Function + { + private DereferenceFieldTransform + { + path = ImmutableList.copyOf(requireNonNull(path, "path is null")); + checkArgument(!path.isEmpty(), "path is empty"); + checkArgument(path.stream().allMatch(element -> element >= 0), "path element is negative"); + requireNonNull(transform, "transform is null"); + } + + @Override + public Block apply(Page page) + { + Block block = page.getBlock(path.get(0)); + for (int dereferenceIndex : path.subList(1, path.size())) { + block = getRowFieldsFromBlock(block).get(dereferenceIndex); + } + if (transform.isPresent()) { + block = transform.get().apply(block); + } + return block; + } + } + + private record TransformBlock(Function transform, int inputColumn) + implements Function + { + @Override + public Block apply(Page page) + { + return transform.apply(page.getBlock(inputColumn)); + } + } + + private record TransformPages(List> functions) + implements Function + { + private TransformPages + { + functions = List.copyOf(requireNonNull(functions, "functions is null")); + } + + @Override + public Page apply(Page page) + { + TransformPage transformPage = new TransformPage(page, functions); + return transformPage.getPage(); + } + } + + private record TransformPage(Page Page, List> transforms, Block[] blocks) + { + private TransformPage(Page page, List> transforms) + { + this(page, transforms, new Block[transforms.size()]); + } + + private TransformPage + { + requireNonNull(Page, "Page is null"); + transforms = List.copyOf(requireNonNull(transforms, "transforms is null")); + requireNonNull(blocks, "blocks is null"); + checkArgument(transforms.size() == blocks.length, "transforms and blocks size mismatch"); + } + + public int getPositionCount() + { + return Page.getPositionCount(); + } + + public int getChannelCount() + { + return blocks.length; + } + + public Block getBlock(int channel) + { + Block block = blocks[channel]; + if (block == null) { + block = transforms.get(channel).apply(Page); + blocks[channel] = block; + } + return block; + } + + public Page getPage() + { + for (int i = 0; i < blocks.length; i++) { + getBlock(i); + } + return new Page(getPositionCount(), blocks); + } + } +} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/TrinoViewHiveMetastore.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/TrinoViewHiveMetastore.java index f009267fdab4..216b37cc3b52 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/TrinoViewHiveMetastore.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/TrinoViewHiveMetastore.java @@ -18,6 +18,7 @@ import io.trino.plugin.hive.metastore.Column; import io.trino.plugin.hive.metastore.HiveMetastore; import io.trino.plugin.hive.metastore.PrincipalPrivileges; +import io.trino.plugin.hive.metastore.TableInfo; import io.trino.spi.TrinoException; import io.trino.spi.connector.ConnectorSession; import io.trino.spi.connector.ConnectorViewDefinition; @@ -32,9 +33,7 @@ import java.util.stream.Stream; import static com.google.common.collect.ImmutableList.toImmutableList; -import static io.trino.plugin.hive.HiveMetadata.PRESTO_VIEW_COMMENT; import static io.trino.plugin.hive.HiveMetadata.PRESTO_VIEW_EXPANDED_TEXT_MARKER; -import static io.trino.plugin.hive.HiveMetadata.TABLE_COMMENT; import static io.trino.plugin.hive.HiveType.HIVE_STRING; import static io.trino.plugin.hive.TableType.VIRTUAL_VIEW; import static io.trino.plugin.hive.TrinoViewUtil.createViewProperties; @@ -73,7 +72,7 @@ public void createView(ConnectorSession session, SchemaTableName schemaViewName, .setTableName(schemaViewName.getTableName()) .setOwner(isUsingSystemSecurity ? Optional.empty() : Optional.of(session.getUser())) .setTableType(VIRTUAL_VIEW.name()) - .setDataColumns(ImmutableList.of(new Column("dummy", HIVE_STRING, Optional.empty()))) + .setDataColumns(ImmutableList.of(new Column("dummy", HIVE_STRING, Optional.empty(), Map.of()))) .setPartitionColumns(ImmutableList.of()) .setParameters(createViewProperties(session, trinoVersion, connectorName)) .setViewOriginalText(Optional.of(encodeViewData(definition))) @@ -91,7 +90,7 @@ public void createView(ConnectorSession session, SchemaTableName schemaViewName, throw new ViewAlreadyExistsException(schemaViewName); } - metastore.replaceTable(schemaViewName.getSchemaName(), schemaViewName.getTableName(), table, principalPrivileges); + metastore.replaceTable(schemaViewName.getSchemaName(), schemaViewName.getTableName(), table, principalPrivileges, ImmutableMap.of()); return; } @@ -157,8 +156,9 @@ public Map getViews(Optional s private Stream listViews(String schema) { // Filter on PRESTO_VIEW_COMMENT to distinguish from materialized views - return metastore.getTablesWithParameter(schema, TABLE_COMMENT, PRESTO_VIEW_COMMENT).stream() - .map(table -> new SchemaTableName(schema, table)); + return metastore.getTables(schema).stream() + .filter(tableInfo -> tableInfo.extendedRelationType() == TableInfo.ExtendedRelationType.TRINO_VIEW) + .map(TableInfo::tableName); } public Optional getView(SchemaTableName viewName) @@ -168,7 +168,6 @@ public Optional getView(SchemaTableName viewName) } return metastore.getTable(viewName.getSchemaName(), viewName.getTableName()) .flatMap(view -> TrinoViewUtil.getView( - viewName, view.getViewOriginalText(), view.getTableType(), view.getParameters(), @@ -180,7 +179,7 @@ public void updateViewComment(ConnectorSession session, SchemaTableName viewName io.trino.plugin.hive.metastore.Table view = metastore.getTable(viewName.getSchemaName(), viewName.getTableName()) .orElseThrow(() -> new ViewNotFoundException(viewName)); - ConnectorViewDefinition definition = TrinoViewUtil.getView(viewName, view.getViewOriginalText(), view.getTableType(), view.getParameters(), view.getOwner()) + ConnectorViewDefinition definition = TrinoViewUtil.getView(view.getViewOriginalText(), view.getTableType(), view.getParameters(), view.getOwner()) .orElseThrow(() -> new ViewNotFoundException(viewName)); ConnectorViewDefinition newDefinition = new ConnectorViewDefinition( definition.getOriginalSql(), @@ -199,7 +198,7 @@ public void updateViewColumnComment(ConnectorSession session, SchemaTableName vi io.trino.plugin.hive.metastore.Table view = metastore.getTable(viewName.getSchemaName(), viewName.getTableName()) .orElseThrow(() -> new ViewNotFoundException(viewName)); - ConnectorViewDefinition definition = TrinoViewUtil.getView(viewName, view.getViewOriginalText(), view.getTableType(), view.getParameters(), view.getOwner()) + ConnectorViewDefinition definition = TrinoViewUtil.getView(view.getViewOriginalText(), view.getTableType(), view.getParameters(), view.getOwner()) .orElseThrow(() -> new ViewNotFoundException(viewName)); ConnectorViewDefinition newDefinition = new ConnectorViewDefinition( definition.getOriginalSql(), @@ -222,6 +221,6 @@ private void replaceView(ConnectorSession session, SchemaTableName viewName, io. PrincipalPrivileges principalPrivileges = isUsingSystemSecurity ? NO_PRIVILEGES : buildInitialPrivilegeSet(session.getUser()); - metastore.replaceTable(viewName.getSchemaName(), viewName.getTableName(), viewBuilder.build(), principalPrivileges); + metastore.replaceTable(viewName.getSchemaName(), viewName.getTableName(), viewBuilder.build(), principalPrivileges, ImmutableMap.of()); } } diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/TrinoViewUtil.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/TrinoViewUtil.java index 189b2caa354d..c0231ca401e8 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/TrinoViewUtil.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/TrinoViewUtil.java @@ -16,42 +16,34 @@ import com.google.common.collect.ImmutableMap; import io.trino.spi.connector.ConnectorSession; import io.trino.spi.connector.ConnectorViewDefinition; -import io.trino.spi.connector.SchemaTableName; import java.util.Map; import java.util.Optional; import static com.google.common.base.Preconditions.checkArgument; -import static io.trino.plugin.hive.HiveMetadata.PRESTO_QUERY_ID_NAME; -import static io.trino.plugin.hive.HiveMetadata.PRESTO_VERSION_NAME; -import static io.trino.plugin.hive.HiveMetadata.PRESTO_VIEW_COMMENT; -import static io.trino.plugin.hive.HiveMetadata.TABLE_COMMENT; import static io.trino.plugin.hive.HiveMetadata.TRINO_CREATED_BY; +import static io.trino.plugin.hive.HiveMetadata.TRINO_QUERY_ID_NAME; +import static io.trino.plugin.hive.HiveMetadata.TRINO_VERSION_NAME; import static io.trino.plugin.hive.ViewReaderUtil.PRESTO_VIEW_FLAG; -import static io.trino.plugin.hive.ViewReaderUtil.isHiveOrPrestoView; -import static io.trino.plugin.hive.ViewReaderUtil.isPrestoView; +import static io.trino.plugin.hive.ViewReaderUtil.isTrinoView; +import static io.trino.plugin.hive.metastore.Table.TABLE_COMMENT; +import static io.trino.plugin.hive.metastore.TableInfo.PRESTO_VIEW_COMMENT; public final class TrinoViewUtil { private TrinoViewUtil() {} public static Optional getView( - SchemaTableName viewName, Optional viewOriginalText, String tableType, Map tableParameters, Optional tableOwner) { - if (!isView(tableType, tableParameters)) { + if (!isTrinoView(tableType, tableParameters)) { // Filter out Tables and Materialized Views return Optional.empty(); } - if (!isPrestoView(tableParameters)) { - // Hive views are not compatible - throw new HiveViewNotSupportedException(viewName); - } - checkArgument(viewOriginalText.isPresent(), "viewOriginalText must be present"); ConnectorViewDefinition definition = ViewReaderUtil.PrestoViewReader.decodeViewData(viewOriginalText.get()); // use owner from table metadata if it exists @@ -68,18 +60,13 @@ public static Optional getView( return Optional.of(definition); } - private static boolean isView(String tableType, Map tableParameters) - { - return isHiveOrPrestoView(tableType) && PRESTO_VIEW_COMMENT.equals(tableParameters.get(TABLE_COMMENT)); - } - public static Map createViewProperties(ConnectorSession session, String trinoVersion, String connectorName) { return ImmutableMap.builder() .put(PRESTO_VIEW_FLAG, "true") .put(TRINO_CREATED_BY, connectorName) - .put(PRESTO_VERSION_NAME, trinoVersion) - .put(PRESTO_QUERY_ID_NAME, session.getQueryId()) + .put(TRINO_VERSION_NAME, trinoVersion) + .put(TRINO_QUERY_ID_NAME, session.getQueryId()) .put(TABLE_COMMENT, PRESTO_VIEW_COMMENT) .buildOrThrow(); } diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/ViewReaderUtil.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/ViewReaderUtil.java index 4eadea90eb02..f46bef1709c8 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/ViewReaderUtil.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/ViewReaderUtil.java @@ -49,13 +49,15 @@ import static com.linkedin.coral.trino.rel2trino.functions.TrinoKeywordsConverter.quoteWordIfNotQuoted; import static io.trino.plugin.hive.HiveErrorCode.HIVE_INVALID_VIEW_DATA; import static io.trino.plugin.hive.HiveErrorCode.HIVE_VIEW_TRANSLATION_ERROR; -import static io.trino.plugin.hive.HiveMetadata.TABLE_COMMENT; import static io.trino.plugin.hive.HiveSessionProperties.isHiveViewsLegacyTranslation; import static io.trino.plugin.hive.HiveStorageFormat.TEXTFILE; import static io.trino.plugin.hive.HiveType.toHiveType; import static io.trino.plugin.hive.TableType.EXTERNAL_TABLE; import static io.trino.plugin.hive.TableType.VIRTUAL_VIEW; import static io.trino.plugin.hive.metastore.StorageFormat.fromHiveStorageFormat; +import static io.trino.plugin.hive.metastore.Table.TABLE_COMMENT; +import static io.trino.plugin.hive.metastore.TableInfo.ICEBERG_MATERIALIZED_VIEW_COMMENT; +import static io.trino.plugin.hive.metastore.TableInfo.PRESTO_VIEW_COMMENT; import static io.trino.plugin.hive.metastore.thrift.ThriftMetastoreUtil.toMetastoreApiTable; import static io.trino.plugin.hive.util.HiveUtil.checkCondition; import static java.lang.String.format; @@ -108,7 +110,7 @@ private static CoralTableRedirectionResolver coralTableRedirectionResolver( format("%s is redirected to %s, but that relation cannot be found", schemaTableName, target))); List columns = tableSchema.getColumns().stream() .filter(columnSchema -> !columnSchema.isHidden()) - .map(columnSchema -> new Column(columnSchema.getName(), toHiveType(columnSchema.getType()), Optional.empty() /* comment */)) + .map(columnSchema -> new Column(columnSchema.getName(), toHiveType(columnSchema.getType()), Optional.empty() /* comment */, Map.of())) .collect(toImmutableList()); Table table = Table.builder() .setDatabaseName(schemaTableName.getSchemaName()) @@ -123,13 +125,24 @@ private static CoralTableRedirectionResolver coralTableRedirectionResolver( }); } - public static final String ICEBERG_MATERIALIZED_VIEW_COMMENT = "Presto Materialized View"; public static final String PRESTO_VIEW_FLAG = "presto_view"; static final String VIEW_PREFIX = "/* Presto View: "; static final String VIEW_SUFFIX = " */"; private static final JsonCodec VIEW_CODEC = new JsonCodecFactory(new ObjectMapperProvider()).jsonCodec(ConnectorViewDefinition.class); + /** + * Returns true if table represents a Hive view, Trino/Presto view, materialized view or anything + * else that gets registered using table type "VIRTUAL_VIEW". + * Note: this method returns false for a table that represents Hive's own materialized view + * ("MATERIALIZED_VIEW" table type). Hive own's materialized views are currently treated as ordinary + * tables by Trino. + */ + public static boolean isSomeKindOfAView(Table table) + { + return table.getTableType().equals(VIRTUAL_VIEW.name()); + } + public static boolean isPrestoView(Table table) { return isPrestoView(table.getParameters()); @@ -141,6 +154,28 @@ public static boolean isPrestoView(Map tableParameters) return "true".equals(tableParameters.get(PRESTO_VIEW_FLAG)); } + /** + * Returns true when the table represents a "Trino view" (AKA "presto view"). + * Returns false for Hive views or Trino materialized views. + */ + public static boolean isTrinoView(Table table) + { + return isTrinoView(table.getTableType(), table.getParameters()); + } + + /** + * Returns true when the table represents a "Trino view" (AKA "presto view"). + * Returns false for Hive views or Trino materialized views. + */ + public static boolean isTrinoView(String tableType, Map tableParameters) + { + // A Trino view can be recognized by table type "VIRTUAL_VIEW" and table parameters presto_view="true" and comment="Presto View" since their first implementation see + // https://github.com/trinodb/trino/blame/38bd0dff736024f3ae01dbbe7d1db5bd1d50c43e/presto-hive/src/main/java/com/facebook/presto/hive/HiveMetadata.java#L902. + return tableType.equals(VIRTUAL_VIEW.name()) && + "true".equals(tableParameters.get(PRESTO_VIEW_FLAG)) && + PRESTO_VIEW_COMMENT.equalsIgnoreCase(tableParameters.get(TABLE_COMMENT)); + } + public static boolean isHiveOrPrestoView(Table table) { return isHiveOrPrestoView(table.getTableType()); diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/acid/AcidSchema.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/acid/AcidSchema.java index e63a5b8df801..149cdb2a2c3e 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/acid/AcidSchema.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/acid/AcidSchema.java @@ -14,6 +14,7 @@ package io.trino.plugin.hive.acid; import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; import io.trino.plugin.hive.HiveType; import io.trino.plugin.hive.HiveTypeName; import io.trino.spi.type.RowType; @@ -21,8 +22,8 @@ import io.trino.spi.type.Type; import java.util.List; +import java.util.Map; import java.util.Optional; -import java.util.Properties; import static com.google.common.base.Preconditions.checkArgument; import static io.trino.plugin.hive.HiveType.HIVE_INT; @@ -61,17 +62,17 @@ public final class AcidSchema private AcidSchema() {} - public static Properties createAcidSchema(HiveType rowType) + public static Map createAcidSchema(HiveType rowType) { - Properties hiveAcidSchema = new Properties(); - hiveAcidSchema.setProperty(LIST_COLUMNS, String.join(",", ACID_COLUMN_NAMES)); - // We must supply an accurate row type, because Apache ORC code we don't control has a consistency - // check that the layout of this "row" must agree with the layout of an inserted row. - hiveAcidSchema.setProperty(LIST_COLUMN_TYPES, createAcidColumnHiveTypes(rowType).stream() - .map(HiveType::getHiveTypeName) - .map(HiveTypeName::toString) - .collect(joining(":"))); - return hiveAcidSchema; + return ImmutableMap.builder() + .put(LIST_COLUMNS, String.join(",", ACID_COLUMN_NAMES)) + // We must supply an accurate row type, because Apache ORC code we don't control has a consistency + // check that the layout of this "row" must agree with the layout of an inserted row. + .put(LIST_COLUMN_TYPES, createAcidColumnHiveTypes(rowType).stream() + .map(HiveType::getHiveTypeName) + .map(HiveTypeName::toString) + .collect(joining(":"))) + .buildOrThrow(); } public static Type createRowType(List names, List types) diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/acid/AcidTransaction.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/acid/AcidTransaction.java index dc442e829baa..fba0559124ab 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/acid/AcidTransaction.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/acid/AcidTransaction.java @@ -17,6 +17,8 @@ import com.fasterxml.jackson.annotation.JsonIgnore; import com.fasterxml.jackson.annotation.JsonProperty; +import java.util.OptionalLong; + import static com.google.common.base.MoreObjects.toStringHelper; import static com.google.common.base.Preconditions.checkState; import static io.trino.plugin.hive.acid.AcidOperation.CREATE_TABLE; @@ -88,6 +90,12 @@ public long getWriteId() return writeId; } + @JsonIgnore + public OptionalLong getOptionalWriteId() + { + return isAcidTransactionRunning() ? OptionalLong.of(writeId) : OptionalLong.empty(); + } + private void ensureTransactionRunning(String description) { checkState(isAcidTransactionRunning(), "Not in ACID transaction while %s", description); diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/avro/AvroFileWriterFactory.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/avro/AvroFileWriterFactory.java index 8b3d420862fe..32df7a43786b 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/avro/AvroFileWriterFactory.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/avro/AvroFileWriterFactory.java @@ -38,18 +38,17 @@ import java.io.Closeable; import java.util.List; +import java.util.Map; import java.util.Optional; import java.util.OptionalInt; -import java.util.Properties; import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.collect.ImmutableList.toImmutableList; import static io.trino.memory.context.AggregatedMemoryContext.newSimpleAggregatedMemoryContext; import static io.trino.plugin.hive.HiveErrorCode.HIVE_WRITER_OPEN_ERROR; -import static io.trino.plugin.hive.HiveMetadata.PRESTO_QUERY_ID_NAME; -import static io.trino.plugin.hive.HiveMetadata.PRESTO_VERSION_NAME; +import static io.trino.plugin.hive.HiveMetadata.TRINO_QUERY_ID_NAME; +import static io.trino.plugin.hive.HiveMetadata.TRINO_VERSION_NAME; import static io.trino.plugin.hive.HiveSessionProperties.getTimestampPrecision; -import static io.trino.plugin.hive.HiveSessionProperties.isAvroNativeWriterEnabled; import static io.trino.plugin.hive.util.HiveClassNames.AVRO_CONTAINER_OUTPUT_FORMAT_CLASS; import static io.trino.plugin.hive.util.HiveUtil.getColumnNames; import static io.trino.plugin.hive.util.HiveUtil.getColumnTypes; @@ -79,16 +78,13 @@ public Optional createFileWriter( List inputColumnNames, StorageFormat storageFormat, HiveCompressionCodec compressionCodec, - Properties schema, + Map schema, ConnectorSession session, OptionalInt bucketNumber, AcidTransaction transaction, boolean useAcidSchema, WriterKind writerKind) { - if (!isAvroNativeWriterEnabled(session)) { - return Optional.empty(); - } if (!AVRO_CONTAINER_OUTPUT_FORMAT_CLASS.equals(storageFormat.getOutputFormat())) { return Optional.empty(); } @@ -130,8 +126,8 @@ public Optional createFileWriter( inputColumnTypes, compressionKind, ImmutableMap.builder() - .put(PRESTO_VERSION_NAME, nodeVersion.toString()) - .put(PRESTO_QUERY_ID_NAME, session.getQueryId()) + .put(TRINO_VERSION_NAME, nodeVersion.toString()) + .put(TRINO_QUERY_ID_NAME, session.getQueryId()) .buildOrThrow())); } catch (Exception e) { diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/avro/AvroHiveFileUtils.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/avro/AvroHiveFileUtils.java index a6ca72874e13..d1cba1e19e73 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/avro/AvroHiveFileUtils.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/avro/AvroHiveFileUtils.java @@ -39,11 +39,10 @@ import java.util.ArrayList; import java.util.Arrays; import java.util.List; +import java.util.Map; import java.util.Optional; -import java.util.Properties; import java.util.concurrent.atomic.AtomicInteger; -import static io.trino.plugin.hive.HiveMetadata.TABLE_COMMENT; import static io.trino.plugin.hive.avro.AvroHiveConstants.CHAR_TYPE_LOGICAL_NAME; import static io.trino.plugin.hive.avro.AvroHiveConstants.SCHEMA_DOC; import static io.trino.plugin.hive.avro.AvroHiveConstants.SCHEMA_LITERAL; @@ -54,6 +53,7 @@ import static io.trino.plugin.hive.avro.AvroHiveConstants.TABLE_NAME; import static io.trino.plugin.hive.avro.AvroHiveConstants.VARCHAR_AND_CHAR_LOGICAL_TYPE_LENGTH_PROP; import static io.trino.plugin.hive.avro.AvroHiveConstants.VARCHAR_TYPE_LOGICAL_NAME; +import static io.trino.plugin.hive.metastore.Table.TABLE_COMMENT; import static io.trino.plugin.hive.util.HiveUtil.getColumnNames; import static io.trino.plugin.hive.util.HiveUtil.getColumnTypes; import static io.trino.plugin.hive.util.SerdeConstants.LIST_COLUMN_COMMENTS; @@ -67,17 +67,17 @@ public final class AvroHiveFileUtils private AvroHiveFileUtils() {} // Lifted and shifted from org.apache.hadoop.hive.serde2.avro.AvroSerdeUtils.determineSchemaOrThrowException - public static Schema determineSchemaOrThrowException(TrinoFileSystem fileSystem, Properties properties) + public static Schema determineSchemaOrThrowException(TrinoFileSystem fileSystem, Map properties) throws IOException { // Try pull schema from literal table property - String schemaString = properties.getProperty(SCHEMA_LITERAL, ""); + String schemaString = properties.getOrDefault(SCHEMA_LITERAL, ""); if (!schemaString.isBlank() && !schemaString.equals(SCHEMA_NONE)) { return getSchemaParser().parse(schemaString); } // Try pull schema directly from URL - String schemaURL = properties.getProperty(SCHEMA_URL, ""); + String schemaURL = properties.getOrDefault(SCHEMA_URL, ""); if (!schemaURL.isBlank()) { TrinoInputFile schemaFile = fileSystem.newInputFile(Location.of(schemaURL)); if (!schemaFile.exists()) { @@ -90,37 +90,35 @@ public static Schema determineSchemaOrThrowException(TrinoFileSystem fileSystem, throw new IOException("Unable to read avro schema file from given path: " + schemaURL, e); } } - Schema schema = getSchemaFromProperties(properties); - properties.setProperty(SCHEMA_LITERAL, schema.toString()); - return schema; + return getSchemaFromProperties(properties); } - private static Schema getSchemaFromProperties(Properties properties) + private static Schema getSchemaFromProperties(Map schema) throws IOException { - List columnNames = getColumnNames(properties); - List columnTypes = getColumnTypes(properties); + List columnNames = getColumnNames(schema); + List columnTypes = getColumnTypes(schema); if (columnNames.isEmpty() || columnTypes.isEmpty()) { throw new IOException("Unable to parse column names or column types from job properties to create Avro Schema"); } if (columnNames.size() != columnTypes.size()) { throw new IllegalArgumentException("Avro Schema initialization failed. Number of column name and column type differs. columnNames = %s, columnTypes = %s".formatted(columnNames, columnTypes)); } - List columnComments = Optional.ofNullable(properties.getProperty(LIST_COLUMN_COMMENTS)) + List columnComments = Optional.ofNullable(schema.get(LIST_COLUMN_COMMENTS)) .filter(not(String::isBlank)) .map(Splitter.on('\0')::splitToList) .orElse(emptyList()); - final String tableName = properties.getProperty(TABLE_NAME); - final String tableComment = properties.getProperty(TABLE_COMMENT); + final String tableName = schema.get(TABLE_NAME); + final String tableComment = schema.get(TABLE_COMMENT); return constructSchemaFromParts( columnNames, columnTypes, columnComments, - Optional.ofNullable(properties.getProperty(SCHEMA_NAMESPACE)), - Optional.ofNullable(properties.getProperty(SCHEMA_NAME, tableName)), - Optional.ofNullable(properties.getProperty(SCHEMA_DOC, tableComment))); + Optional.ofNullable(schema.get(SCHEMA_NAMESPACE)), + Optional.ofNullable(schema.getOrDefault(SCHEMA_NAME, tableName)), + Optional.ofNullable(schema.getOrDefault(SCHEMA_DOC, tableComment))); } private static Schema constructSchemaFromParts(List columnNames, List columnTypes, diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/avro/AvroPageSourceFactory.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/avro/AvroPageSourceFactory.java index 887ba146cae2..2e904935a6ad 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/avro/AvroPageSourceFactory.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/avro/AvroPageSourceFactory.java @@ -24,14 +24,12 @@ import io.trino.filesystem.memory.MemoryInputFile; import io.trino.hive.formats.avro.AvroTypeException; import io.trino.plugin.hive.AcidInfo; -import io.trino.plugin.hive.FileFormatDataSourceStats; import io.trino.plugin.hive.HiveColumnHandle; import io.trino.plugin.hive.HivePageSourceFactory; import io.trino.plugin.hive.HiveTimestampPrecision; import io.trino.plugin.hive.ReaderColumns; import io.trino.plugin.hive.ReaderPageSource; import io.trino.plugin.hive.acid.AcidTransaction; -import io.trino.plugin.hive.fs.MonitoredInputFile; import io.trino.spi.TrinoException; import io.trino.spi.connector.ConnectorSession; import io.trino.spi.connector.EmptyPageSource; @@ -48,7 +46,6 @@ import java.util.Objects; import java.util.Optional; import java.util.OptionalInt; -import java.util.Properties; import java.util.Set; import java.util.UUID; @@ -58,11 +55,9 @@ import static io.trino.plugin.hive.HiveErrorCode.HIVE_CANNOT_OPEN_SPLIT; import static io.trino.plugin.hive.HivePageSourceProvider.projectBaseColumns; import static io.trino.plugin.hive.HiveSessionProperties.getTimestampPrecision; -import static io.trino.plugin.hive.HiveSessionProperties.isAvroNativeReaderEnabled; import static io.trino.plugin.hive.ReaderPageSource.noProjectionAdaptation; import static io.trino.plugin.hive.avro.AvroHiveFileUtils.wrapInUnionWithNull; import static io.trino.plugin.hive.util.HiveClassNames.AVRO_SERDE_CLASS; -import static io.trino.plugin.hive.util.HiveUtil.getDeserializerClassName; import static io.trino.plugin.hive.util.HiveUtil.splitError; import static java.lang.Math.min; import static java.util.Objects.requireNonNull; @@ -73,13 +68,11 @@ public class AvroPageSourceFactory private static final DataSize BUFFER_SIZE = DataSize.of(8, DataSize.Unit.MEGABYTE); private final TrinoFileSystemFactory trinoFileSystemFactory; - private final FileFormatDataSourceStats stats; @Inject - public AvroPageSourceFactory(TrinoFileSystemFactory trinoFileSystemFactory, FileFormatDataSourceStats stats) + public AvroPageSourceFactory(TrinoFileSystemFactory trinoFileSystemFactory) { this.trinoFileSystemFactory = requireNonNull(trinoFileSystemFactory, "trinoFileSystemFactory is null"); - this.stats = requireNonNull(stats, "stats is null"); } @Override @@ -89,7 +82,8 @@ public Optional createPageSource( long start, long length, long estimatedFileSize, - Properties schema, + long fileModifiedTime, + io.trino.plugin.hive.Schema schema, List columns, TupleDomain effectivePredicate, Optional acidInfo, @@ -97,10 +91,7 @@ public Optional createPageSource( boolean originalFile, AcidTransaction transaction) { - if (!isAvroNativeReaderEnabled(session)) { - return Optional.empty(); - } - else if (!AVRO_SERDE_CLASS.equals(getDeserializerClassName(schema))) { + if (!AVRO_SERDE_CLASS.equals(schema.serializationLibraryName())) { return Optional.empty(); } checkArgument(acidInfo.isEmpty(), "Acid is not supported"); @@ -115,12 +106,12 @@ else if (!AVRO_SERDE_CLASS.equals(getDeserializerClassName(schema))) { } TrinoFileSystem trinoFileSystem = trinoFileSystemFactory.create(session.getIdentity()); - TrinoInputFile inputFile = new MonitoredInputFile(stats, trinoFileSystem.newInputFile(path)); + TrinoInputFile inputFile = trinoFileSystem.newInputFile(path); HiveTimestampPrecision hiveTimestampPrecision = getTimestampPrecision(session); Schema tableSchema; try { - tableSchema = AvroHiveFileUtils.determineSchemaOrThrowException(trinoFileSystem, schema); + tableSchema = AvroHiveFileUtils.determineSchemaOrThrowException(trinoFileSystem, schema.serdeProperties()); } catch (IOException | org.apache.avro.AvroTypeException e) { throw new TrinoException(HIVE_CANNOT_OPEN_SPLIT, "Unable to load or parse schema", e); diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/aws/AwsSdkClientCoreStats.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/aws/AwsSdkClientCoreStats.java deleted file mode 100644 index 0fb4d6aa918e..000000000000 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/aws/AwsSdkClientCoreStats.java +++ /dev/null @@ -1,183 +0,0 @@ -/* - * 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.hive.aws; - -import com.amazonaws.Request; -import com.amazonaws.Response; -import com.amazonaws.metrics.RequestMetricCollector; -import com.amazonaws.util.AWSRequestMetrics; -import com.amazonaws.util.TimingInfo; -import com.google.errorprone.annotations.ThreadSafe; -import io.airlift.stats.CounterStat; -import io.airlift.stats.TimeStat; -import org.weakref.jmx.Managed; -import org.weakref.jmx.Nested; - -import java.util.List; -import java.util.concurrent.atomic.AtomicLong; - -import static com.amazonaws.util.AWSRequestMetrics.Field.ClientExecuteTime; -import static com.amazonaws.util.AWSRequestMetrics.Field.HttpClientPoolAvailableCount; -import static com.amazonaws.util.AWSRequestMetrics.Field.HttpClientPoolLeasedCount; -import static com.amazonaws.util.AWSRequestMetrics.Field.HttpClientPoolPendingCount; -import static com.amazonaws.util.AWSRequestMetrics.Field.HttpClientRetryCount; -import static com.amazonaws.util.AWSRequestMetrics.Field.HttpRequestTime; -import static com.amazonaws.util.AWSRequestMetrics.Field.RequestCount; -import static com.amazonaws.util.AWSRequestMetrics.Field.RetryPauseTime; -import static com.amazonaws.util.AWSRequestMetrics.Field.ThrottleException; -import static java.util.Objects.requireNonNull; -import static java.util.concurrent.TimeUnit.MILLISECONDS; - -@ThreadSafe -public final class AwsSdkClientCoreStats -{ - private final CounterStat awsRequestCount = new CounterStat(); - private final CounterStat awsRetryCount = new CounterStat(); - private final CounterStat awsThrottleExceptions = new CounterStat(); - private final TimeStat awsRequestTime = new TimeStat(MILLISECONDS); - private final TimeStat awsClientExecuteTime = new TimeStat(MILLISECONDS); - private final TimeStat awsClientRetryPauseTime = new TimeStat(MILLISECONDS); - private final AtomicLong awsHttpClientPoolAvailableCount = new AtomicLong(); - private final AtomicLong awsHttpClientPoolLeasedCount = new AtomicLong(); - private final AtomicLong awsHttpClientPoolPendingCount = new AtomicLong(); - - @Managed - @Nested - public CounterStat getAwsRequestCount() - { - return awsRequestCount; - } - - @Managed - @Nested - public CounterStat getAwsRetryCount() - { - return awsRetryCount; - } - - @Managed - @Nested - public CounterStat getAwsThrottleExceptions() - { - return awsThrottleExceptions; - } - - @Managed - @Nested - public TimeStat getAwsRequestTime() - { - return awsRequestTime; - } - - @Managed - @Nested - public TimeStat getAwsClientExecuteTime() - { - return awsClientExecuteTime; - } - - @Managed - @Nested - public TimeStat getAwsClientRetryPauseTime() - { - return awsClientRetryPauseTime; - } - - @Managed - public long getAwsHttpClientPoolAvailableCount() - { - return awsHttpClientPoolAvailableCount.get(); - } - - @Managed - public long getAwsHttpClientPoolLeasedCount() - { - return awsHttpClientPoolLeasedCount.get(); - } - - @Managed - public long getAwsHttpClientPoolPendingCount() - { - return awsHttpClientPoolPendingCount.get(); - } - - public AwsSdkClientCoreRequestMetricCollector newRequestMetricCollector() - { - return new AwsSdkClientCoreRequestMetricCollector(this); - } - - public static class AwsSdkClientCoreRequestMetricCollector - extends RequestMetricCollector - { - private final AwsSdkClientCoreStats stats; - - protected AwsSdkClientCoreRequestMetricCollector(AwsSdkClientCoreStats stats) - { - this.stats = requireNonNull(stats, "stats is null"); - } - - @Override - public void collectMetrics(Request request, Response response) - { - TimingInfo timingInfo = request.getAWSRequestMetrics().getTimingInfo(); - - Number requestCounts = timingInfo.getCounter(RequestCount.name()); - if (requestCounts != null) { - stats.awsRequestCount.update(requestCounts.longValue()); - } - - Number retryCounts = timingInfo.getCounter(HttpClientRetryCount.name()); - if (retryCounts != null) { - stats.awsRetryCount.update(retryCounts.longValue()); - } - - Number throttleExceptions = timingInfo.getCounter(ThrottleException.name()); - if (throttleExceptions != null) { - stats.awsThrottleExceptions.update(throttleExceptions.longValue()); - } - - Number httpClientPoolAvailableCount = timingInfo.getCounter(HttpClientPoolAvailableCount.name()); - if (httpClientPoolAvailableCount != null) { - stats.awsHttpClientPoolAvailableCount.set(httpClientPoolAvailableCount.longValue()); - } - - Number httpClientPoolLeasedCount = timingInfo.getCounter(HttpClientPoolLeasedCount.name()); - if (httpClientPoolLeasedCount != null) { - stats.awsHttpClientPoolLeasedCount.set(httpClientPoolLeasedCount.longValue()); - } - - Number httpClientPoolPendingCount = timingInfo.getCounter(HttpClientPoolPendingCount.name()); - if (httpClientPoolPendingCount != null) { - stats.awsHttpClientPoolPendingCount.set(httpClientPoolPendingCount.longValue()); - } - - recordSubTimingDurations(timingInfo, HttpRequestTime, stats.awsRequestTime); - recordSubTimingDurations(timingInfo, ClientExecuteTime, stats.awsClientExecuteTime); - recordSubTimingDurations(timingInfo, RetryPauseTime, stats.awsClientRetryPauseTime); - } - - private static void recordSubTimingDurations(TimingInfo timingInfo, AWSRequestMetrics.Field field, TimeStat timeStat) - { - List subTimings = timingInfo.getAllSubMeasurements(field.name()); - if (subTimings != null) { - for (TimingInfo subTiming : subTimings) { - Long endTimeNanos = subTiming.getEndTimeNanoIfKnown(); - if (endTimeNanos != null) { - timeStat.addNanos(endTimeNanos - subTiming.getStartTimeNano()); - } - } - } - } - } -} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/aws/athena/PartitionProjectionMetastoreDecorator.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/aws/athena/PartitionProjectionMetastoreDecorator.java deleted file mode 100644 index c221e87e5ad9..000000000000 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/aws/athena/PartitionProjectionMetastoreDecorator.java +++ /dev/null @@ -1,97 +0,0 @@ -/* - * 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.hive.aws.athena; - -import com.google.inject.Inject; -import io.trino.plugin.hive.metastore.ForwardingHiveMetastore; -import io.trino.plugin.hive.metastore.HiveMetastore; -import io.trino.plugin.hive.metastore.HiveMetastoreDecorator; -import io.trino.plugin.hive.metastore.Partition; -import io.trino.plugin.hive.metastore.Table; -import io.trino.spi.TrinoException; -import io.trino.spi.predicate.TupleDomain; - -import java.util.List; -import java.util.Map; -import java.util.Optional; - -import static io.trino.plugin.hive.HiveErrorCode.HIVE_TABLE_DROPPED_DURING_QUERY; -import static java.util.Objects.requireNonNull; - -public class PartitionProjectionMetastoreDecorator - implements HiveMetastoreDecorator -{ - private final PartitionProjectionService partitionProjectionService; - - @Inject - public PartitionProjectionMetastoreDecorator(PartitionProjectionService partitionProjectionService) - { - this.partitionProjectionService = requireNonNull(partitionProjectionService, "partitionProjectionService is null"); - } - - @Override - public int getPriority() - { - return PRIORITY_PARTITION_PROJECTION; - } - - @Override - public HiveMetastore decorate(HiveMetastore hiveMetastore) - { - return new PartitionProjectionMetastore(hiveMetastore, partitionProjectionService); - } - - private static class PartitionProjectionMetastore - extends ForwardingHiveMetastore - { - private final PartitionProjectionService partitionProjectionService; - - public PartitionProjectionMetastore(HiveMetastore hiveMetastore, PartitionProjectionService partitionProjectionService) - { - super(hiveMetastore); - this.partitionProjectionService = requireNonNull(partitionProjectionService, "partitionProjectionService is null"); - } - - @Override - public Optional> getPartitionNamesByFilter(String databaseName, String tableName, List columnNames, TupleDomain partitionKeysFilter) - { - Table table = super.getTable(databaseName, tableName) - .orElseThrow(() -> new TrinoException(HIVE_TABLE_DROPPED_DURING_QUERY, "Table does not exists: " + tableName)); - - Optional projection = getPartitionProjection(table); - if (projection.isPresent()) { - return projection.get().getProjectedPartitionNamesByFilter(columnNames, partitionKeysFilter); - } - - return super.getPartitionNamesByFilter(databaseName, tableName, columnNames, partitionKeysFilter); - } - - @Override - public Map> getPartitionsByNames(Table table, List partitionNames) - { - Optional projection = getPartitionProjection(table); - if (projection.isPresent()) { - return projection.get().getProjectedPartitionsByNames(table, partitionNames); - } - return super.getPartitionsByNames(table, partitionNames); - } - - private Optional getPartitionProjection(Table table) - { - return partitionProjectionService.getPartitionProjectionFromTable(table) - .filter(PartitionProjection::isEnabled); - } - } -} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/aws/athena/PartitionProjectionModule.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/aws/athena/PartitionProjectionModule.java deleted file mode 100644 index e1c3ae1ff185..000000000000 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/aws/athena/PartitionProjectionModule.java +++ /dev/null @@ -1,51 +0,0 @@ -/* - * 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.hive.aws.athena; - -import com.google.inject.Binder; -import com.google.inject.Scopes; -import com.google.inject.multibindings.MapBinder; -import io.airlift.configuration.AbstractConfigurationAwareModule; -import io.trino.plugin.hive.aws.athena.projection.DateProjectionFactory; -import io.trino.plugin.hive.aws.athena.projection.EnumProjectionFactory; -import io.trino.plugin.hive.aws.athena.projection.InjectedProjectionFactory; -import io.trino.plugin.hive.aws.athena.projection.IntegerProjectionFactory; -import io.trino.plugin.hive.aws.athena.projection.ProjectionFactory; -import io.trino.plugin.hive.aws.athena.projection.ProjectionType; -import io.trino.plugin.hive.metastore.HiveMetastoreDecorator; - -import static com.google.inject.multibindings.MapBinder.newMapBinder; -import static com.google.inject.multibindings.Multibinder.newSetBinder; - -public class PartitionProjectionModule - extends AbstractConfigurationAwareModule -{ - @Override - public void setup(Binder binder) - { - MapBinder projectionFactoriesBinder = - newMapBinder(binder, ProjectionType.class, ProjectionFactory.class); - projectionFactoriesBinder.addBinding(ProjectionType.ENUM).to(EnumProjectionFactory.class).in(Scopes.SINGLETON); - projectionFactoriesBinder.addBinding(ProjectionType.INTEGER).to(IntegerProjectionFactory.class).in(Scopes.SINGLETON); - projectionFactoriesBinder.addBinding(ProjectionType.DATE).to(DateProjectionFactory.class).in(Scopes.SINGLETON); - projectionFactoriesBinder.addBinding(ProjectionType.INJECTED).to(InjectedProjectionFactory.class).in(Scopes.SINGLETON); - - binder.bind(PartitionProjectionService.class).in(Scopes.SINGLETON); - - newSetBinder(binder, HiveMetastoreDecorator.class) - .addBinding() - .to(PartitionProjectionMetastoreDecorator.class) - .in(Scopes.SINGLETON); - } -} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/aws/athena/PartitionProjectionProperties.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/aws/athena/PartitionProjectionProperties.java deleted file mode 100644 index 38733e0f1f74..000000000000 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/aws/athena/PartitionProjectionProperties.java +++ /dev/null @@ -1,109 +0,0 @@ -/* - * 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.hive.aws.athena; - -import java.util.Map; -import java.util.Optional; -import java.util.function.Function; - -import static io.trino.plugin.hive.aws.athena.projection.Projection.invalidProjectionException; -import static java.lang.String.format; - -public final class PartitionProjectionProperties -{ - /** - * General properties suffixes - */ - static final String TABLE_PROJECTION_ENABLED_SUFFIX = "enabled"; - /** - * Forces Trino to not use Athena table projection for a given table. - * Kill switch to be used as a workaround if compatibility issues are found. - */ - static final String TABLE_PROJECTION_IGNORE_SUFFIX = "ignore"; - static final String COLUMN_PROJECTION_TYPE_SUFFIX = "type"; - static final String COLUMN_PROJECTION_VALUES_SUFFIX = "values"; - static final String COLUMN_PROJECTION_RANGE_SUFFIX = "range"; - static final String COLUMN_PROJECTION_INTERVAL_SUFFIX = "interval"; - static final String COLUMN_PROJECTION_DIGITS_SUFFIX = "digits"; - static final String COLUMN_PROJECTION_FORMAT_SUFFIX = "format"; - - /** - * Metastore table properties - */ - private static final String METASTORE_PROPERTY_SEPARATOR = "."; - - private static final String METASTORE_PROPERTY_PREFIX = "projection" + METASTORE_PROPERTY_SEPARATOR; - - static final String METASTORE_PROPERTY_PROJECTION_INTERVAL_UNIT_SUFFIX = "interval" + METASTORE_PROPERTY_SEPARATOR + "unit"; - - static final String METASTORE_PROPERTY_PROJECTION_ENABLED = METASTORE_PROPERTY_PREFIX + TABLE_PROJECTION_ENABLED_SUFFIX; - static final String METASTORE_PROPERTY_PROJECTION_LOCATION_TEMPLATE = "storage" + METASTORE_PROPERTY_SEPARATOR + "location" + METASTORE_PROPERTY_SEPARATOR + "template"; - /** - * See {@link #TABLE_PROJECTION_IGNORE_SUFFIX } to understand duplication with enable property - **/ - static final String METASTORE_PROPERTY_PROJECTION_IGNORE = "trino" + METASTORE_PROPERTY_SEPARATOR + "partition_projection" + METASTORE_PROPERTY_SEPARATOR + TABLE_PROJECTION_IGNORE_SUFFIX; - - /** - * Trino table properties - */ - private static final String PROPERTY_KEY_SEPARATOR = "_"; - - static final String PROPERTY_KEY_PREFIX = "partition" + PROPERTY_KEY_SEPARATOR + "projection" + PROPERTY_KEY_SEPARATOR; - - private static final String PROPERTY_KEY_SUFFIX_COLUMN_PROJECTION_INTERVAL_UNIT = "interval" + PROPERTY_KEY_SEPARATOR + "unit"; - - public static final String PARTITION_PROJECTION_ENABLED = PROPERTY_KEY_PREFIX + TABLE_PROJECTION_ENABLED_SUFFIX; - public static final String PARTITION_PROJECTION_LOCATION_TEMPLATE = PROPERTY_KEY_PREFIX + "location" + PROPERTY_KEY_SEPARATOR + "template"; - /** - * See {@link #TABLE_PROJECTION_IGNORE_SUFFIX } to understand duplication with enable property - **/ - public static final String PARTITION_PROJECTION_IGNORE = PROPERTY_KEY_PREFIX + TABLE_PROJECTION_IGNORE_SUFFIX; - - public static final String COLUMN_PROJECTION_TYPE = PROPERTY_KEY_PREFIX + COLUMN_PROJECTION_TYPE_SUFFIX; - public static final String COLUMN_PROJECTION_VALUES = PROPERTY_KEY_PREFIX + COLUMN_PROJECTION_VALUES_SUFFIX; - public static final String COLUMN_PROJECTION_RANGE = PROPERTY_KEY_PREFIX + COLUMN_PROJECTION_RANGE_SUFFIX; - public static final String COLUMN_PROJECTION_INTERVAL = PROPERTY_KEY_PREFIX + COLUMN_PROJECTION_INTERVAL_SUFFIX; - public static final String COLUMN_PROJECTION_INTERVAL_UNIT = PROPERTY_KEY_PREFIX + PROPERTY_KEY_SUFFIX_COLUMN_PROJECTION_INTERVAL_UNIT; - public static final String COLUMN_PROJECTION_DIGITS = PROPERTY_KEY_PREFIX + COLUMN_PROJECTION_DIGITS_SUFFIX; - public static final String COLUMN_PROJECTION_FORMAT = PROPERTY_KEY_PREFIX + COLUMN_PROJECTION_FORMAT_SUFFIX; - - static String getMetastoreProjectionPropertyKey(String columnName, String propertyKeySuffix) - { - return METASTORE_PROPERTY_PREFIX + columnName + METASTORE_PROPERTY_SEPARATOR + propertyKeySuffix; - } - - public static T getProjectionPropertyRequiredValue( - String columnName, - Map columnProjectionProperties, - String propertyKey, - Function decoder) - { - return getProjectionPropertyValue(columnProjectionProperties, propertyKey, decoder) - .orElseThrow(() -> invalidProjectionException(columnName, format("Missing required property: '%s'", propertyKey))); - } - - public static Optional getProjectionPropertyValue( - Map columnProjectionProperties, - String propertyKey, - Function decoder) - { - return Optional.ofNullable( - columnProjectionProperties.get(propertyKey)) - .map(value -> decoder.apply(value)); - } - - private PartitionProjectionProperties() - { - } -} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/aws/athena/PartitionProjectionService.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/aws/athena/PartitionProjectionService.java deleted file mode 100644 index d9175050adec..000000000000 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/aws/athena/PartitionProjectionService.java +++ /dev/null @@ -1,381 +0,0 @@ -/* - * 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.hive.aws.athena; - -import com.google.common.base.Joiner; -import com.google.common.base.Splitter; -import com.google.common.collect.ImmutableMap; -import com.google.common.collect.ImmutableSet; -import com.google.inject.Inject; -import io.trino.plugin.hive.HiveConfig; -import io.trino.plugin.hive.aws.athena.projection.Projection; -import io.trino.plugin.hive.aws.athena.projection.ProjectionFactory; -import io.trino.plugin.hive.aws.athena.projection.ProjectionType; -import io.trino.plugin.hive.metastore.Column; -import io.trino.plugin.hive.metastore.Table; -import io.trino.spi.TrinoException; -import io.trino.spi.connector.ColumnMetadata; -import io.trino.spi.connector.ConnectorTableMetadata; -import io.trino.spi.type.Type; -import io.trino.spi.type.TypeManager; - -import java.time.temporal.ChronoUnit; -import java.util.List; -import java.util.Locale; -import java.util.Map; -import java.util.Objects; -import java.util.Optional; -import java.util.Set; -import java.util.function.Function; -import java.util.stream.Collectors; - -import static com.google.common.collect.ImmutableList.toImmutableList; -import static com.google.common.collect.ImmutableMap.toImmutableMap; -import static io.trino.plugin.hive.HiveTableProperties.getPartitionedBy; -import static io.trino.plugin.hive.HiveTimestampPrecision.DEFAULT_PRECISION; -import static io.trino.plugin.hive.aws.athena.PartitionProjectionProperties.COLUMN_PROJECTION_DIGITS; -import static io.trino.plugin.hive.aws.athena.PartitionProjectionProperties.COLUMN_PROJECTION_DIGITS_SUFFIX; -import static io.trino.plugin.hive.aws.athena.PartitionProjectionProperties.COLUMN_PROJECTION_FORMAT; -import static io.trino.plugin.hive.aws.athena.PartitionProjectionProperties.COLUMN_PROJECTION_FORMAT_SUFFIX; -import static io.trino.plugin.hive.aws.athena.PartitionProjectionProperties.COLUMN_PROJECTION_INTERVAL; -import static io.trino.plugin.hive.aws.athena.PartitionProjectionProperties.COLUMN_PROJECTION_INTERVAL_SUFFIX; -import static io.trino.plugin.hive.aws.athena.PartitionProjectionProperties.COLUMN_PROJECTION_INTERVAL_UNIT; -import static io.trino.plugin.hive.aws.athena.PartitionProjectionProperties.COLUMN_PROJECTION_RANGE; -import static io.trino.plugin.hive.aws.athena.PartitionProjectionProperties.COLUMN_PROJECTION_RANGE_SUFFIX; -import static io.trino.plugin.hive.aws.athena.PartitionProjectionProperties.COLUMN_PROJECTION_TYPE; -import static io.trino.plugin.hive.aws.athena.PartitionProjectionProperties.COLUMN_PROJECTION_TYPE_SUFFIX; -import static io.trino.plugin.hive.aws.athena.PartitionProjectionProperties.COLUMN_PROJECTION_VALUES; -import static io.trino.plugin.hive.aws.athena.PartitionProjectionProperties.COLUMN_PROJECTION_VALUES_SUFFIX; -import static io.trino.plugin.hive.aws.athena.PartitionProjectionProperties.METASTORE_PROPERTY_PROJECTION_ENABLED; -import static io.trino.plugin.hive.aws.athena.PartitionProjectionProperties.METASTORE_PROPERTY_PROJECTION_IGNORE; -import static io.trino.plugin.hive.aws.athena.PartitionProjectionProperties.METASTORE_PROPERTY_PROJECTION_INTERVAL_UNIT_SUFFIX; -import static io.trino.plugin.hive.aws.athena.PartitionProjectionProperties.METASTORE_PROPERTY_PROJECTION_LOCATION_TEMPLATE; -import static io.trino.plugin.hive.aws.athena.PartitionProjectionProperties.PARTITION_PROJECTION_ENABLED; -import static io.trino.plugin.hive.aws.athena.PartitionProjectionProperties.PARTITION_PROJECTION_IGNORE; -import static io.trino.plugin.hive.aws.athena.PartitionProjectionProperties.PARTITION_PROJECTION_LOCATION_TEMPLATE; -import static io.trino.plugin.hive.aws.athena.PartitionProjectionProperties.PROPERTY_KEY_PREFIX; -import static io.trino.plugin.hive.aws.athena.PartitionProjectionProperties.getMetastoreProjectionPropertyKey; -import static io.trino.plugin.hive.aws.athena.projection.Projection.unsupportedProjectionColumnTypeException; -import static io.trino.spi.StandardErrorCode.INVALID_COLUMN_PROPERTY; -import static java.lang.String.format; -import static java.util.Objects.requireNonNull; -import static java.util.function.Function.identity; - -public final class PartitionProjectionService -{ - private final boolean partitionProjectionEnabled; - private final Map projectionFactories; - private final TypeManager typeManager; - - @Inject - public PartitionProjectionService( - HiveConfig hiveConfig, - Map projectionFactories, - TypeManager typeManager) - { - this.partitionProjectionEnabled = hiveConfig.isPartitionProjectionEnabled(); - this.typeManager = requireNonNull(typeManager, "typeManager is null"); - this.projectionFactories = ImmutableMap.copyOf(requireNonNull(projectionFactories, "projectionFactories is null")); - } - - public Map getPartitionProjectionTrinoTableProperties(Table table) - { - Map metastoreTableProperties = table.getParameters(); - ImmutableMap.Builder trinoTablePropertiesBuilder = ImmutableMap.builder(); - rewriteProperty(metastoreTableProperties, trinoTablePropertiesBuilder, METASTORE_PROPERTY_PROJECTION_IGNORE, PARTITION_PROJECTION_IGNORE, Boolean::valueOf); - rewriteProperty(metastoreTableProperties, trinoTablePropertiesBuilder, METASTORE_PROPERTY_PROJECTION_ENABLED, PARTITION_PROJECTION_ENABLED, Boolean::valueOf); - rewriteProperty(metastoreTableProperties, trinoTablePropertiesBuilder, METASTORE_PROPERTY_PROJECTION_LOCATION_TEMPLATE, PARTITION_PROJECTION_LOCATION_TEMPLATE, String::valueOf); - return trinoTablePropertiesBuilder.buildOrThrow(); - } - - public static Map getPartitionProjectionTrinoColumnProperties(Table table, String columnName) - { - Map metastoreTableProperties = table.getParameters(); - return rewriteColumnProjectionProperties(metastoreTableProperties, columnName); - } - - public Map getPartitionProjectionHiveTableProperties(ConnectorTableMetadata tableMetadata) - { - // If partition projection is globally disabled we don't allow defining its properties - if (!partitionProjectionEnabled && isAnyPartitionProjectionPropertyUsed(tableMetadata)) { - throw columnProjectionException("Partition projection is disabled. Enable it in configuration by setting " - + HiveConfig.CONFIGURATION_HIVE_PARTITION_PROJECTION_ENABLED + "=true"); - } - - ImmutableMap.Builder metastoreTablePropertiesBuilder = ImmutableMap.builder(); - // Handle Table Properties - Map trinoTableProperties = tableMetadata.getProperties(); - rewriteProperty( - trinoTableProperties, - metastoreTablePropertiesBuilder, - PARTITION_PROJECTION_IGNORE, - METASTORE_PROPERTY_PROJECTION_IGNORE, - value -> value.toString().toLowerCase(Locale.ENGLISH)); - rewriteProperty( - trinoTableProperties, - metastoreTablePropertiesBuilder, - PARTITION_PROJECTION_ENABLED, - METASTORE_PROPERTY_PROJECTION_ENABLED, - value -> value.toString().toLowerCase(Locale.ENGLISH)); - rewriteProperty( - trinoTableProperties, - metastoreTablePropertiesBuilder, - PARTITION_PROJECTION_LOCATION_TEMPLATE, - METASTORE_PROPERTY_PROJECTION_LOCATION_TEMPLATE, - Object::toString); - - // Handle Column Properties - tableMetadata.getColumns().stream() - .filter(columnMetadata -> !columnMetadata.getProperties().isEmpty()) - .forEach(columnMetadata -> { - Map columnProperties = columnMetadata.getProperties(); - String columnName = columnMetadata.getName(); - rewriteProperty( - columnProperties, - metastoreTablePropertiesBuilder, - COLUMN_PROJECTION_TYPE, - getMetastoreProjectionPropertyKey(columnName, COLUMN_PROJECTION_TYPE_SUFFIX), - value -> ((ProjectionType) value).name().toLowerCase(Locale.ENGLISH)); - rewriteProperty( - columnProperties, - metastoreTablePropertiesBuilder, - COLUMN_PROJECTION_VALUES, - getMetastoreProjectionPropertyKey(columnName, COLUMN_PROJECTION_VALUES_SUFFIX), - value -> Joiner.on(",").join((List) value)); - rewriteProperty( - columnProperties, - metastoreTablePropertiesBuilder, - COLUMN_PROJECTION_RANGE, - getMetastoreProjectionPropertyKey(columnName, COLUMN_PROJECTION_RANGE_SUFFIX), - value -> Joiner.on(",").join((List) value)); - rewriteProperty( - columnProperties, - metastoreTablePropertiesBuilder, - COLUMN_PROJECTION_INTERVAL, - getMetastoreProjectionPropertyKey(columnName, COLUMN_PROJECTION_INTERVAL_SUFFIX), - value -> ((Integer) value).toString()); - rewriteProperty( - columnProperties, - metastoreTablePropertiesBuilder, - COLUMN_PROJECTION_INTERVAL_UNIT, - getMetastoreProjectionPropertyKey(columnName, METASTORE_PROPERTY_PROJECTION_INTERVAL_UNIT_SUFFIX), - value -> ((ChronoUnit) value).name().toLowerCase(Locale.ENGLISH)); - rewriteProperty( - columnProperties, - metastoreTablePropertiesBuilder, - COLUMN_PROJECTION_DIGITS, - getMetastoreProjectionPropertyKey(columnName, COLUMN_PROJECTION_DIGITS_SUFFIX), - value -> ((Integer) value).toString()); - rewriteProperty( - columnProperties, - metastoreTablePropertiesBuilder, - COLUMN_PROJECTION_FORMAT, - getMetastoreProjectionPropertyKey(columnName, COLUMN_PROJECTION_FORMAT_SUFFIX), - String.class::cast); - }); - - // We initialize partition projection to validate properties. - Map metastoreTableProperties = metastoreTablePropertiesBuilder.buildOrThrow(); - List partitionColumnNames = getPartitionedBy(tableMetadata.getProperties()); - createPartitionProjection( - tableMetadata.getColumns() - .stream() - .map(ColumnMetadata::getName) - .collect(toImmutableList()), - tableMetadata.getColumns().stream() - .filter(columnMetadata -> partitionColumnNames.contains(columnMetadata.getName())) - .collect(toImmutableMap(ColumnMetadata::getName, ColumnMetadata::getType)), - metastoreTableProperties); - - return metastoreTableProperties; - } - - private boolean isAnyPartitionProjectionPropertyUsed(ConnectorTableMetadata tableMetadata) - { - if (tableMetadata.getProperties().keySet().stream() - .anyMatch(propertyKey -> propertyKey.startsWith(PROPERTY_KEY_PREFIX))) { - return true; - } - return tableMetadata.getColumns().stream() - .map(columnMetadata -> columnMetadata.getProperties().keySet()) - .flatMap(Set::stream) - .anyMatch(propertyKey -> propertyKey.startsWith(PROPERTY_KEY_PREFIX)); - } - - public Optional getPartitionProjectionFromTable(Table table) - { - if (!partitionProjectionEnabled) { - return Optional.empty(); - } - - Map tableProperties = table.getParameters(); - if (Optional.ofNullable(tableProperties.get(METASTORE_PROPERTY_PROJECTION_IGNORE)) - .map(Boolean::valueOf) - .orElse(false)) { - return Optional.empty(); - } - - return Optional.of( - createPartitionProjection( - table.getDataColumns() - .stream() - .map(Column::getName) - .collect(toImmutableList()), - table.getPartitionColumns() - .stream().collect(toImmutableMap( - Column::getName, - column -> column.getType().getType( - typeManager, - DEFAULT_PRECISION))), - tableProperties)); - } - - private PartitionProjection createPartitionProjection(List dataColumns, Map partitionColumns, Map tableProperties) - { - Optional projectionEnabledProperty = Optional.ofNullable(tableProperties.get(METASTORE_PROPERTY_PROJECTION_ENABLED)).map(Boolean::valueOf); - if (projectionEnabledProperty.orElse(false) && partitionColumns.size() < 1) { - throw columnProjectionException("Partition projection can't be enabled when no partition columns are defined."); - } - - Map columnProjections = ImmutableSet.builder() - .addAll(partitionColumns.keySet()) - .addAll(dataColumns) - .build() - .stream() - .collect(toImmutableMap( - identity(), - columnName -> rewriteColumnProjectionProperties(tableProperties, columnName))) - .entrySet() - .stream() - .filter(entry -> !entry.getValue().isEmpty()) - .collect(toImmutableMap( - Map.Entry::getKey, - entry -> { - String columnName = entry.getKey(); - if (partitionColumns.containsKey(columnName)) { - return parseColumnProjection(columnName, partitionColumns.get(columnName), entry.getValue()); - } - throw columnProjectionException("Partition projection can't be defined for non partition column: '" + columnName + "'"); - })); - - Optional storageLocationTemplate = Optional.ofNullable(tableProperties.get(METASTORE_PROPERTY_PROJECTION_LOCATION_TEMPLATE)); - if (projectionEnabledProperty.isPresent()) { - for (String columnName : partitionColumns.keySet()) { - if (!columnProjections.containsKey(columnName)) { - throw columnProjectionException("Partition projection definition for column: '" + columnName + "' missing"); - } - if (storageLocationTemplate.isPresent()) { - String locationTemplate = storageLocationTemplate.get(); - if (!locationTemplate.contains("${" + columnName + "}")) { - throw columnProjectionException(format("Partition projection location template: %s is missing partition column: '%s' placeholder", locationTemplate, columnName)); - } - } - } - } - else if (!columnProjections.isEmpty()) { - throw columnProjectionException(format( - "Columns %s projections are disallowed when partition projection property '%s' is missing", - columnProjections.keySet().stream().collect(Collectors.joining("', '", "['", "']")), - PARTITION_PROJECTION_ENABLED)); - } - - return new PartitionProjection(projectionEnabledProperty.orElse(false), storageLocationTemplate, columnProjections); - } - - private static Map rewriteColumnProjectionProperties(Map metastoreTableProperties, String columnName) - { - ImmutableMap.Builder trinoTablePropertiesBuilder = ImmutableMap.builder(); - rewriteProperty( - metastoreTableProperties, - trinoTablePropertiesBuilder, - getMetastoreProjectionPropertyKey(columnName, COLUMN_PROJECTION_TYPE_SUFFIX), - COLUMN_PROJECTION_TYPE, - value -> ProjectionType.valueOf(value.toUpperCase(Locale.ENGLISH))); - rewriteProperty( - metastoreTableProperties, - trinoTablePropertiesBuilder, - getMetastoreProjectionPropertyKey(columnName, COLUMN_PROJECTION_VALUES_SUFFIX), - COLUMN_PROJECTION_VALUES, - PartitionProjectionService::splitCommaSeparatedString); - rewriteProperty( - metastoreTableProperties, - trinoTablePropertiesBuilder, - getMetastoreProjectionPropertyKey(columnName, COLUMN_PROJECTION_RANGE_SUFFIX), - COLUMN_PROJECTION_RANGE, - PartitionProjectionService::splitCommaSeparatedString); - rewriteProperty( - metastoreTableProperties, - trinoTablePropertiesBuilder, - getMetastoreProjectionPropertyKey(columnName, COLUMN_PROJECTION_INTERVAL_SUFFIX), - COLUMN_PROJECTION_INTERVAL, - Integer::valueOf); - rewriteProperty( - metastoreTableProperties, - trinoTablePropertiesBuilder, - getMetastoreProjectionPropertyKey(columnName, METASTORE_PROPERTY_PROJECTION_INTERVAL_UNIT_SUFFIX), - COLUMN_PROJECTION_INTERVAL_UNIT, - value -> ChronoUnit.valueOf(value.toUpperCase(Locale.ENGLISH))); - rewriteProperty( - metastoreTableProperties, - trinoTablePropertiesBuilder, - getMetastoreProjectionPropertyKey(columnName, COLUMN_PROJECTION_DIGITS_SUFFIX), - COLUMN_PROJECTION_DIGITS, - Integer::valueOf); - rewriteProperty( - metastoreTableProperties, - trinoTablePropertiesBuilder, - getMetastoreProjectionPropertyKey(columnName, COLUMN_PROJECTION_FORMAT_SUFFIX), - COLUMN_PROJECTION_FORMAT, - value -> value); - return trinoTablePropertiesBuilder.buildOrThrow(); - } - - private Projection parseColumnProjection(String columnName, Type columnType, Map columnProperties) - { - ProjectionType projectionType = (ProjectionType) columnProperties.get(COLUMN_PROJECTION_TYPE); - if (Objects.isNull(projectionType)) { - throw columnProjectionException("Projection type property missing for column: '" + columnName + "'"); - } - ProjectionFactory projectionFactory = Optional.ofNullable(projectionFactories.get(projectionType)) - .orElseThrow(() -> columnProjectionException(format("Partition projection type %s for column: '%s' not supported", projectionType, columnName))); - if (!projectionFactory.isSupportedColumnType(columnType)) { - throw unsupportedProjectionColumnTypeException(columnName, columnType); - } - return projectionFactory.create(columnName, columnType, columnProperties); - } - - private static void rewriteProperty( - Map sourceProperties, - ImmutableMap.Builder targetPropertiesBuilder, - String sourcePropertyKey, - String targetPropertyKey, - Function valueMapper) - { - Optional.ofNullable(sourceProperties.get(sourcePropertyKey)) - .ifPresent(value -> targetPropertiesBuilder.put(targetPropertyKey, valueMapper.apply(value))); - } - - private TrinoException columnProjectionException(String message) - { - return new TrinoException(INVALID_COLUMN_PROPERTY, message); - } - - private static List splitCommaSeparatedString(String value) - { - return Splitter.on(',') - .trimResults() - .omitEmptyStrings() - .splitToList(value); - } -} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/aws/athena/projection/DateProjection.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/aws/athena/projection/DateProjection.java deleted file mode 100644 index 20d9408f3aa3..000000000000 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/aws/athena/projection/DateProjection.java +++ /dev/null @@ -1,114 +0,0 @@ -/* - * 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.hive.aws.athena.projection; - -import com.google.common.collect.ImmutableList; -import io.trino.spi.predicate.Domain; -import io.trino.spi.type.DateType; -import io.trino.spi.type.TimestampType; -import io.trino.spi.type.Type; -import io.trino.spi.type.VarcharType; - -import java.text.DateFormat; -import java.text.ParseException; -import java.time.Instant; -import java.time.temporal.ChronoField; -import java.time.temporal.ChronoUnit; -import java.util.Date; -import java.util.List; -import java.util.Optional; -import java.util.function.Supplier; - -import static io.airlift.slice.Slices.utf8Slice; -import static io.trino.plugin.hive.aws.athena.projection.DateProjectionFactory.UTC_TIME_ZONE_ID; -import static io.trino.spi.predicate.Domain.singleValue; -import static java.util.Objects.requireNonNull; -import static java.util.concurrent.TimeUnit.MILLISECONDS; - -public class DateProjection - extends Projection -{ - private final DateFormat dateFormat; - private final Supplier leftBound; - private final Supplier rightBound; - private final int interval; - private final ChronoUnit intervalUnit; - - public DateProjection(String columnName, DateFormat dateFormat, Supplier leftBound, Supplier rightBound, int interval, ChronoUnit intervalUnit) - { - super(columnName); - this.dateFormat = requireNonNull(dateFormat, "dateFormatPattern is null"); - this.leftBound = requireNonNull(leftBound, "leftBound is null"); - this.rightBound = requireNonNull(rightBound, "rightBound is null"); - this.interval = interval; - this.intervalUnit = requireNonNull(intervalUnit, "intervalUnit is null"); - } - - @Override - public List getProjectedValues(Optional partitionValueFilter) - { - ImmutableList.Builder builder = ImmutableList.builder(); - - Instant leftBound = adjustBoundToDateFormat(this.leftBound.get()); - Instant rightBound = adjustBoundToDateFormat(this.rightBound.get()); - - Instant currentValue = leftBound; - while (!currentValue.isAfter(rightBound)) { - String currentValueFormatted = formatValue(currentValue); - if (isValueInDomain(partitionValueFilter, currentValue, currentValueFormatted)) { - builder.add(currentValueFormatted); - } - currentValue = currentValue.atZone(UTC_TIME_ZONE_ID) - .plus(interval, intervalUnit) - .toInstant(); - } - - return builder.build(); - } - - private Instant adjustBoundToDateFormat(Instant value) - { - String formatted = formatValue(value.with(ChronoField.MILLI_OF_SECOND, 0)); - try { - return dateFormat.parse(formatted).toInstant(); - } - catch (ParseException e) { - throw invalidProjectionException(formatted, e.getMessage()); - } - } - - private String formatValue(Instant current) - { - return dateFormat.format(new Date(current.toEpochMilli())); - } - - private boolean isValueInDomain(Optional valueDomain, Instant value, String formattedValue) - { - if (valueDomain.isEmpty() || valueDomain.get().isAll()) { - return true; - } - Domain domain = valueDomain.get(); - Type type = domain.getType(); - if (type instanceof VarcharType) { - return domain.contains(singleValue(type, utf8Slice(formattedValue))); - } - if (type instanceof DateType) { - return domain.contains(singleValue(type, MILLISECONDS.toDays(value.toEpochMilli()))); - } - if (type instanceof TimestampType && ((TimestampType) type).isShort()) { - return domain.contains(singleValue(type, MILLISECONDS.toMicros(value.toEpochMilli()))); - } - throw unsupportedProjectionColumnTypeException(type); - } -} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/aws/athena/projection/EnumProjectionFactory.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/aws/athena/projection/EnumProjectionFactory.java deleted file mode 100644 index 13e897f78204..000000000000 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/aws/athena/projection/EnumProjectionFactory.java +++ /dev/null @@ -1,48 +0,0 @@ -/* - * 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.hive.aws.athena.projection; - -import io.trino.spi.type.Type; -import io.trino.spi.type.VarcharType; - -import java.util.List; -import java.util.Map; - -import static com.google.common.collect.ImmutableList.toImmutableList; -import static io.trino.plugin.hive.aws.athena.PartitionProjectionProperties.COLUMN_PROJECTION_VALUES; -import static io.trino.plugin.hive.aws.athena.PartitionProjectionProperties.getProjectionPropertyRequiredValue; - -public class EnumProjectionFactory - implements ProjectionFactory -{ - @Override - public boolean isSupportedColumnType(Type columnType) - { - return columnType instanceof VarcharType; - } - - @Override - public Projection create(String columnName, Type columnType, Map columnProperties) - { - return new EnumProjection( - columnName, - getProjectionPropertyRequiredValue( - columnName, - columnProperties, - COLUMN_PROJECTION_VALUES, - value -> ((List) value).stream() - .map(String::valueOf) - .collect(toImmutableList()))); - } -} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/aws/athena/projection/IntegerProjectionFactory.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/aws/athena/projection/IntegerProjectionFactory.java deleted file mode 100644 index 7117221c6fea..000000000000 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/aws/athena/projection/IntegerProjectionFactory.java +++ /dev/null @@ -1,66 +0,0 @@ -/* - * 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.hive.aws.athena.projection; - -import io.trino.spi.type.BigintType; -import io.trino.spi.type.IntegerType; -import io.trino.spi.type.Type; -import io.trino.spi.type.VarcharType; - -import java.util.List; -import java.util.Map; - -import static com.google.common.collect.ImmutableList.toImmutableList; -import static io.trino.plugin.hive.aws.athena.PartitionProjectionProperties.COLUMN_PROJECTION_DIGITS; -import static io.trino.plugin.hive.aws.athena.PartitionProjectionProperties.COLUMN_PROJECTION_INTERVAL; -import static io.trino.plugin.hive.aws.athena.PartitionProjectionProperties.COLUMN_PROJECTION_RANGE; -import static io.trino.plugin.hive.aws.athena.PartitionProjectionProperties.getProjectionPropertyRequiredValue; -import static io.trino.plugin.hive.aws.athena.PartitionProjectionProperties.getProjectionPropertyValue; -import static io.trino.plugin.hive.aws.athena.projection.Projection.invalidProjectionException; -import static java.lang.String.format; - -public class IntegerProjectionFactory - implements ProjectionFactory -{ - @Override - public boolean isSupportedColumnType(Type columnType) - { - return columnType instanceof VarcharType - || columnType instanceof IntegerType - || columnType instanceof BigintType; - } - - @Override - public Projection create(String columnName, Type columnType, Map columnProperties) - { - List range = getProjectionPropertyRequiredValue( - columnName, - columnProperties, - COLUMN_PROJECTION_RANGE, - value -> ((List) value).stream() - .map(element -> Integer.valueOf((String) element)) - .collect(toImmutableList())); - if (range.size() != 2) { - invalidProjectionException( - columnName, - format("Property: '%s' needs to be list of 2 integers", COLUMN_PROJECTION_RANGE)); - } - return new IntegerProjection( - columnName, - range.get(0), - range.get(1), - getProjectionPropertyValue(columnProperties, COLUMN_PROJECTION_INTERVAL, Integer.class::cast).orElse(1), - getProjectionPropertyValue(columnProperties, COLUMN_PROJECTION_DIGITS, Integer.class::cast)); - } -} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/aws/athena/projection/Projection.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/aws/athena/projection/Projection.java deleted file mode 100644 index 40066baf357f..000000000000 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/aws/athena/projection/Projection.java +++ /dev/null @@ -1,62 +0,0 @@ -/* - * 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.hive.aws.athena.projection; - -import io.trino.spi.TrinoException; -import io.trino.spi.predicate.Domain; -import io.trino.spi.type.Type; - -import java.util.List; -import java.util.Optional; - -import static io.trino.spi.StandardErrorCode.INVALID_COLUMN_PROPERTY; -import static java.lang.String.format; -import static java.util.Objects.requireNonNull; - -public abstract class Projection -{ - private final String columnName; - - public Projection(String columnName) - { - this.columnName = requireNonNull(columnName, "columnName is null"); - } - - public String getColumnName() - { - return columnName; - } - - public abstract List getProjectedValues(Optional partitionValueFilter); - - protected TrinoException unsupportedProjectionColumnTypeException(Type columnType) - { - return unsupportedProjectionColumnTypeException(columnName, columnType); - } - - public static TrinoException unsupportedProjectionColumnTypeException(String columnName, Type columnType) - { - return invalidProjectionException(columnName, "Unsupported column type: " + columnType.getDisplayName()); - } - - public static TrinoException invalidProjectionException(String columnName, String message) - { - throw new TrinoException(INVALID_COLUMN_PROPERTY, invalidProjectionMessage(columnName, message)); - } - - public static String invalidProjectionMessage(String columnName, String message) - { - return format("Column projection for column '%s' failed. %s", columnName, message); - } -} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/coercions/DoubleToVarcharCoercers.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/coercions/DoubleToVarcharCoercers.java new file mode 100644 index 000000000000..d44641812012 --- /dev/null +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/coercions/DoubleToVarcharCoercers.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.hive.coercions; + +import io.airlift.slice.Slice; +import io.airlift.slice.Slices; +import io.trino.spi.TrinoException; +import io.trino.spi.block.Block; +import io.trino.spi.block.BlockBuilder; +import io.trino.spi.type.DoubleType; +import io.trino.spi.type.VarcharType; + +import static io.airlift.slice.SliceUtf8.countCodePoints; +import static io.trino.spi.StandardErrorCode.INVALID_ARGUMENTS; +import static io.trino.spi.type.DoubleType.DOUBLE; +import static java.lang.String.format; + +public final class DoubleToVarcharCoercers +{ + private DoubleToVarcharCoercers() {} + + public static TypeCoercer createDoubleToVarcharCoercer(VarcharType toType, boolean isOrcFile) + { + return isOrcFile ? new OrcDoubleToVarcharCoercer(toType) : new DoubleToVarcharCoercer(toType); + } + + public static class DoubleToVarcharCoercer + extends TypeCoercer + { + public DoubleToVarcharCoercer(VarcharType toType) + { + super(DOUBLE, toType); + } + + @Override + protected void applyCoercedValue(BlockBuilder blockBuilder, Block block, int position) + { + writeDoubleAsSlice(DOUBLE.getDouble(block, position), blockBuilder, toType); + } + } + + public static class OrcDoubleToVarcharCoercer + extends TypeCoercer + { + public OrcDoubleToVarcharCoercer(VarcharType toType) + { + super(DOUBLE, toType); + } + + @Override + protected void applyCoercedValue(BlockBuilder blockBuilder, Block block, int position) + { + double doubleValue = DOUBLE.getDouble(block, position); + if (Double.isNaN(doubleValue)) { + blockBuilder.appendNull(); + return; + } + writeDoubleAsSlice(doubleValue, blockBuilder, toType); + } + } + + private static void writeDoubleAsSlice(double value, BlockBuilder blockBuilder, VarcharType varcharType) + { + Slice converted = Slices.utf8Slice(Double.toString(value)); + if (!varcharType.isUnbounded() && countCodePoints(converted) > varcharType.getBoundedLength()) { + throw new TrinoException(INVALID_ARGUMENTS, format("Varchar representation of %s exceeds %s bounds", value, varcharType)); + } + varcharType.writeSlice(blockBuilder, converted); + } +} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/coercions/FloatToVarcharCoercers.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/coercions/FloatToVarcharCoercers.java new file mode 100644 index 000000000000..570a3a228c00 --- /dev/null +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/coercions/FloatToVarcharCoercers.java @@ -0,0 +1,82 @@ +/* + * 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.hive.coercions; + +import io.airlift.slice.Slice; +import io.airlift.slice.Slices; +import io.trino.spi.TrinoException; +import io.trino.spi.block.Block; +import io.trino.spi.block.BlockBuilder; +import io.trino.spi.type.RealType; +import io.trino.spi.type.VarcharType; + +import static io.airlift.slice.SliceUtf8.countCodePoints; +import static io.trino.spi.StandardErrorCode.INVALID_ARGUMENTS; +import static io.trino.spi.type.RealType.REAL; +import static java.lang.String.format; + +public final class FloatToVarcharCoercers +{ + private FloatToVarcharCoercers() {} + + public static TypeCoercer createFloatToVarcharCoercer(VarcharType toType, boolean isOrcFile) + { + return isOrcFile ? new OrcFloatToVarcharCoercer(toType) : new FloatToVarcharCoercer(toType); + } + + public static class FloatToVarcharCoercer + extends TypeCoercer + { + public FloatToVarcharCoercer(VarcharType toType) + { + super(REAL, toType); + } + + @Override + protected void applyCoercedValue(BlockBuilder blockBuilder, Block block, int position) + { + writeFloatAsSlice(REAL.getFloat(block, position), blockBuilder, toType); + } + } + + public static class OrcFloatToVarcharCoercer + extends TypeCoercer + { + public OrcFloatToVarcharCoercer(VarcharType toType) + { + super(REAL, toType); + } + + @Override + protected void applyCoercedValue(BlockBuilder blockBuilder, Block block, int position) + { + float floatValue = REAL.getFloat(block, position); + + if (Float.isNaN(floatValue)) { + blockBuilder.appendNull(); + return; + } + writeFloatAsSlice(floatValue, blockBuilder, toType); + } + } + + private static void writeFloatAsSlice(float value, BlockBuilder blockBuilder, VarcharType varcharType) + { + Slice converted = Slices.utf8Slice(Float.toString(value)); + if (!varcharType.isUnbounded() && countCodePoints(converted) > varcharType.getBoundedLength()) { + throw new TrinoException(INVALID_ARGUMENTS, format("Varchar representation of %s exceeds %s bounds", value, varcharType)); + } + varcharType.writeSlice(blockBuilder, converted); + } +} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/coercions/IntegerNumberToDoubleCoercer.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/coercions/IntegerNumberToDoubleCoercer.java new file mode 100644 index 000000000000..9577b41bbd3d --- /dev/null +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/coercions/IntegerNumberToDoubleCoercer.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.hive.coercions; + +import io.trino.spi.TrinoException; +import io.trino.spi.block.Block; +import io.trino.spi.block.BlockBuilder; +import io.trino.spi.type.DoubleType; +import io.trino.spi.type.Type; + +import static io.trino.spi.StandardErrorCode.INVALID_CAST_ARGUMENT; +import static io.trino.spi.type.DoubleType.DOUBLE; + +public class IntegerNumberToDoubleCoercer + extends TypeCoercer +{ + private static final long MIN_EXACT_DOUBLE = -(1L << 52); // -2^52 + private static final long MAX_EXACT_DOUBLE = (1L << 52) - 1; // 2^52 - 1 + + public IntegerNumberToDoubleCoercer(F fromType) + { + super(fromType, DOUBLE); + } + + @Override + protected void applyCoercedValue(BlockBuilder blockBuilder, Block block, int position) + { + long value = fromType.getLong(block, position); + // IEEE 754 double-precision can guarantee this for up to 53 bits (52 bits of significand + the implicit leading 1 bit) + // https://stackoverflow.com/questions/43655668/are-all-integer-values-perfectly-represented-as-doubles + if (overflow(value)) { + throw new TrinoException(INVALID_CAST_ARGUMENT, "Cannot read value '%s' as DOUBLE".formatted(value)); + } + DOUBLE.writeDouble(blockBuilder, value); + } + + private static boolean overflow(long value) + { + return value < MIN_EXACT_DOUBLE || value > MAX_EXACT_DOUBLE; + } +} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/fs/CachingDirectoryLister.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/fs/CachingDirectoryLister.java index 2c212394888b..00cad2f2290a 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/fs/CachingDirectoryLister.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/fs/CachingDirectoryLister.java @@ -13,7 +13,6 @@ */ package io.trino.plugin.hive.fs; -import com.google.common.annotations.VisibleForTesting; import com.google.common.cache.Cache; import com.google.common.cache.Weigher; import com.google.common.collect.ImmutableList; @@ -21,6 +20,7 @@ import io.airlift.units.DataSize; import io.airlift.units.Duration; import io.trino.cache.EvictableCacheBuilder; +import io.trino.filesystem.FileEntry; import io.trino.filesystem.Location; import io.trino.filesystem.TrinoFileSystem; import io.trino.plugin.hive.HiveConfig; @@ -36,6 +36,7 @@ import java.util.List; import java.util.Optional; import java.util.concurrent.TimeUnit; +import java.util.function.Predicate; import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.collect.ImmutableList.toImmutableList; @@ -53,14 +54,15 @@ public class CachingDirectoryLister // to deal more efficiently with cache invalidation scenarios for partitioned tables. private final Cache cache; private final List tablePrefixes; + private final Predicate filterPredicate; @Inject public CachingDirectoryLister(HiveConfig hiveClientConfig) { - this(hiveClientConfig.getFileStatusCacheExpireAfterWrite(), hiveClientConfig.getFileStatusCacheMaxRetainedSize(), hiveClientConfig.getFileStatusCacheTables()); + this(hiveClientConfig.getFileStatusCacheExpireAfterWrite(), hiveClientConfig.getFileStatusCacheMaxRetainedSize(), hiveClientConfig.getFileStatusCacheTables(), hiveClientConfig.getS3StorageClassFilter().toFileEntryPredicate()); } - public CachingDirectoryLister(Duration expireAfterWrite, DataSize maxSize, List tables) + public CachingDirectoryLister(Duration expireAfterWrite, DataSize maxSize, List tables, Predicate filterPredicate) { this.cache = EvictableCacheBuilder.newBuilder() .maximumWeight(maxSize.toBytes()) @@ -72,6 +74,7 @@ public CachingDirectoryLister(Duration expireAfterWrite, DataSize maxSize, List< this.tablePrefixes = tables.stream() .map(CachingDirectoryLister::parseTableName) .collect(toImmutableList()); + this.filterPredicate = filterPredicate; } private static SchemaTablePrefix parseTableName(String tableName) @@ -205,8 +208,8 @@ public long getRequestCount() return cache.stats().requestCount(); } - @VisibleForTesting - boolean isCached(Location location) + @Override + public boolean isCached(Location location) { ValueHolder cached = cache.getIfPresent(location); return cached != null && cached.getFiles().isPresent(); diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/fs/HiveFileIterator.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/fs/HiveFileIterator.java index 098b75183086..f17789b40ee4 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/fs/HiveFileIterator.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/fs/HiveFileIterator.java @@ -15,10 +15,8 @@ import com.google.common.annotations.VisibleForTesting; import com.google.common.collect.AbstractIterator; -import io.airlift.stats.TimeStat; import io.trino.filesystem.Location; import io.trino.filesystem.TrinoFileSystem; -import io.trino.hdfs.HdfsNamenodeStats; import io.trino.plugin.hive.metastore.Table; import io.trino.spi.TrinoException; @@ -43,11 +41,7 @@ public enum NestedDirectoryPolicy FAIL } - private final Table table; private final Location location; - private final TrinoFileSystem fileSystem; - private final DirectoryLister directoryLister; - private final HdfsNamenodeStats namenodeStats; private final NestedDirectoryPolicy nestedDirectoryPolicy; private final Iterator remoteIterator; @@ -56,23 +50,18 @@ public HiveFileIterator( Location location, TrinoFileSystem fileSystem, DirectoryLister directoryLister, - HdfsNamenodeStats namenodeStats, NestedDirectoryPolicy nestedDirectoryPolicy) { - this.table = requireNonNull(table, "table is null"); this.location = requireNonNull(location, "location is null"); - this.fileSystem = requireNonNull(fileSystem, "fileSystem is null"); - this.directoryLister = requireNonNull(directoryLister, "directoryLister is null"); - this.namenodeStats = requireNonNull(namenodeStats, "namenodeStats is null"); this.nestedDirectoryPolicy = requireNonNull(nestedDirectoryPolicy, "nestedDirectoryPolicy is null"); - this.remoteIterator = getLocatedFileStatusRemoteIterator(location); + this.remoteIterator = new FileStatusIterator(table, location, fileSystem, directoryLister, nestedDirectoryPolicy); } @Override protected TrinoFileStatus computeNext() { while (remoteIterator.hasNext()) { - TrinoFileStatus status = getLocatedFileStatus(remoteIterator); + TrinoFileStatus status = remoteIterator.next(); // Ignore hidden files and directories if (nestedDirectoryPolicy == RECURSE) { @@ -91,23 +80,6 @@ else if (isHiddenFileOrDirectory(Location.of(status.getPath()))) { return endOfData(); } - private Iterator getLocatedFileStatusRemoteIterator(Location location) - { - try (TimeStat.BlockTimer ignored = namenodeStats.getListLocatedStatus().time()) { - return new FileStatusIterator(table, location, fileSystem, directoryLister, namenodeStats, nestedDirectoryPolicy); - } - catch (IOException e) { - throw new TrinoException(HIVE_FILESYSTEM_ERROR, "Failed to list files for location: " + location, e); - } - } - - private TrinoFileStatus getLocatedFileStatus(Iterator iterator) - { - try (TimeStat.BlockTimer ignored = namenodeStats.getRemoteIteratorNext().time()) { - return iterator.next(); - } - } - @VisibleForTesting static boolean isHiddenFileOrDirectory(Location location) { @@ -148,7 +120,6 @@ private static class FileStatusIterator implements Iterator { private final Location location; - private final HdfsNamenodeStats namenodeStats; private final RemoteIterator fileStatusIterator; private FileStatusIterator( @@ -156,12 +127,9 @@ private FileStatusIterator( Location location, TrinoFileSystem fileSystem, DirectoryLister directoryLister, - HdfsNamenodeStats namenodeStats, NestedDirectoryPolicy nestedDirectoryPolicy) - throws IOException { - this.location = location; - this.namenodeStats = namenodeStats; + this.location = requireNonNull(location, "location is null"); try { if (nestedDirectoryPolicy == RECURSE) { this.fileStatusIterator = directoryLister.listFilesRecursively(fileSystem, table, location); @@ -202,7 +170,6 @@ public TrinoFileStatus next() private TrinoException processException(IOException exception) { - namenodeStats.getRemoteIteratorNext().recordException(exception); if (exception instanceof FileNotFoundException) { return new TrinoException(HIVE_FILE_NOT_FOUND, "Partition location does not exist: " + location); } diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/fs/TransactionScopeCachingDirectoryLister.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/fs/TransactionScopeCachingDirectoryLister.java index 1f83b0ada254..3b7bc11f1eda 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/fs/TransactionScopeCachingDirectoryLister.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/fs/TransactionScopeCachingDirectoryLister.java @@ -157,10 +157,10 @@ public TrinoFileStatus next() }; } - @VisibleForTesting - boolean isCached(Location location) + @Override + public boolean isCached(Location location) { - return isCached(new TransactionDirectoryListingCacheKey(transactionId, location)); + return isCached(new TransactionDirectoryListingCacheKey(transactionId, location)) || delegate.isCached(location); } @VisibleForTesting diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/line/CsvFileWriterFactory.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/line/CsvFileWriterFactory.java index 94d500376de1..3453ec78827a 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/line/CsvFileWriterFactory.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/line/CsvFileWriterFactory.java @@ -17,7 +17,6 @@ import io.trino.filesystem.TrinoFileSystemFactory; import io.trino.hive.formats.line.csv.CsvSerializerFactory; import io.trino.hive.formats.line.text.TextLineWriterFactory; -import io.trino.plugin.hive.HiveSessionProperties; import io.trino.spi.type.TypeManager; public class CsvFileWriterFactory @@ -30,7 +29,6 @@ public CsvFileWriterFactory(TrinoFileSystemFactory trinoFileSystemFactory, TypeM typeManager, new CsvSerializerFactory(), new TextLineWriterFactory(), - HiveSessionProperties::isCsvNativeWriterEnabled, true); } } diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/line/CsvPageSourceFactory.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/line/CsvPageSourceFactory.java index b3a54bc35b07..1990dc670c47 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/line/CsvPageSourceFactory.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/line/CsvPageSourceFactory.java @@ -17,9 +17,7 @@ import io.trino.filesystem.TrinoFileSystemFactory; import io.trino.hive.formats.line.csv.CsvDeserializerFactory; import io.trino.hive.formats.line.text.TextLineReaderFactory; -import io.trino.plugin.hive.FileFormatDataSourceStats; import io.trino.plugin.hive.HiveConfig; -import io.trino.plugin.hive.HiveSessionProperties; import static java.lang.Math.toIntExact; @@ -27,12 +25,10 @@ public class CsvPageSourceFactory extends LinePageSourceFactory { @Inject - public CsvPageSourceFactory(TrinoFileSystemFactory trinoFileSystemFactory, FileFormatDataSourceStats stats, HiveConfig config) + public CsvPageSourceFactory(TrinoFileSystemFactory trinoFileSystemFactory, HiveConfig config) { super(trinoFileSystemFactory, - stats, new CsvDeserializerFactory(), - new TextLineReaderFactory(1024, 1024, toIntExact(config.getTextMaxLineLength().toBytes())), - HiveSessionProperties::isCsvNativeReaderEnabled); + new TextLineReaderFactory(1024, 1024, toIntExact(config.getTextMaxLineLength().toBytes()))); } } diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/line/JsonFileWriterFactory.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/line/JsonFileWriterFactory.java index 77b32379c6d4..c5bfdb309afd 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/line/JsonFileWriterFactory.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/line/JsonFileWriterFactory.java @@ -17,7 +17,6 @@ import io.trino.filesystem.TrinoFileSystemFactory; import io.trino.hive.formats.line.json.JsonSerializerFactory; import io.trino.hive.formats.line.text.TextLineWriterFactory; -import io.trino.plugin.hive.HiveSessionProperties; import io.trino.spi.type.TypeManager; public class JsonFileWriterFactory @@ -30,7 +29,6 @@ public JsonFileWriterFactory(TrinoFileSystemFactory trinoFileSystemFactory, Type typeManager, new JsonSerializerFactory(), new TextLineWriterFactory(), - HiveSessionProperties::isJsonNativeWriterEnabled, false); } } diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/line/JsonPageSourceFactory.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/line/JsonPageSourceFactory.java index aeeeef9e1364..7f9e794ab1ab 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/line/JsonPageSourceFactory.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/line/JsonPageSourceFactory.java @@ -17,9 +17,7 @@ import io.trino.filesystem.TrinoFileSystemFactory; import io.trino.hive.formats.line.json.JsonDeserializerFactory; import io.trino.hive.formats.line.text.TextLineReaderFactory; -import io.trino.plugin.hive.FileFormatDataSourceStats; import io.trino.plugin.hive.HiveConfig; -import io.trino.plugin.hive.HiveSessionProperties; import static java.lang.Math.toIntExact; @@ -27,12 +25,10 @@ public class JsonPageSourceFactory extends LinePageSourceFactory { @Inject - public JsonPageSourceFactory(TrinoFileSystemFactory trinoFileSystemFactory, FileFormatDataSourceStats stats, HiveConfig config) + public JsonPageSourceFactory(TrinoFileSystemFactory trinoFileSystemFactory, HiveConfig config) { super(trinoFileSystemFactory, - stats, new JsonDeserializerFactory(), - new TextLineReaderFactory(1024, 1024, toIntExact(config.getTextMaxLineLength().toBytes())), - HiveSessionProperties::isJsonNativeReaderEnabled); + new TextLineReaderFactory(1024, 1024, toIntExact(config.getTextMaxLineLength().toBytes()))); } } diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/line/LineFileWriterFactory.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/line/LineFileWriterFactory.java index 68abafb1891c..29bfe9839018 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/line/LineFileWriterFactory.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/line/LineFileWriterFactory.java @@ -41,14 +41,12 @@ import java.io.IOException; import java.io.OutputStream; import java.util.List; +import java.util.Map; import java.util.Optional; import java.util.OptionalInt; -import java.util.Properties; -import java.util.function.Predicate; import java.util.stream.IntStream; import static com.google.common.collect.ImmutableList.toImmutableList; -import static com.google.common.collect.Maps.fromProperties; import static io.trino.memory.context.AggregatedMemoryContext.newSimpleAggregatedMemoryContext; import static io.trino.plugin.hive.HiveErrorCode.HIVE_UNSUPPORTED_FORMAT; import static io.trino.plugin.hive.HiveErrorCode.HIVE_WRITER_OPEN_ERROR; @@ -64,7 +62,6 @@ public abstract class LineFileWriterFactory { private final TrinoFileSystemFactory fileSystemFactory; private final TypeManager typeManager; - private final Predicate activation; private final LineSerializerFactory lineSerializerFactory; private final LineWriterFactory lineWriterFactory; private final boolean headerSupported; @@ -74,12 +71,10 @@ protected LineFileWriterFactory( TypeManager typeManager, LineSerializerFactory lineSerializerFactory, LineWriterFactory lineWriterFactory, - Predicate activation, boolean headerSupported) { this.fileSystemFactory = requireNonNull(fileSystemFactory, "fileSystemFactory is null"); this.typeManager = requireNonNull(typeManager, "typeManager is null"); - this.activation = requireNonNull(activation, "activation is null"); this.lineSerializerFactory = requireNonNull(lineSerializerFactory, "lineSerializerFactory is null"); this.lineWriterFactory = requireNonNull(lineWriterFactory, "lineWriterFactory is null"); this.headerSupported = headerSupported; @@ -91,7 +86,7 @@ public Optional createFileWriter( List inputColumnNames, StorageFormat storageFormat, HiveCompressionCodec compressionCodec, - Properties schema, + Map schema, ConnectorSession session, OptionalInt bucketNumber, AcidTransaction transaction, @@ -99,8 +94,7 @@ public Optional createFileWriter( WriterKind writerKind) { if (!lineWriterFactory.getHiveOutputFormatClassName().equals(storageFormat.getOutputFormat()) || - !lineSerializerFactory.getHiveSerDeClassNames().contains(storageFormat.getSerde()) || - !activation.test(session)) { + !lineSerializerFactory.getHiveSerDeClassNames().contains(storageFormat.getSerde())) { return Optional.empty(); } @@ -119,7 +113,7 @@ public Optional createFileWriter( .mapToObj(ordinal -> new Column(fileColumnNames.get(ordinal), fileColumnTypes.get(ordinal), ordinal)) .toList(); - LineSerializer lineSerializer = lineSerializerFactory.create(columns, fromProperties(schema)); + LineSerializer lineSerializer = lineSerializerFactory.create(columns, schema); try { TrinoFileSystem fileSystem = fileSystemFactory.create(session.getIdentity()); @@ -146,10 +140,10 @@ public Optional createFileWriter( } } - private Optional getFileHeader(Properties schema, List columns) + private Optional getFileHeader(Map schema, List columns) throws IOException { - String skipHeaderCount = schema.getProperty(SKIP_HEADER_COUNT_KEY, "0"); + String skipHeaderCount = schema.getOrDefault(SKIP_HEADER_COUNT_KEY, "0"); if (skipHeaderCount.equals("0")) { return Optional.empty(); } @@ -162,7 +156,7 @@ private Optional getFileHeader(Properties schema, List columns) columns.stream() .map(column -> new Column(column.name(), VARCHAR, column.ordinal())) .collect(toImmutableList()), - fromProperties(schema)); + schema); PageBuilder pageBuilder = new PageBuilder(headerSerializer.getTypes()); pageBuilder.declarePosition(); diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/line/LinePageSourceFactory.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/line/LinePageSourceFactory.java index af3e7dc738a2..ec9dd3af853a 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/line/LinePageSourceFactory.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/line/LinePageSourceFactory.java @@ -13,7 +13,6 @@ */ package io.trino.plugin.hive.line; -import com.google.common.collect.Maps; import io.airlift.slice.Slices; import io.airlift.units.DataSize; import io.airlift.units.DataSize.Unit; @@ -28,13 +27,12 @@ import io.trino.hive.formats.line.LineReader; import io.trino.hive.formats.line.LineReaderFactory; import io.trino.plugin.hive.AcidInfo; -import io.trino.plugin.hive.FileFormatDataSourceStats; import io.trino.plugin.hive.HiveColumnHandle; import io.trino.plugin.hive.HivePageSourceFactory; import io.trino.plugin.hive.ReaderColumns; import io.trino.plugin.hive.ReaderPageSource; +import io.trino.plugin.hive.Schema; import io.trino.plugin.hive.acid.AcidTransaction; -import io.trino.plugin.hive.fs.MonitoredInputFile; import io.trino.spi.TrinoException; import io.trino.spi.connector.ConnectorSession; import io.trino.spi.connector.EmptyPageSource; @@ -44,8 +42,6 @@ import java.util.List; import java.util.Optional; import java.util.OptionalInt; -import java.util.Properties; -import java.util.function.Predicate; import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.collect.ImmutableList.toImmutableList; @@ -54,7 +50,6 @@ import static io.trino.plugin.hive.HiveErrorCode.HIVE_CANNOT_OPEN_SPLIT; import static io.trino.plugin.hive.HivePageSourceProvider.projectBaseColumns; import static io.trino.plugin.hive.ReaderPageSource.noProjectionAdaptation; -import static io.trino.plugin.hive.util.HiveUtil.getDeserializerClassName; import static io.trino.plugin.hive.util.HiveUtil.getFooterCount; import static io.trino.plugin.hive.util.HiveUtil.getHeaderCount; import static io.trino.plugin.hive.util.HiveUtil.splitError; @@ -67,22 +62,16 @@ public abstract class LinePageSourceFactory private static final DataSize SMALL_FILE_SIZE = DataSize.of(8, Unit.MEGABYTE); private final TrinoFileSystemFactory fileSystemFactory; - private final FileFormatDataSourceStats stats; private final LineDeserializerFactory lineDeserializerFactory; private final LineReaderFactory lineReaderFactory; - private final Predicate activation; protected LinePageSourceFactory( TrinoFileSystemFactory fileSystemFactory, - FileFormatDataSourceStats stats, LineDeserializerFactory lineDeserializerFactory, - LineReaderFactory lineReaderFactory, - Predicate activation) + LineReaderFactory lineReaderFactory) { this.fileSystemFactory = requireNonNull(fileSystemFactory, "fileSystemFactory is null"); - this.stats = requireNonNull(stats, "stats is null"); this.lineDeserializerFactory = requireNonNull(lineDeserializerFactory, "lineDeserializerFactory is null"); - this.activation = requireNonNull(activation, "activation is null"); this.lineReaderFactory = requireNonNull(lineReaderFactory, "lineReaderFactory is null"); } @@ -93,7 +82,8 @@ public Optional createPageSource( long start, long length, long estimatedFileSize, - Properties schema, + long fileModifiedTime, + Schema schema, List columns, TupleDomain effectivePredicate, Optional acidInfo, @@ -101,20 +91,19 @@ public Optional createPageSource( boolean originalFile, AcidTransaction transaction) { - if (!lineReaderFactory.getHiveOutputFormatClassName().equals(schema.getProperty(FILE_INPUT_FORMAT)) || - !lineDeserializerFactory.getHiveSerDeClassNames().contains(getDeserializerClassName(schema)) || - !activation.test(session)) { + if (!lineReaderFactory.getHiveInputFormatClassNames().contains(schema.serdeProperties().get(FILE_INPUT_FORMAT)) || + !lineDeserializerFactory.getHiveSerDeClassNames().contains(schema.serializationLibraryName())) { return Optional.empty(); } checkArgument(acidInfo.isEmpty(), "Acid is not supported"); // get header and footer count - int headerCount = getHeaderCount(schema); + int headerCount = getHeaderCount(schema.serdeProperties()); if (headerCount > 1) { checkArgument(start == 0, "Multiple header rows are not supported for a split file"); } - int footerCount = getFooterCount(schema); + int footerCount = getFooterCount(schema.serdeProperties()); if (footerCount > 0) { checkArgument(start == 0, "Footer not supported for a split file"); } @@ -135,12 +124,12 @@ public Optional createPageSource( projectedReaderColumns.stream() .map(column -> new Column(column.getName(), column.getType(), column.getBaseHiveColumnIndex())) .collect(toImmutableList()), - Maps.fromProperties(schema)); + schema.serdeProperties()); } // buffer file if small TrinoFileSystem trinoFileSystem = fileSystemFactory.create(session.getIdentity()); - TrinoInputFile inputFile = new MonitoredInputFile(stats, trinoFileSystem.newInputFile(path)); + TrinoInputFile inputFile = trinoFileSystem.newInputFile(path); try { length = min(inputFile.length() - start, length); if (!inputFile.exists()) { diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/line/OpenXJsonFileWriterFactory.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/line/OpenXJsonFileWriterFactory.java index 68291237bd1b..e4097455d475 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/line/OpenXJsonFileWriterFactory.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/line/OpenXJsonFileWriterFactory.java @@ -17,7 +17,6 @@ import io.trino.filesystem.TrinoFileSystemFactory; import io.trino.hive.formats.line.openxjson.OpenXJsonSerializerFactory; import io.trino.hive.formats.line.text.TextLineWriterFactory; -import io.trino.plugin.hive.HiveSessionProperties; import io.trino.spi.type.TypeManager; public class OpenXJsonFileWriterFactory @@ -30,7 +29,6 @@ public OpenXJsonFileWriterFactory(TrinoFileSystemFactory trinoFileSystemFactory, typeManager, new OpenXJsonSerializerFactory(), new TextLineWriterFactory(), - HiveSessionProperties::isOpenXJsonNativeWriterEnabled, true); } } diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/line/OpenXJsonPageSourceFactory.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/line/OpenXJsonPageSourceFactory.java index 3b6c135cec51..1598ef6f0942 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/line/OpenXJsonPageSourceFactory.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/line/OpenXJsonPageSourceFactory.java @@ -17,9 +17,7 @@ import io.trino.filesystem.TrinoFileSystemFactory; import io.trino.hive.formats.line.openxjson.OpenXJsonDeserializerFactory; import io.trino.hive.formats.line.text.TextLineReaderFactory; -import io.trino.plugin.hive.FileFormatDataSourceStats; import io.trino.plugin.hive.HiveConfig; -import io.trino.plugin.hive.HiveSessionProperties; import static java.lang.Math.toIntExact; @@ -27,12 +25,10 @@ public class OpenXJsonPageSourceFactory extends LinePageSourceFactory { @Inject - public OpenXJsonPageSourceFactory(TrinoFileSystemFactory trinoFileSystemFactory, FileFormatDataSourceStats stats, HiveConfig config) + public OpenXJsonPageSourceFactory(TrinoFileSystemFactory trinoFileSystemFactory, HiveConfig config) { super(trinoFileSystemFactory, - stats, new OpenXJsonDeserializerFactory(), - new TextLineReaderFactory(1024, 1024, toIntExact(config.getTextMaxLineLength().toBytes())), - HiveSessionProperties::isOpenXJsonNativeReaderEnabled); + new TextLineReaderFactory(1024, 1024, toIntExact(config.getTextMaxLineLength().toBytes()))); } } diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/line/RegexFileWriterFactory.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/line/RegexFileWriterFactory.java index ae3bc7d2a984..6aa8cc6517c0 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/line/RegexFileWriterFactory.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/line/RegexFileWriterFactory.java @@ -24,9 +24,9 @@ import io.trino.spi.connector.ConnectorSession; import java.util.List; +import java.util.Map; import java.util.Optional; import java.util.OptionalInt; -import java.util.Properties; import static io.trino.plugin.hive.HiveErrorCode.HIVE_WRITER_OPEN_ERROR; import static io.trino.plugin.hive.util.HiveClassNames.REGEX_SERDE_CLASS; @@ -40,7 +40,7 @@ public Optional createFileWriter( List inputColumnNames, StorageFormat storageFormat, HiveCompressionCodec compressionCodec, - Properties schema, + Map schema, ConnectorSession session, OptionalInt bucketNumber, AcidTransaction transaction, diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/line/RegexPageSourceFactory.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/line/RegexPageSourceFactory.java index 9408b679b1ee..5af670e37210 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/line/RegexPageSourceFactory.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/line/RegexPageSourceFactory.java @@ -17,9 +17,7 @@ import io.trino.filesystem.TrinoFileSystemFactory; import io.trino.hive.formats.line.regex.RegexDeserializerFactory; import io.trino.hive.formats.line.text.TextLineReaderFactory; -import io.trino.plugin.hive.FileFormatDataSourceStats; import io.trino.plugin.hive.HiveConfig; -import io.trino.plugin.hive.HiveSessionProperties; import static java.lang.Math.toIntExact; @@ -27,12 +25,10 @@ public class RegexPageSourceFactory extends LinePageSourceFactory { @Inject - public RegexPageSourceFactory(TrinoFileSystemFactory trinoFileSystemFactory, FileFormatDataSourceStats stats, HiveConfig config) + public RegexPageSourceFactory(TrinoFileSystemFactory trinoFileSystemFactory, HiveConfig config) { super(trinoFileSystemFactory, - stats, new RegexDeserializerFactory(), - new TextLineReaderFactory(1024, 1024, toIntExact(config.getTextMaxLineLength().toBytes())), - HiveSessionProperties::isRegexNativeReaderEnabled); + new TextLineReaderFactory(1024, 1024, toIntExact(config.getTextMaxLineLength().toBytes()))); } } diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/line/SimpleSequenceFilePageSourceFactory.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/line/SimpleSequenceFilePageSourceFactory.java index 9c682159e593..f77a91b408b8 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/line/SimpleSequenceFilePageSourceFactory.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/line/SimpleSequenceFilePageSourceFactory.java @@ -17,9 +17,7 @@ import io.trino.filesystem.TrinoFileSystemFactory; import io.trino.hive.formats.line.sequence.SequenceFileReaderFactory; import io.trino.hive.formats.line.simple.SimpleDeserializerFactory; -import io.trino.plugin.hive.FileFormatDataSourceStats; import io.trino.plugin.hive.HiveConfig; -import io.trino.plugin.hive.HiveSessionProperties; import static java.lang.Math.toIntExact; @@ -27,12 +25,10 @@ public class SimpleSequenceFilePageSourceFactory extends LinePageSourceFactory { @Inject - public SimpleSequenceFilePageSourceFactory(TrinoFileSystemFactory trinoFileSystemFactory, FileFormatDataSourceStats stats, HiveConfig config) + public SimpleSequenceFilePageSourceFactory(TrinoFileSystemFactory trinoFileSystemFactory, HiveConfig config) { super(trinoFileSystemFactory, - stats, new SimpleDeserializerFactory(), - new SequenceFileReaderFactory(1024, toIntExact(config.getTextMaxLineLength().toBytes())), - HiveSessionProperties::isSequenceFileNativeReaderEnabled); + new SequenceFileReaderFactory(1024, toIntExact(config.getTextMaxLineLength().toBytes()))); } } diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/line/SimpleSequenceFileWriterFactory.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/line/SimpleSequenceFileWriterFactory.java index 77bc24f43a31..291a3c4c0edf 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/line/SimpleSequenceFileWriterFactory.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/line/SimpleSequenceFileWriterFactory.java @@ -17,7 +17,6 @@ import io.trino.filesystem.TrinoFileSystemFactory; import io.trino.hive.formats.line.sequence.SequenceFileWriterFactory; import io.trino.hive.formats.line.simple.SimpleSerializerFactory; -import io.trino.plugin.hive.HiveSessionProperties; import io.trino.plugin.hive.NodeVersion; import io.trino.spi.type.TypeManager; @@ -31,7 +30,6 @@ public SimpleSequenceFileWriterFactory(TrinoFileSystemFactory trinoFileSystemFac typeManager, new SimpleSerializerFactory(), new SequenceFileWriterFactory(nodeVersion.toString()), - HiveSessionProperties::isSequenceFileNativeWriterEnabled, false); } } diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/line/SimpleTextFilePageSourceFactory.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/line/SimpleTextFilePageSourceFactory.java index 08d285d9cf77..e283b9668c8d 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/line/SimpleTextFilePageSourceFactory.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/line/SimpleTextFilePageSourceFactory.java @@ -17,9 +17,7 @@ import io.trino.filesystem.TrinoFileSystemFactory; import io.trino.hive.formats.line.simple.SimpleDeserializerFactory; import io.trino.hive.formats.line.text.TextLineReaderFactory; -import io.trino.plugin.hive.FileFormatDataSourceStats; import io.trino.plugin.hive.HiveConfig; -import io.trino.plugin.hive.HiveSessionProperties; import static java.lang.Math.toIntExact; @@ -27,12 +25,10 @@ public class SimpleTextFilePageSourceFactory extends LinePageSourceFactory { @Inject - public SimpleTextFilePageSourceFactory(TrinoFileSystemFactory trinoFileSystemFactory, FileFormatDataSourceStats stats, HiveConfig config) + public SimpleTextFilePageSourceFactory(TrinoFileSystemFactory trinoFileSystemFactory, HiveConfig config) { super(trinoFileSystemFactory, - stats, new SimpleDeserializerFactory(), - new TextLineReaderFactory(1024, 1024, toIntExact(config.getTextMaxLineLength().toBytes())), - HiveSessionProperties::isTextFileNativeReaderEnabled); + new TextLineReaderFactory(1024, 1024, toIntExact(config.getTextMaxLineLength().toBytes()))); } } diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/line/SimpleTextFileWriterFactory.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/line/SimpleTextFileWriterFactory.java index 68ebad92ef85..523eb8a35246 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/line/SimpleTextFileWriterFactory.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/line/SimpleTextFileWriterFactory.java @@ -17,7 +17,6 @@ import io.trino.filesystem.TrinoFileSystemFactory; import io.trino.hive.formats.line.simple.SimpleSerializerFactory; import io.trino.hive.formats.line.text.TextLineWriterFactory; -import io.trino.plugin.hive.HiveSessionProperties; import io.trino.spi.type.TypeManager; public class SimpleTextFileWriterFactory @@ -30,7 +29,6 @@ public SimpleTextFileWriterFactory(TrinoFileSystemFactory trinoFileSystemFactory typeManager, new SimpleSerializerFactory(), new TextLineWriterFactory(), - HiveSessionProperties::isTextFileNativeWriterEnabled, false); } } diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/CachingHiveMetastoreModule.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/CachingHiveMetastoreModule.java new file mode 100644 index 000000000000..a44963f8389e --- /dev/null +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/CachingHiveMetastoreModule.java @@ -0,0 +1,66 @@ +/* + * 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.hive.metastore; + +import com.google.inject.Binder; +import com.google.inject.Provides; +import com.google.inject.Scopes; +import com.google.inject.Singleton; +import io.airlift.configuration.AbstractConfigurationAwareModule; +import io.trino.plugin.hive.metastore.cache.CachingHiveMetastore; +import io.trino.plugin.hive.metastore.cache.CachingHiveMetastoreConfig; +import io.trino.plugin.hive.metastore.cache.ImpersonationCachingConfig; +import io.trino.plugin.hive.metastore.cache.SharedHiveMetastoreCache; +import io.trino.plugin.hive.metastore.cache.SharedHiveMetastoreCache.CachingHiveMetastoreFactory; + +import java.util.Optional; + +import static io.airlift.configuration.ConfigBinder.configBinder; +import static org.weakref.jmx.guice.ExportBinder.newExporter; + +public class CachingHiveMetastoreModule + extends AbstractConfigurationAwareModule +{ + @Override + protected void setup(Binder binder) + { + configBinder(binder).bindConfig(CachingHiveMetastoreConfig.class); + // TODO this should only be bound when impersonation is actually enabled + configBinder(binder).bindConfig(ImpersonationCachingConfig.class); + binder.bind(SharedHiveMetastoreCache.class).in(Scopes.SINGLETON); + // export under the old name, for backwards compatibility + newExporter(binder).export(HiveMetastoreFactory.class) + .as(generator -> generator.generatedNameOf(CachingHiveMetastore.class)); + } + + @Provides + @Singleton + public static HiveMetastoreFactory createHiveMetastore( + @RawHiveMetastoreFactory HiveMetastoreFactory metastoreFactory, + SharedHiveMetastoreCache sharedHiveMetastoreCache) + { + // cross TX metastore cache is enabled wrapper with caching metastore + return sharedHiveMetastoreCache.createCachingHiveMetastoreFactory(metastoreFactory); + } + + @Provides + @Singleton + public static Optional createHiveMetastore(HiveMetastoreFactory metastoreFactory) + { + if (metastoreFactory instanceof CachingHiveMetastoreFactory) { + return Optional.of(((CachingHiveMetastoreFactory) metastoreFactory).getMetastore()); + } + return Optional.empty(); + } +} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/Column.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/Column.java index af1440c121c5..b07d6200bd5f 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/Column.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/Column.java @@ -111,4 +111,14 @@ public int hashCode() { return Objects.hash(name, type, comment, properties); } + + public Column withName(String newColumnName) + { + return new Column(newColumnName, type, comment, properties); + } + + public Column withComment(Optional comment) + { + return new Column(name, type, comment, properties); + } } diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/CoralSemiTransactionalHiveMSCAdapter.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/CoralSemiTransactionalHiveMSCAdapter.java index c432e3663f56..3306dac70bfd 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/CoralSemiTransactionalHiveMSCAdapter.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/CoralSemiTransactionalHiveMSCAdapter.java @@ -67,7 +67,9 @@ public com.linkedin.coral.hive.metastore.api.Database getDatabase(String dbName) @Override public List getAllTables(String dbName) { - return delegate.getAllTables(dbName); + return delegate.getTables(dbName).stream() + .map(tableInfo -> tableInfo.tableName().getTableName()) + .toList(); } @Override diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/ForwardingHiveMetastore.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/ForwardingHiveMetastore.java index f08be97fa7ac..1a8708e16a32 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/ForwardingHiveMetastore.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/ForwardingHiveMetastore.java @@ -13,25 +13,21 @@ */ package io.trino.plugin.hive.metastore; -import io.trino.hive.thrift.metastore.DataOperationType; -import io.trino.plugin.hive.HiveColumnStatisticType; +import com.google.common.collect.ImmutableSet; import io.trino.plugin.hive.HivePartition; import io.trino.plugin.hive.HiveType; import io.trino.plugin.hive.PartitionStatistics; import io.trino.plugin.hive.acid.AcidOperation; -import io.trino.plugin.hive.acid.AcidTransaction; import io.trino.plugin.hive.metastore.HivePrivilegeInfo.HivePrivilege; import io.trino.spi.connector.SchemaTableName; import io.trino.spi.predicate.TupleDomain; import io.trino.spi.security.RoleGrant; -import io.trino.spi.type.Type; import java.util.List; import java.util.Map; import java.util.Optional; import java.util.OptionalLong; import java.util.Set; -import java.util.function.Function; import static java.util.Objects.requireNonNull; @@ -64,78 +60,39 @@ public Optional
getTable(String databaseName, String tableName) } @Override - public Set getSupportedColumnStatistics(Type type) + public Map getTableColumnStatistics(String databaseName, String tableName, Set columnNames) { - return delegate.getSupportedColumnStatistics(type); + return delegate.getTableColumnStatistics(databaseName, tableName, columnNames); } @Override - public PartitionStatistics getTableStatistics(Table table) + public Map> getPartitionColumnStatistics(String databaseName, String tableName, Set partitionNames, Set columnNames) { - return delegate.getTableStatistics(table); + return delegate.getPartitionColumnStatistics(databaseName, tableName, partitionNames, columnNames); } @Override - public Map getPartitionStatistics(Table table, List partitions) + public void updateTableStatistics(String databaseName, String tableName, OptionalLong acidWriteId, StatisticsUpdateMode mode, PartitionStatistics statisticsUpdate) { - return delegate.getPartitionStatistics(table, partitions); + delegate.updateTableStatistics(databaseName, tableName, acidWriteId, mode, statisticsUpdate); } @Override - public void updateTableStatistics( - String databaseName, - String tableName, - AcidTransaction transaction, - Function update) - { - delegate.updateTableStatistics(databaseName, tableName, transaction, update); - } - - @Override - public void updatePartitionStatistics( - Table table, - String partitionName, - Function update) - { - delegate.updatePartitionStatistics(table, partitionName, update); - } - - @Override - public void updatePartitionStatistics( - Table table, - Map> updates) - { - delegate.updatePartitionStatistics(table, updates); - } - - @Override - public List getAllTables(String databaseName) - { - return delegate.getAllTables(databaseName); - } - - @Override - public Optional> getAllTables() - { - return delegate.getAllTables(); - } - - @Override - public List getTablesWithParameter(String databaseName, String parameterKey, String parameterValue) + public void updatePartitionStatistics(Table table, StatisticsUpdateMode mode, Map partitionUpdates) { - return delegate.getTablesWithParameter(databaseName, parameterKey, parameterValue); + delegate.updatePartitionStatistics(table, mode, partitionUpdates); } @Override - public List getAllViews(String databaseName) + public List getTables(String databaseName) { - return delegate.getAllViews(databaseName); + return delegate.getTables(databaseName); } @Override - public Optional> getAllViews() + public List getTableNamesWithParameters(String databaseName, String parameterKey, ImmutableSet parameterValues) { - return delegate.getAllViews(); + return delegate.getTableNamesWithParameters(databaseName, parameterKey, parameterValues); } @Override @@ -179,9 +136,10 @@ public void replaceTable( String databaseName, String tableName, Table newTable, - PrincipalPrivileges principalPrivileges) + PrincipalPrivileges principalPrivileges, + Map environmentContext) { - delegate.replaceTable(databaseName, tableName, newTable, principalPrivileges); + delegate.replaceTable(databaseName, tableName, newTable, principalPrivileges, environmentContext); } @Override @@ -313,12 +271,6 @@ public void revokeRoles(Set roles, Set grantees, boolean delegate.revokeRoles(roles, grantees, adminOption, grantor); } - @Override - public Set listGrantedPrincipals(String role) - { - return delegate.listGrantedPrincipals(role); - } - @Override public Set listRoleGrants(HivePrincipal principal) { @@ -410,19 +362,6 @@ public long allocateWriteId(String dbName, String tableName, long transactionId) return delegate.allocateWriteId(dbName, tableName, transactionId); } - @Override - public void acquireTableWriteLock( - AcidTransactionOwner transactionOwner, - String queryId, - long transactionId, - String dbName, - String tableName, - DataOperationType operation, - boolean isDynamicPartitionWrite) - { - delegate.acquireTableWriteLock(transactionOwner, queryId, transactionId, dbName, tableName, operation, isDynamicPartitionWrite); - } - @Override public void updateTableWriteId( String dbName, @@ -434,16 +373,6 @@ public void updateTableWriteId( delegate.updateTableWriteId(dbName, tableName, transactionId, writeId, rowCountChange); } - @Override - public void alterPartitions( - String dbName, - String tableName, - List partitions, - long writeId) - { - delegate.alterPartitions(dbName, tableName, partitions, writeId); - } - @Override public void addDynamicPartitions( String dbName, diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/HiveColumnStatistics.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/HiveColumnStatistics.java index c68c19f17413..76c1996a2231 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/HiveColumnStatistics.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/HiveColumnStatistics.java @@ -41,9 +41,9 @@ public class HiveColumnStatistics private final Optional dateStatistics; private final Optional booleanStatistics; private final OptionalLong maxValueSizeInBytes; - private final OptionalLong totalSizeInBytes; + private final OptionalDouble averageColumnLength; private final OptionalLong nullsCount; - private final OptionalLong distinctValuesCount; + private final OptionalLong distinctValuesWithNullCount; public static HiveColumnStatistics empty() { @@ -58,9 +58,9 @@ public HiveColumnStatistics( @JsonProperty("dateStatistics") Optional dateStatistics, @JsonProperty("booleanStatistics") Optional booleanStatistics, @JsonProperty("maxValueSizeInBytes") OptionalLong maxValueSizeInBytes, - @JsonProperty("totalSizeInBytes") OptionalLong totalSizeInBytes, + @JsonProperty("averageColumnLength") OptionalDouble averageColumnLength, @JsonProperty("nullsCount") OptionalLong nullsCount, - @JsonProperty("distinctValuesCount") OptionalLong distinctValuesCount) + @JsonProperty("distinctValuesWithNullCount") OptionalLong distinctValuesWithNullCount) { this.integerStatistics = requireNonNull(integerStatistics, "integerStatistics is null"); this.doubleStatistics = requireNonNull(doubleStatistics, "doubleStatistics is null"); @@ -68,9 +68,9 @@ public HiveColumnStatistics( this.dateStatistics = requireNonNull(dateStatistics, "dateStatistics is null"); this.booleanStatistics = requireNonNull(booleanStatistics, "booleanStatistics is null"); this.maxValueSizeInBytes = requireNonNull(maxValueSizeInBytes, "maxValueSizeInBytes is null"); - this.totalSizeInBytes = requireNonNull(totalSizeInBytes, "totalSizeInBytes is null"); + this.averageColumnLength = requireNonNull(averageColumnLength, "averageColumnLength is null"); this.nullsCount = requireNonNull(nullsCount, "nullsCount is null"); - this.distinctValuesCount = requireNonNull(distinctValuesCount, "distinctValuesCount is null"); + this.distinctValuesWithNullCount = requireNonNull(distinctValuesWithNullCount, "distinctValuesWithNullCount is null"); List presentStatistics = new ArrayList<>(); integerStatistics.ifPresent(s -> presentStatistics.add("integerStatistics")); @@ -118,9 +118,9 @@ public OptionalLong getMaxValueSizeInBytes() } @JsonProperty - public OptionalLong getTotalSizeInBytes() + public OptionalDouble getAverageColumnLength() { - return totalSizeInBytes; + return averageColumnLength; } @JsonProperty @@ -130,9 +130,9 @@ public OptionalLong getNullsCount() } @JsonProperty - public OptionalLong getDistinctValuesCount() + public OptionalLong getDistinctValuesWithNullCount() { - return distinctValuesCount; + return distinctValuesWithNullCount; } @Override @@ -151,9 +151,9 @@ public boolean equals(Object o) Objects.equals(dateStatistics, that.dateStatistics) && Objects.equals(booleanStatistics, that.booleanStatistics) && Objects.equals(maxValueSizeInBytes, that.maxValueSizeInBytes) && - Objects.equals(totalSizeInBytes, that.totalSizeInBytes) && + Objects.equals(averageColumnLength, that.averageColumnLength) && Objects.equals(nullsCount, that.nullsCount) && - Objects.equals(distinctValuesCount, that.distinctValuesCount); + Objects.equals(distinctValuesWithNullCount, that.distinctValuesWithNullCount); } @Override @@ -166,9 +166,9 @@ public int hashCode() dateStatistics, booleanStatistics, maxValueSizeInBytes, - totalSizeInBytes, + averageColumnLength, nullsCount, - distinctValuesCount); + distinctValuesWithNullCount); } @Override @@ -181,27 +181,27 @@ public String toString() .add("dateStatistics", dateStatistics) .add("booleanStatistics", booleanStatistics) .add("maxValueSizeInBytes", maxValueSizeInBytes) - .add("totalSizeInBytes", totalSizeInBytes) + .add("averageColumnLength", averageColumnLength) .add("nullsCount", nullsCount) - .add("distinctValuesCount", distinctValuesCount) + .add("distinctValuesWithNullCount", distinctValuesWithNullCount) .toString(); } - public static HiveColumnStatistics createIntegerColumnStatistics(OptionalLong min, OptionalLong max, OptionalLong nullsCount, OptionalLong distinctValuesCount) + public static HiveColumnStatistics createIntegerColumnStatistics(OptionalLong min, OptionalLong max, OptionalLong nullsCount, OptionalLong distinctValuesWithNullCount) { return builder() .setIntegerStatistics(new IntegerStatistics(min, max)) .setNullsCount(nullsCount) - .setDistinctValuesCount(distinctValuesCount) + .setDistinctValuesWithNullCount(distinctValuesWithNullCount) .build(); } - public static HiveColumnStatistics createDoubleColumnStatistics(OptionalDouble min, OptionalDouble max, OptionalLong nullsCount, OptionalLong distinctValuesCount) + public static HiveColumnStatistics createDoubleColumnStatistics(OptionalDouble min, OptionalDouble max, OptionalLong nullsCount, OptionalLong distinctValuesWithNullCount) { return builder() .setDoubleStatistics(new DoubleStatistics(min, max)) .setNullsCount(nullsCount) - .setDistinctValuesCount(distinctValuesCount) + .setDistinctValuesWithNullCount(distinctValuesWithNullCount) .build(); } @@ -210,16 +210,16 @@ public static HiveColumnStatistics createDecimalColumnStatistics(Optional min, Optional max, OptionalLong nullsCount, OptionalLong distinctValuesCount) + public static HiveColumnStatistics createDateColumnStatistics(Optional min, Optional max, OptionalLong nullsCount, OptionalLong distinctValuesWithNullCount) { return builder() .setDateStatistics(new DateStatistics(min, max)) .setNullsCount(nullsCount) - .setDistinctValuesCount(distinctValuesCount) + .setDistinctValuesWithNullCount(distinctValuesWithNullCount) .build(); } @@ -233,23 +233,23 @@ public static HiveColumnStatistics createBooleanColumnStatistics(OptionalLong tr public static HiveColumnStatistics createStringColumnStatistics( OptionalLong maxValueSizeInBytes, - OptionalLong totalSizeInBytes, + OptionalDouble averageColumnLength, OptionalLong nullsCount, - OptionalLong distinctValuesCount) + OptionalLong distinctValuesWithNullCount) { return builder() .setMaxValueSizeInBytes(maxValueSizeInBytes) - .setTotalSizeInBytes(totalSizeInBytes) + .setAverageColumnLength(averageColumnLength) .setNullsCount(nullsCount) - .setDistinctValuesCount(distinctValuesCount) + .setDistinctValuesWithNullCount(distinctValuesWithNullCount) .build(); } - public static HiveColumnStatistics createBinaryColumnStatistics(OptionalLong maxValueSizeInBytes, OptionalLong totalSizeInBytes, OptionalLong nullsCount) + public static HiveColumnStatistics createBinaryColumnStatistics(OptionalLong maxValueSizeInBytes, OptionalDouble averageColumnLength, OptionalLong nullsCount) { return builder() .setMaxValueSizeInBytes(maxValueSizeInBytes) - .setTotalSizeInBytes(totalSizeInBytes) + .setAverageColumnLength(averageColumnLength) .setNullsCount(nullsCount) .build(); } @@ -272,9 +272,9 @@ public static class Builder private Optional dateStatistics = Optional.empty(); private Optional booleanStatistics = Optional.empty(); private OptionalLong maxValueSizeInBytes = OptionalLong.empty(); - private OptionalLong totalSizeInBytes = OptionalLong.empty(); + private OptionalDouble averageColumnLength = OptionalDouble.empty(); private OptionalLong nullsCount = OptionalLong.empty(); - private OptionalLong distinctValuesCount = OptionalLong.empty(); + private OptionalLong distinctValuesWithNullCount = OptionalLong.empty(); private Builder() {} @@ -286,9 +286,9 @@ private Builder(HiveColumnStatistics other) this.dateStatistics = other.getDateStatistics(); this.booleanStatistics = other.getBooleanStatistics(); this.maxValueSizeInBytes = other.getMaxValueSizeInBytes(); - this.totalSizeInBytes = other.getTotalSizeInBytes(); + this.averageColumnLength = other.getAverageColumnLength(); this.nullsCount = other.getNullsCount(); - this.distinctValuesCount = other.getDistinctValuesCount(); + this.distinctValuesWithNullCount = other.getDistinctValuesWithNullCount(); } public Builder setIntegerStatistics(Optional integerStatistics) @@ -363,15 +363,15 @@ public Builder setMaxValueSizeInBytes(OptionalLong maxValueSizeInBytes) return this; } - public Builder setTotalSizeInBytes(long totalSizeInBytes) + public Builder setAverageColumnLength(double averageColumnLength) { - this.totalSizeInBytes = OptionalLong.of(totalSizeInBytes); + this.averageColumnLength = OptionalDouble.of(averageColumnLength); return this; } - public Builder setTotalSizeInBytes(OptionalLong totalSizeInBytes) + public Builder setAverageColumnLength(OptionalDouble averageColumnLength) { - this.totalSizeInBytes = totalSizeInBytes; + this.averageColumnLength = averageColumnLength; return this; } @@ -387,15 +387,15 @@ public Builder setNullsCount(long nullsCount) return this; } - public Builder setDistinctValuesCount(OptionalLong distinctValuesCount) + public Builder setDistinctValuesWithNullCount(OptionalLong distinctValuesWithNullCount) { - this.distinctValuesCount = distinctValuesCount; + this.distinctValuesWithNullCount = distinctValuesWithNullCount; return this; } - public Builder setDistinctValuesCount(long distinctValuesCount) + public Builder setDistinctValuesWithNullCount(long distinctValuesWithNullCount) { - this.distinctValuesCount = OptionalLong.of(distinctValuesCount); + this.distinctValuesWithNullCount = OptionalLong.of(distinctValuesWithNullCount); return this; } @@ -408,9 +408,9 @@ public HiveColumnStatistics build() dateStatistics, booleanStatistics, maxValueSizeInBytes, - totalSizeInBytes, + averageColumnLength, nullsCount, - distinctValuesCount); + distinctValuesWithNullCount); } } } diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/HiveMetastore.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/HiveMetastore.java index 55c23e5fc1e8..cc38362ef9a0 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/HiveMetastore.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/HiveMetastore.java @@ -13,27 +13,22 @@ */ package io.trino.plugin.hive.metastore; -import com.google.common.collect.ImmutableMap; -import io.trino.hive.thrift.metastore.DataOperationType; -import io.trino.plugin.hive.HiveColumnStatisticType; +import com.google.common.collect.ImmutableSet; import io.trino.plugin.hive.HivePartition; import io.trino.plugin.hive.HiveType; import io.trino.plugin.hive.PartitionStatistics; import io.trino.plugin.hive.acid.AcidOperation; -import io.trino.plugin.hive.acid.AcidTransaction; import io.trino.plugin.hive.metastore.HivePrivilegeInfo.HivePrivilege; import io.trino.spi.TrinoException; import io.trino.spi.connector.SchemaTableName; import io.trino.spi.predicate.TupleDomain; import io.trino.spi.security.RoleGrant; -import io.trino.spi.type.Type; import java.util.List; import java.util.Map; import java.util.Optional; import java.util.OptionalLong; import java.util.Set; -import java.util.function.Function; import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED; @@ -45,36 +40,40 @@ public interface HiveMetastore Optional
getTable(String databaseName, String tableName); - Set getSupportedColumnStatistics(Type type); - - PartitionStatistics getTableStatistics(Table table); - - Map getPartitionStatistics(Table table, List partitions); + /** + * @param columnNames Must not be empty. + */ + Map getTableColumnStatistics(String databaseName, String tableName, Set columnNames); - void updateTableStatistics(String databaseName, String tableName, AcidTransaction transaction, Function update); + /** + * @param columnNames Must not be empty. + */ + Map> getPartitionColumnStatistics( + String databaseName, + String tableName, + Set partitionNames, + Set columnNames); - default void updatePartitionStatistics(Table table, String partitionName, Function update) + /** + * If true, callers should inspect table and partition parameters for spark stats. + * This method really only exists for the ThriftHiveMetastore implementation. Spark mixes table and column statistics into the table parameters, and this breaks + * the abstractions of the metastore interface. + */ + default boolean useSparkTableStatistics() { - updatePartitionStatistics(table, ImmutableMap.of(partitionName, update)); + return false; } - void updatePartitionStatistics(Table table, Map> updates); + void updateTableStatistics(String databaseName, String tableName, OptionalLong acidWriteId, StatisticsUpdateMode mode, PartitionStatistics statisticsUpdate); - List getAllTables(String databaseName); + void updatePartitionStatistics(Table table, StatisticsUpdateMode mode, Map partitionUpdates); - /** - * @return List of tables, views and materialized views names from all schemas or Optional.empty if operation is not supported - */ - Optional> getAllTables(); - - List getTablesWithParameter(String databaseName, String parameterKey, String parameterValue); - - List getAllViews(String databaseName); + List getTables(String databaseName); /** - * @return List of views including materialized views names from all schemas or Optional.empty if operation is not supported + * @param parameterValues is using ImmutableSet to mark that this api does not support filtering by null parameter value. */ - Optional> getAllViews(); + List getTableNamesWithParameters(String databaseName, String parameterKey, ImmutableSet parameterValues); void createDatabase(Database database); @@ -89,11 +88,11 @@ default void updatePartitionStatistics(Table table, String partitionName, Functi void dropTable(String databaseName, String tableName, boolean deleteData); /** - * This should only be used if the semantic here is drop and add. Trying to + * This should only be used if the semantic here is to drop and add. Trying to * alter one field of a table object previously acquired from getTable is * probably not what you want. */ - void replaceTable(String databaseName, String tableName, Table newTable, PrincipalPrivileges principalPrivileges); + void replaceTable(String databaseName, String tableName, Table newTable, PrincipalPrivileges principalPrivileges, Map environmentContext); void renameTable(String databaseName, String tableName, String newDatabaseName, String newTableName); @@ -118,7 +117,7 @@ default void updatePartitionStatistics(Table table, String partitionName, Functi * @param tableName the name of the table * @param columnNames the list of partition column names * @param partitionKeysFilter optional filter for the partition column values - * @return a list of partition names as created by {@link MetastoreUtil#toPartitionName} + * @return a list of partition names as created by {@code MetastoreUtil#toPartitionName} * @see TupleDomain */ Optional> getPartitionNamesByFilter(String databaseName, String tableName, List columnNames, TupleDomain partitionKeysFilter); @@ -141,8 +140,6 @@ default void updatePartitionStatistics(Table table, String partitionName, Functi void revokeRoles(Set roles, Set grantees, boolean adminOption, HivePrincipal grantor); - Set listGrantedPrincipals(String role); - Set listRoleGrants(HivePrincipal principal); void grantTablePrivileges(String databaseName, String tableName, String tableOwner, HivePrincipal grantee, HivePrincipal grantor, Set privileges, boolean grantOption); @@ -210,7 +207,7 @@ default void acquireTableWriteLock( long transactionId, String dbName, String tableName, - DataOperationType operation, + AcidOperation operation, boolean isDynamicPartitionWrite) { throw new UnsupportedOperationException(); @@ -221,11 +218,6 @@ default void updateTableWriteId(String dbName, String tableName, long transactio throw new UnsupportedOperationException(); } - default void alterPartitions(String dbName, String tableName, List partitions, long writeId) - { - throw new UnsupportedOperationException(); - } - default void addDynamicPartitions(String dbName, String tableName, List partitionNames, long transactionId, long writeId, AcidOperation operation) { throw new UnsupportedOperationException(); diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/HiveMetastoreFactory.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/HiveMetastoreFactory.java index 070e30367e77..5ba1a1558987 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/HiveMetastoreFactory.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/HiveMetastoreFactory.java @@ -21,6 +21,11 @@ public interface HiveMetastoreFactory { + default boolean hasBuiltInCaching() + { + return false; + } + boolean isImpersonationEnabled(); /** diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/HiveMetastoreModule.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/HiveMetastoreModule.java index 7735f2cb2bb7..c0b923b16911 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/HiveMetastoreModule.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/HiveMetastoreModule.java @@ -16,11 +16,8 @@ import com.google.inject.Binder; import com.google.inject.Key; import com.google.inject.Module; -import com.google.inject.Provides; -import com.google.inject.Singleton; import io.airlift.configuration.AbstractConfigurationAwareModule; import io.trino.plugin.hive.AllowHiveTableRename; -import io.trino.plugin.hive.HideDeltaLakeTables; import io.trino.plugin.hive.metastore.file.FileMetastoreModule; import io.trino.plugin.hive.metastore.glue.GlueMetastoreModule; import io.trino.plugin.hive.metastore.thrift.ThriftMetastoreModule; @@ -62,12 +59,4 @@ private void bindMetastoreModule(String name, Module module) metastore -> name.equalsIgnoreCase(metastore.getMetastoreType()), module)); } - - @HideDeltaLakeTables - @Singleton - @Provides - public boolean hideDeltaLakeTables(HiveMetastoreConfig hiveMetastoreConfig) - { - return hiveMetastoreConfig.isHideDeltaLakeTables(); - } } diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/HivePageSinkMetadataProvider.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/HivePageSinkMetadataProvider.java index dc1f14d0a9b7..3bc240af4154 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/HivePageSinkMetadataProvider.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/HivePageSinkMetadataProvider.java @@ -13,7 +13,6 @@ */ package io.trino.plugin.hive.metastore; -import io.trino.plugin.hive.HiveMetastoreClosure; import io.trino.spi.connector.SchemaTableName; import java.util.List; @@ -25,12 +24,12 @@ public class HivePageSinkMetadataProvider { - private final HiveMetastoreClosure delegate; + private final HiveMetastore delegate; private final SchemaTableName schemaTableName; private final Optional
table; private final Map, Optional> modifiedPartitions; - public HivePageSinkMetadataProvider(HivePageSinkMetadata pageSinkMetadata, HiveMetastoreClosure delegate) + public HivePageSinkMetadataProvider(HivePageSinkMetadata pageSinkMetadata, HiveMetastore delegate) { requireNonNull(pageSinkMetadata, "pageSinkMetadata is null"); this.delegate = delegate; @@ -52,7 +51,7 @@ public Optional getPartition(List partitionValues) } Optional modifiedPartition = modifiedPartitions.get(partitionValues); if (modifiedPartition == null) { - return delegate.getPartition(schemaTableName.getSchemaName(), schemaTableName.getTableName(), partitionValues); + return delegate.getPartition(table.get(), partitionValues); } return modifiedPartition; } diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/HiveTransaction.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/HiveTransaction.java index 5bf83100a1b1..d2605281bb5c 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/HiveTransaction.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/HiveTransaction.java @@ -14,7 +14,6 @@ package io.trino.plugin.hive.metastore; import com.google.common.collect.ImmutableList; -import io.trino.plugin.hive.HiveMetastoreClosure; import io.trino.plugin.hive.HivePartition; import io.trino.plugin.hive.HiveTableHandle; import io.trino.plugin.hive.acid.AcidTransaction; @@ -62,7 +61,7 @@ public AcidTransaction getTransaction() public ValidTxnWriteIdList getValidWriteIds( AcidTransactionOwner transactionOwner, - HiveMetastoreClosure metastore, + HiveMetastore metastore, HiveTableHandle tableHandle) { List lockedTables; diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/MetastoreMethod.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/MetastoreMethod.java new file mode 100644 index 000000000000..44b38dd63bda --- /dev/null +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/MetastoreMethod.java @@ -0,0 +1,50 @@ +/* + * 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.hive.metastore; + +import static com.google.common.base.CaseFormat.LOWER_CAMEL; +import static com.google.common.base.CaseFormat.UPPER_UNDERSCORE; + +public enum MetastoreMethod +{ + CREATE_DATABASE, + DROP_DATABASE, + CREATE_TABLE, + GET_ALL_DATABASES, + GET_DATABASE, + GET_TABLE, + GET_TABLES, + GET_ALL_TABLES, + GET_ALL_TABLES_FROM_DATABASE, + GET_TABLE_WITH_PARAMETER, + GET_TABLE_STATISTICS, + GET_ALL_VIEWS, + GET_ALL_VIEWS_FROM_DATABASE, + UPDATE_TABLE_STATISTICS, + ADD_PARTITIONS, + GET_PARTITION_NAMES_BY_FILTER, + GET_PARTITIONS_BY_NAMES, + GET_PARTITION, + GET_PARTITION_STATISTICS, + GET_PARTITION_COLUMN_STATISTICS, + UPDATE_PARTITION_STATISTICS, + REPLACE_TABLE, + DROP_TABLE, + /**/; + + public static MetastoreMethod fromMethodName(String name) + { + return valueOf(LOWER_CAMEL.to(UPPER_UNDERSCORE, name)); + } +} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/MetastoreUtil.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/MetastoreUtil.java index 405796542872..17944c61d897 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/MetastoreUtil.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/MetastoreUtil.java @@ -17,8 +17,10 @@ import com.google.common.base.VerifyException; import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableMultimap; +import com.google.common.collect.ImmutableSet; import com.google.common.primitives.Longs; import io.airlift.slice.Slice; +import io.trino.plugin.hive.HiveBasicStatistics; import io.trino.plugin.hive.HiveColumnHandle; import io.trino.plugin.hive.PartitionOfflineException; import io.trino.plugin.hive.TableOfflineException; @@ -42,6 +44,7 @@ import io.trino.spi.type.TinyintType; import io.trino.spi.type.Type; import io.trino.spi.type.VarcharType; +import jakarta.annotation.Nullable; import org.joda.time.format.DateTimeFormatter; import org.joda.time.format.ISODateTimeFormat; @@ -51,7 +54,8 @@ import java.util.Map; import java.util.Objects; import java.util.Optional; -import java.util.Properties; +import java.util.OptionalLong; +import java.util.Set; import java.util.concurrent.TimeUnit; import static com.google.common.base.Preconditions.checkArgument; @@ -72,7 +76,6 @@ import static io.trino.plugin.hive.HiveMetadata.AVRO_SCHEMA_URL_KEY; import static io.trino.plugin.hive.HiveSplitManager.PRESTO_OFFLINE; import static io.trino.plugin.hive.HiveStorageFormat.AVRO; -import static io.trino.plugin.hive.metastore.thrift.ThriftMetastoreUtil.NUM_ROWS; import static io.trino.plugin.hive.util.HiveClassNames.AVRO_SERDE_CLASS; import static io.trino.plugin.hive.util.HiveUtil.makePartName; import static io.trino.plugin.hive.util.SerdeConstants.SERIALIZATION_LIB; @@ -87,10 +90,15 @@ public final class MetastoreUtil { private static final String HIVE_PARTITION_VALUE_WILDCARD = ""; + public static final String NUM_ROWS = "numRows"; + public static final String NUM_FILES = "numFiles"; + public static final String RAW_DATA_SIZE = "rawDataSize"; + public static final String TOTAL_SIZE = "totalSize"; + public static final Set STATS_PROPERTIES = ImmutableSet.of(NUM_FILES, NUM_ROWS, RAW_DATA_SIZE, TOTAL_SIZE); private MetastoreUtil() {} - public static Properties getHiveSchema(Table table) + public static Map getHiveSchema(Table table) { // Mimics function in Hive: MetaStoreUtils.getTableMetadata(Table) return getHiveSchema( @@ -103,7 +111,7 @@ public static Properties getHiveSchema(Table table) table.getPartitionColumns()); } - public static Properties getHiveSchema(Partition partition, Table table) + public static Map getHiveSchema(Partition partition, Table table) { // Mimics function in Hive: MetaStoreUtils.getSchema(Partition, Table) return getHiveSchema( @@ -116,7 +124,7 @@ public static Properties getHiveSchema(Partition partition, Table table) table.getPartitionColumns()); } - private static Properties getHiveSchema( + private static Map getHiveSchema( Storage sd, Optional tableSd, List tableDataColumns, @@ -128,33 +136,33 @@ private static Properties getHiveSchema( // Mimics function in Hive: // MetaStoreUtils.getSchema(StorageDescriptor, StorageDescriptor, Map, String, String, List) - Properties schema = new Properties(); + ImmutableMap.Builder schema = ImmutableMap.builder(); - schema.setProperty(FILE_INPUT_FORMAT, sd.getStorageFormat().getInputFormat()); - schema.setProperty(FILE_OUTPUT_FORMAT, sd.getStorageFormat().getOutputFormat()); + schema.put(FILE_INPUT_FORMAT, sd.getStorageFormat().getInputFormat()); + schema.put(FILE_OUTPUT_FORMAT, sd.getStorageFormat().getOutputFormat()); - schema.setProperty(META_TABLE_NAME, databaseName + "." + tableName); - schema.setProperty(META_TABLE_LOCATION, sd.getLocation()); + schema.put(META_TABLE_NAME, databaseName + "." + tableName); + schema.put(META_TABLE_LOCATION, sd.getLocation()); if (sd.getBucketProperty().isPresent()) { - schema.setProperty(BUCKET_FIELD_NAME, Joiner.on(",").join(sd.getBucketProperty().get().getBucketedBy())); - schema.setProperty(BUCKET_COUNT, Integer.toString(sd.getBucketProperty().get().getBucketCount())); + schema.put(BUCKET_FIELD_NAME, Joiner.on(",").join(sd.getBucketProperty().get().getBucketedBy())); + schema.put(BUCKET_COUNT, Integer.toString(sd.getBucketProperty().get().getBucketCount())); } else { - schema.setProperty(BUCKET_COUNT, "0"); + schema.put(BUCKET_COUNT, "0"); } for (Map.Entry param : sd.getSerdeParameters().entrySet()) { - schema.setProperty(param.getKey(), (param.getValue() != null) ? param.getValue() : ""); + schema.put(param.getKey(), (param.getValue() != null) ? param.getValue() : ""); } if (sd.getStorageFormat().getSerde().equals(AVRO_SERDE_CLASS) && tableSd.isPresent()) { for (Map.Entry param : tableSd.get().getSerdeParameters().entrySet()) { - schema.setProperty(param.getKey(), nullToEmpty(param.getValue())); + schema.put(param.getKey(), nullToEmpty(param.getValue())); } } - schema.setProperty(SERIALIZATION_LIB, sd.getStorageFormat().getSerde()); + schema.put(SERIALIZATION_LIB, sd.getStorageFormat().getSerde()); StringBuilder columnNameBuilder = new StringBuilder(); StringBuilder columnTypeBuilder = new StringBuilder(); @@ -173,9 +181,9 @@ private static Properties getHiveSchema( } String columnNames = columnNameBuilder.toString(); String columnTypes = columnTypeBuilder.toString(); - schema.setProperty(META_TABLE_COLUMNS, columnNames); - schema.setProperty(META_TABLE_COLUMN_TYPES, columnTypes); - schema.setProperty("columns.comments", columnCommentBuilder.toString()); + schema.put(META_TABLE_COLUMNS, columnNames); + schema.put(META_TABLE_COLUMN_TYPES, columnTypes); + schema.put("columns.comments", columnCommentBuilder.toString()); StringBuilder partString = new StringBuilder(); String partStringSep = ""; @@ -192,20 +200,20 @@ private static Properties getHiveSchema( } } if (partString.length() > 0) { - schema.setProperty(META_TABLE_PARTITION_COLUMNS, partString.toString()); - schema.setProperty(META_TABLE_PARTITION_COLUMN_TYPES, partTypesString.toString()); + schema.put(META_TABLE_PARTITION_COLUMNS, partString.toString()); + schema.put(META_TABLE_PARTITION_COLUMN_TYPES, partTypesString.toString()); } if (parameters != null) { for (Map.Entry entry : parameters.entrySet()) { // add non-null parameters to the schema if (entry.getValue() != null) { - schema.setProperty(entry.getKey(), entry.getValue()); + schema.put(entry.getKey(), entry.getValue()); } } } - return schema; + return schema.buildKeepingLast(); } public static ProtectMode getProtectMode(Partition partition) @@ -429,4 +437,43 @@ public static Map adjustRowCount(Map parameters, copiedParameters.put(NUM_ROWS, String.valueOf(newRowCount)); return ImmutableMap.copyOf(copiedParameters); } + + public static HiveBasicStatistics getHiveBasicStatistics(Map parameters) + { + OptionalLong numFiles = toLong(parameters.get(NUM_FILES)); + OptionalLong numRows = toLong(parameters.get(NUM_ROWS)); + OptionalLong inMemoryDataSizeInBytes = toLong(parameters.get(RAW_DATA_SIZE)); + OptionalLong onDiskDataSizeInBytes = toLong(parameters.get(TOTAL_SIZE)); + return new HiveBasicStatistics(numFiles, numRows, inMemoryDataSizeInBytes, onDiskDataSizeInBytes); + } + + public static Map updateStatisticsParameters(Map parameters, HiveBasicStatistics statistics) + { + ImmutableMap.Builder result = ImmutableMap.builder(); + + parameters.forEach((key, value) -> { + if (!STATS_PROPERTIES.contains(key)) { + result.put(key, value); + } + }); + + statistics.getFileCount().ifPresent(count -> result.put(NUM_FILES, Long.toString(count))); + statistics.getRowCount().ifPresent(count -> result.put(NUM_ROWS, Long.toString(count))); + statistics.getInMemoryDataSizeInBytes().ifPresent(size -> result.put(RAW_DATA_SIZE, Long.toString(size))); + statistics.getOnDiskDataSizeInBytes().ifPresent(size -> result.put(TOTAL_SIZE, Long.toString(size))); + + return result.buildOrThrow(); + } + + private static OptionalLong toLong(@Nullable String parameterValue) + { + if (parameterValue == null) { + return OptionalLong.empty(); + } + Long longValue = Longs.tryParse(parameterValue); + if (longValue == null || longValue < 0) { + return OptionalLong.empty(); + } + return OptionalLong.of(longValue); + } } diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/Partition.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/Partition.java index 19e3e02bf80e..d577f1d43483 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/Partition.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/Partition.java @@ -143,6 +143,13 @@ public static Builder builder(Partition partition) return new Builder(partition); } + public Partition withParameters(Map parameters) + { + return builder(this) + .setParameters(parameters) + .build(); + } + public static class Builder { private final Storage.Builder storageBuilder; diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/SemiTransactionalHiveMetastore.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/SemiTransactionalHiveMetastore.java index 7d79a6fa7ff7..fb731a946d6a 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/SemiTransactionalHiveMetastore.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/SemiTransactionalHiveMetastore.java @@ -18,34 +18,36 @@ import com.google.common.base.Joiner; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableMultiset; import com.google.common.collect.ImmutableSet; import com.google.common.collect.Iterables; import com.google.common.collect.Lists; +import com.google.common.collect.Multiset; import com.google.errorprone.annotations.FormatMethod; import com.google.errorprone.annotations.concurrent.GuardedBy; +import dev.failsafe.Failsafe; +import dev.failsafe.FailsafeException; +import dev.failsafe.RetryPolicy; import io.airlift.log.Logger; import io.airlift.units.Duration; +import io.trino.filesystem.FileIterator; import io.trino.filesystem.Location; -import io.trino.hdfs.HdfsContext; -import io.trino.hdfs.HdfsEnvironment; -import io.trino.hive.thrift.metastore.DataOperationType; +import io.trino.filesystem.TrinoFileSystem; +import io.trino.filesystem.TrinoFileSystemFactory; import io.trino.plugin.hive.HiveBasicStatistics; -import io.trino.plugin.hive.HiveColumnStatisticType; -import io.trino.plugin.hive.HiveMetastoreClosure; import io.trino.plugin.hive.HiveTableHandle; import io.trino.plugin.hive.HiveType; import io.trino.plugin.hive.LocationHandle.WriteMode; import io.trino.plugin.hive.PartitionNotFoundException; import io.trino.plugin.hive.PartitionStatistics; import io.trino.plugin.hive.PartitionUpdateAndMergeResults; -import io.trino.plugin.hive.SchemaAlreadyExistsException; import io.trino.plugin.hive.TableAlreadyExistsException; import io.trino.plugin.hive.TableInvalidationCallback; import io.trino.plugin.hive.acid.AcidOperation; import io.trino.plugin.hive.acid.AcidTransaction; import io.trino.plugin.hive.metastore.HivePrivilegeInfo.HivePrivilege; +import io.trino.plugin.hive.projection.PartitionProjection; import io.trino.plugin.hive.security.SqlStandardAccessControlMetadataMetastore; -import io.trino.plugin.hive.util.RetryDriver; import io.trino.plugin.hive.util.ValidTxnWriteIdList; import io.trino.spi.TrinoException; import io.trino.spi.connector.ConnectorSession; @@ -53,26 +55,22 @@ import io.trino.spi.connector.SchemaTableName; import io.trino.spi.connector.TableNotFoundException; import io.trino.spi.predicate.TupleDomain; +import io.trino.spi.security.ConnectorIdentity; import io.trino.spi.security.PrincipalType; import io.trino.spi.security.RoleGrant; -import io.trino.spi.type.Type; -import org.apache.hadoop.fs.FileStatus; -import org.apache.hadoop.fs.FileSystem; -import org.apache.hadoop.fs.LocatedFileStatus; -import org.apache.hadoop.fs.Path; -import org.apache.hadoop.fs.RemoteIterator; +import io.trino.spi.type.TypeManager; import java.io.FileNotFoundException; import java.io.IOException; +import java.util.AbstractMap; import java.util.ArrayList; import java.util.Collection; import java.util.HashMap; import java.util.HashSet; -import java.util.Iterator; import java.util.LinkedHashSet; -import java.util.LinkedList; import java.util.List; import java.util.Map; +import java.util.Map.Entry; import java.util.Optional; import java.util.OptionalLong; import java.util.Queue; @@ -87,7 +85,7 @@ import java.util.function.Supplier; import java.util.regex.Matcher; import java.util.regex.Pattern; -import java.util.stream.Collectors; +import java.util.stream.Stream; import static com.google.common.base.MoreObjects.toStringHelper; import static com.google.common.base.Preconditions.checkArgument; @@ -102,24 +100,28 @@ import static io.trino.plugin.hive.HiveErrorCode.HIVE_METASTORE_ERROR; import static io.trino.plugin.hive.HiveErrorCode.HIVE_PATH_ALREADY_EXISTS; import static io.trino.plugin.hive.HiveErrorCode.HIVE_TABLE_DROPPED_DURING_QUERY; -import static io.trino.plugin.hive.HiveMetadata.PRESTO_QUERY_ID_NAME; +import static io.trino.plugin.hive.HiveMetadata.TRINO_QUERY_ID_NAME; +import static io.trino.plugin.hive.HivePartitionManager.extractPartitionValues; import static io.trino.plugin.hive.LocationHandle.WriteMode.DIRECT_TO_TARGET_NEW_DIRECTORY; import static io.trino.plugin.hive.TableType.MANAGED_TABLE; -import static io.trino.plugin.hive.ViewReaderUtil.isPrestoView; +import static io.trino.plugin.hive.ViewReaderUtil.isTrinoMaterializedView; +import static io.trino.plugin.hive.ViewReaderUtil.isTrinoView; import static io.trino.plugin.hive.acid.AcidTransaction.NO_ACID_TRANSACTION; import static io.trino.plugin.hive.metastore.HivePrivilegeInfo.HivePrivilege.OWNERSHIP; import static io.trino.plugin.hive.metastore.MetastoreUtil.buildInitialPrivilegeSet; +import static io.trino.plugin.hive.metastore.MetastoreUtil.getHiveBasicStatistics; import static io.trino.plugin.hive.metastore.PrincipalPrivileges.NO_PRIVILEGES; +import static io.trino.plugin.hive.metastore.StatisticsUpdateMode.MERGE_INCREMENTAL; +import static io.trino.plugin.hive.metastore.StatisticsUpdateMode.OVERWRITE_ALL; +import static io.trino.plugin.hive.metastore.StatisticsUpdateMode.OVERWRITE_SOME_COLUMNS; +import static io.trino.plugin.hive.metastore.StatisticsUpdateMode.UNDO_MERGE_INCREMENTAL; +import static io.trino.plugin.hive.metastore.thrift.ThriftMetastoreUtil.getBasicStatisticsWithSparkFallback; +import static io.trino.plugin.hive.metastore.thrift.ThriftSparkMetastoreUtil.getSparkTableStatistics; +import static io.trino.plugin.hive.projection.PartitionProjectionProperties.getPartitionProjectionFromTable; import static io.trino.plugin.hive.util.AcidTables.isTransactionalTable; import static io.trino.plugin.hive.util.HiveUtil.makePartName; import static io.trino.plugin.hive.util.HiveUtil.toPartitionValues; -import static io.trino.plugin.hive.util.HiveWriteUtils.checkedDelete; -import static io.trino.plugin.hive.util.HiveWriteUtils.createDirectory; import static io.trino.plugin.hive.util.HiveWriteUtils.isFileCreatedByQuery; -import static io.trino.plugin.hive.util.HiveWriteUtils.pathExists; -import static io.trino.plugin.hive.util.Statistics.ReduceOperator.SUBTRACT; -import static io.trino.plugin.hive.util.Statistics.merge; -import static io.trino.plugin.hive.util.Statistics.reduce; import static io.trino.spi.StandardErrorCode.ALREADY_EXISTS; import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED; import static io.trino.spi.StandardErrorCode.TRANSACTION_CONFLICT; @@ -141,18 +143,23 @@ public class SemiTransactionalHiveMetastore { private static final Logger log = Logger.get(SemiTransactionalHiveMetastore.class); private static final int PARTITION_COMMIT_BATCH_SIZE = 20; - private static final Pattern DELTA_DIRECTORY_MATCHER = Pattern.compile("(delete_)?delta_[\\d]+_[\\d]+_[\\d]+$"); + private static final Pattern DELTA_DIRECTORY_MATCHER = Pattern.compile("(delete_)?delta_\\d+_\\d+_\\d+$"); - private static final RetryDriver DELETE_RETRY = RetryDriver.retry() - .maxAttempts(3) - .exponentialBackoff(new Duration(1, SECONDS), new Duration(1, SECONDS), new Duration(10, SECONDS), 2.0); + private static final RetryPolicy DELETE_RETRY_POLICY = RetryPolicy.builder() + .withDelay(java.time.Duration.ofSeconds(1)) + .withMaxDuration(java.time.Duration.ofSeconds(30)) + .withMaxAttempts(3) + .abortOn(TrinoFileSystem::isUnrecoverableException) + .build(); private static final Map ACID_OPERATION_ACTION_TYPES = ImmutableMap.of( AcidOperation.INSERT, ActionType.INSERT_EXISTING, AcidOperation.MERGE, ActionType.MERGE); - private final HiveMetastoreClosure delegate; - private final HdfsEnvironment hdfsEnvironment; + private final HiveMetastore delegate; + private final TypeManager typeManager; + private final boolean partitionProjectionEnabled; + private final TrinoFileSystemFactory fileSystemFactory; private final Executor fileSystemExecutor; private final Executor dropExecutor; private final Executor updateExecutor; @@ -163,8 +170,6 @@ public class SemiTransactionalHiveMetastore private final Optional configuredTransactionHeartbeatInterval; private final TableInvalidationCallback tableInvalidationCallback; - private boolean throwOnCleanupFailure; - @GuardedBy("this") private final Map> tableActions = new HashMap<>(); @GuardedBy("this") @@ -189,8 +194,10 @@ public class SemiTransactionalHiveMetastore private Optional currentHiveTransaction = Optional.empty(); public SemiTransactionalHiveMetastore( - HdfsEnvironment hdfsEnvironment, - HiveMetastoreClosure delegate, + TypeManager typeManager, + boolean partitionProjectionEnabled, + TrinoFileSystemFactory fileSystemFactory, + HiveMetastore delegate, Executor fileSystemExecutor, Executor dropExecutor, Executor updateExecutor, @@ -201,7 +208,9 @@ public SemiTransactionalHiveMetastore( ScheduledExecutorService heartbeatService, TableInvalidationCallback tableInvalidationCallback) { - this.hdfsEnvironment = requireNonNull(hdfsEnvironment, "hdfsEnvironment is null"); + this.typeManager = requireNonNull(typeManager, "typeManager is null"); + this.partitionProjectionEnabled = partitionProjectionEnabled; + this.fileSystemFactory = requireNonNull(fileSystemFactory, "fileSystemFactory is null"); this.delegate = requireNonNull(delegate, "delegate is null"); this.fileSystemExecutor = requireNonNull(fileSystemExecutor, "fileSystemExecutor is null"); this.dropExecutor = requireNonNull(dropExecutor, "dropExecutor is null"); @@ -214,17 +223,19 @@ public SemiTransactionalHiveMetastore( this.tableInvalidationCallback = requireNonNull(tableInvalidationCallback, "tableInvalidationCallback is null"); } - public synchronized List getAllDatabases() + public List getAllDatabases() { - checkReadable(); + synchronized (this) { + checkReadable(); + } return delegate.getAllDatabases(); } /** - * Get the underlying metastore closure. Use this method with caution as it bypasses the current transactional state, + * Get the underlying metastore. Use this method with caution as it bypasses the current transactional state, * so modifications made in the transaction are visible. */ - public HiveMetastoreClosure unsafeGetRawHiveMetastoreClosure() + public HiveMetastore unsafeGetRawHiveMetastore() { return delegate; } @@ -235,22 +246,15 @@ public synchronized Optional getDatabase(String databaseName) return delegate.getDatabase(databaseName); } - public synchronized List getAllTables(String databaseName) + public List getTables(String databaseName) { - checkReadable(); - if (!tableActions.isEmpty()) { - throw new UnsupportedOperationException("Listing all tables after adding/dropping/altering tables/views in a transaction is not supported"); - } - return delegate.getAllTables(databaseName); - } - - public synchronized Optional> getAllTables() - { - checkReadable(); - if (!tableActions.isEmpty()) { - throw new UnsupportedOperationException("Listing all tables after adding/dropping/altering tables/views in a transaction is not supported"); + synchronized (this) { + checkReadable(); + if (!tableActions.isEmpty()) { + throw new UnsupportedOperationException("Listing all tables after adding/dropping/altering tables/views in a transaction is not supported"); + } } - return delegate.getAllTables(); + return delegate.getTables(databaseName); } public synchronized Optional
getTable(String databaseName, String tableName) @@ -260,19 +264,11 @@ public synchronized Optional
getTable(String databaseName, String tableNa if (tableAction == null) { return delegate.getTable(databaseName, tableName); } - switch (tableAction.getType()) { - case ADD: - case ALTER: - case INSERT_EXISTING: - case MERGE: - return Optional.of(tableAction.getData().getTable()); - case DROP: - return Optional.empty(); - case DROP_PRESERVE_DATA: - // TODO - break; - } - throw new IllegalStateException("Unknown action type: " + tableAction.getType()); + return switch (tableAction.type()) { + case ADD, ALTER, INSERT_EXISTING, MERGE -> Optional.of(tableAction.data().getTable()); + case DROP -> Optional.empty(); + case DROP_PRESERVE_DATA -> throw new IllegalStateException("Unsupported action type: " + tableAction.type()); + }; } public synchronized boolean isReadableWithinTransaction(String databaseName, String tableName) @@ -281,24 +277,11 @@ public synchronized boolean isReadableWithinTransaction(String databaseName, Str if (tableAction == null) { return true; } - switch (tableAction.getType()) { - case ADD: - case ALTER: - return true; - case INSERT_EXISTING: - case MERGE: - // Until transaction is committed, the table data may or may not be visible. - return false; - case DROP: - case DROP_PRESERVE_DATA: - return false; - } - throw new IllegalStateException("Unknown action type: " + tableAction.getType()); - } - - public synchronized Set getSupportedColumnStatistics(Type type) - { - return delegate.getSupportedColumnStatistics(type); + return switch (tableAction.type()) { + case ADD, ALTER -> true; + case INSERT_EXISTING, MERGE -> false; // Until the transaction is committed, the table data may or may not be visible. + case DROP, DROP_PRESERVE_DATA -> false; + }; } public synchronized PartitionStatistics getTableStatistics(String databaseName, String tableName, Optional> columns) @@ -306,21 +289,32 @@ public synchronized PartitionStatistics getTableStatistics(String databaseName, checkReadable(); Action tableAction = tableActions.get(new SchemaTableName(databaseName, tableName)); if (tableAction == null) { - return delegate.getTableStatistics(databaseName, tableName, columns); - } - switch (tableAction.getType()) { - case ADD: - case ALTER: - case INSERT_EXISTING: - case MERGE: - return tableAction.getData().getStatistics(); - case DROP: - return PartitionStatistics.empty(); - case DROP_PRESERVE_DATA: - // TODO - break; + Table table = getExistingTable(databaseName, tableName); + Set columnNames = columns.orElseGet(() -> Stream.concat(table.getDataColumns().stream(), table.getPartitionColumns().stream()) + .map(Column::getName) + .collect(toImmutableSet())); + + if (delegate.useSparkTableStatistics()) { + Optional sparkTableStatistics = getSparkTableStatistics(table.getParameters(), columnNames.stream() + .map(table::getColumn) + .flatMap(Optional::stream) + .collect(toImmutableMap(Column::getName, Column::getType))); + if (sparkTableStatistics.isPresent()) { + return sparkTableStatistics.get(); + } + } + + HiveBasicStatistics basicStatistics = getHiveBasicStatistics(table.getParameters()); + if (columnNames.isEmpty()) { + return new PartitionStatistics(basicStatistics, ImmutableMap.of()); + } + return new PartitionStatistics(basicStatistics, delegate.getTableColumnStatistics(databaseName, tableName, columnNames)); } - throw new IllegalStateException("Unknown action type: " + tableAction.getType()); + return switch (tableAction.type()) { + case ADD, ALTER, INSERT_EXISTING, MERGE -> tableAction.data().getStatistics(); + case DROP -> PartitionStatistics.empty(); + case DROP_PRESERVE_DATA -> throw new IllegalStateException("Unsupported action type: " + tableAction.type()); + }; } public synchronized Map getPartitionStatistics(String databaseName, String tableName, Set columns, Set partitionNames) @@ -339,29 +333,75 @@ public synchronized Map getPartitionStatistics(Stri Action partitionAction = partitionActionsOfTable.get(partitionValues); if (partitionAction == null) { switch (tableSource) { - case PRE_EXISTING_TABLE: - partitionNamesToQuery.add(partitionName); - break; - case CREATED_IN_THIS_TRANSACTION: - resultBuilder.put(partitionName, PartitionStatistics.empty()); - break; - default: - throw new UnsupportedOperationException("unknown table source"); + case PRE_EXISTING_TABLE -> partitionNamesToQuery.add(partitionName); + case CREATED_IN_THIS_TRANSACTION -> resultBuilder.put(partitionName, PartitionStatistics.empty()); } } else { - resultBuilder.put(partitionName, partitionAction.getData().getStatistics()); + resultBuilder.put(partitionName, partitionAction.data().statistics()); } } - Map delegateResult = delegate.getPartitionStatistics(databaseName, tableName, partitionNamesToQuery.build(), Optional.of(columns)); - if (!delegateResult.isEmpty()) { - resultBuilder.putAll(delegateResult); + Set missingPartitions = partitionNamesToQuery.build(); + if (missingPartitions.isEmpty()) { + return resultBuilder.buildOrThrow(); } - else { - partitionNamesToQuery.build().forEach(partitionName -> resultBuilder.put(partitionName, PartitionStatistics.empty())); + + Map existingPartitions = getExistingPartitions(databaseName, tableName, partitionNames); + if (delegate.useSparkTableStatistics()) { + Map unprocessedPartitions = new HashMap<>(); + existingPartitions.forEach((partitionName, partition) -> { + Optional sparkPartitionStatistics = getSparkTableStatistics(partition.getParameters(), columns.stream() + .map(table.get()::getColumn) + .flatMap(Optional::stream) + .collect(toImmutableMap(Column::getName, Column::getType))); + sparkPartitionStatistics.ifPresentOrElse( + statistics -> resultBuilder.put(partitionName, statistics), + () -> unprocessedPartitions.put(partitionName, partition)); + }); + existingPartitions = unprocessedPartitions; } - return resultBuilder.buildOrThrow(); + + if (!existingPartitions.isEmpty()) { + Map basicStats = existingPartitions.entrySet().stream() + .collect(toImmutableMap(Entry::getKey, entry -> { + if (delegate.useSparkTableStatistics()) { + return getBasicStatisticsWithSparkFallback(entry.getValue().getParameters()); + } + return getHiveBasicStatistics(entry.getValue().getParameters()); + })); + + if (columns.isEmpty()) { + basicStats.forEach((partitionName, basicStatistics) -> resultBuilder.put(partitionName, new PartitionStatistics(basicStatistics, ImmutableMap.of()))); + } + else { + Map> columnStats = delegate.getPartitionColumnStatistics(databaseName, tableName, basicStats.keySet(), columns); + basicStats.forEach((key, value) -> resultBuilder.put(key, new PartitionStatistics(value, columnStats.getOrDefault(key, ImmutableMap.of())))); + } + } + return clearRowCountWhenAllPartitionsHaveNoRows(resultBuilder.buildOrThrow()); + } + + private static Map clearRowCountWhenAllPartitionsHaveNoRows(Map partitionStatistics) + { + if (partitionStatistics.isEmpty()) { + return partitionStatistics; + } + + // When the table has partitions, but row count statistics are set to zero, we treat this case as empty + // statistics to avoid underestimation in the CBO. This scenario may be caused when other engines are + // used to ingest data into partitioned hive tables. + long tableRowCount = partitionStatistics.values().stream() + .mapToLong(statistics -> statistics.getBasicStatistics().getRowCount().orElse(0)) + .sum(); + if (tableRowCount != 0) { + return partitionStatistics; + } + return partitionStatistics.entrySet().stream() + .map(entry -> new AbstractMap.SimpleEntry<>( + entry.getKey(), + entry.getValue().withBasicStatistics(entry.getValue().getBasicStatistics().withEmptyRowCount()))) + .collect(toImmutableMap(AbstractMap.SimpleEntry::getKey, AbstractMap.SimpleEntry::getValue)); } /** @@ -377,20 +417,12 @@ private TableSource getTableSource(String databaseName, String tableName) if (tableAction == null) { return TableSource.PRE_EXISTING_TABLE; } - switch (tableAction.getType()) { - case ADD: - return TableSource.CREATED_IN_THIS_TRANSACTION; - case DROP: - throw new TableNotFoundException(new SchemaTableName(databaseName, tableName)); - case ALTER: - case INSERT_EXISTING: - case MERGE: - return TableSource.PRE_EXISTING_TABLE; - case DROP_PRESERVE_DATA: - // TODO - break; - } - throw new IllegalStateException("Unknown action type: " + tableAction.getType()); + return switch (tableAction.type()) { + case ADD -> TableSource.CREATED_IN_THIS_TRANSACTION; + case DROP -> throw new TableNotFoundException(new SchemaTableName(databaseName, tableName)); + case ALTER, INSERT_EXISTING, MERGE -> TableSource.PRE_EXISTING_TABLE; + case DROP_PRESERVE_DATA -> throw new IllegalStateException("Unsupported action type: " + tableAction.type()); + }; } public synchronized HivePageSinkMetadata generatePageSinkMetadata(SchemaTableName schemaTableName) @@ -407,7 +439,7 @@ public synchronized HivePageSinkMetadata generatePageSinkMetadata(SchemaTableNam } else { ImmutableMap.Builder, Optional> modifiedPartitionMapBuilder = ImmutableMap.builder(); - for (Map.Entry, Action> entry : partitionActionMap.entrySet()) { + for (Entry, Action> entry : partitionActionMap.entrySet()) { modifiedPartitionMapBuilder.put(entry.getKey(), getPartitionFromPartitionAction(entry.getValue())); } modifiedPartitionMap = modifiedPartitionMapBuilder.buildOrThrow(); @@ -418,24 +450,6 @@ public synchronized HivePageSinkMetadata generatePageSinkMetadata(SchemaTableNam modifiedPartitionMap); } - public synchronized List getAllViews(String databaseName) - { - checkReadable(); - if (!tableActions.isEmpty()) { - throw new UnsupportedOperationException("Listing all tables after adding/dropping/altering tables/views in a transaction is not supported"); - } - return delegate.getAllViews(databaseName); - } - - public synchronized Optional> getAllViews() - { - checkReadable(); - if (!tableActions.isEmpty()) { - throw new UnsupportedOperationException("Listing all tables after adding/dropping/altering tables/views in a transaction is not supported"); - } - return delegate.getAllViews(); - } - public synchronized void createDatabase(ConnectorSession session, Database database) { String queryId = session.getQueryId(); @@ -446,31 +460,12 @@ public synchronized void createDatabase(ConnectorSession session, Database datab "Database '%s' does not have correct query id set", database.getDatabaseName()); - setExclusive((delegate, hdfsEnvironment) -> { - try { - delegate.createDatabase(database); - } - catch (SchemaAlreadyExistsException e) { - // Ignore SchemaAlreadyExistsException when database looks like created by us. - // This may happen when an actually successful metastore create call is retried - // e.g. because of a timeout on our side. - Optional existingDatabase = delegate.getDatabase(database.getDatabaseName()); - if (existingDatabase.isEmpty() || !isCreatedBy(existingDatabase.get(), queryId)) { - throw e; - } - } - }); - } - - private static boolean isCreatedBy(Database database, String queryId) - { - Optional databaseQueryId = getQueryId(database); - return databaseQueryId.isPresent() && databaseQueryId.get().equals(queryId); + setExclusive(delegate -> delegate.createDatabase(database)); } public synchronized void dropDatabase(ConnectorSession session, String schemaName) { - setExclusive((delegate, hdfsEnvironment) -> { + setExclusive(delegate -> { boolean deleteData = shouldDeleteDatabaseData(session, schemaName); delegate.dropDatabase(schemaName, deleteData); }); @@ -478,20 +473,21 @@ public synchronized void dropDatabase(ConnectorSession session, String schemaNam public boolean shouldDeleteDatabaseData(ConnectorSession session, String schemaName) { - Optional location = delegate.getDatabase(schemaName) + Optional location = delegate.getDatabase(schemaName) .orElseThrow(() -> new SchemaNotFoundException(schemaName)) .getLocation() - .map(Path::new); + .map(Location::of); // If we see files in the schema location, don't delete it. // If we see no files, request deletion. // If we fail to check the schema location, behave according to fallback. return location.map(path -> { try { - return !hdfsEnvironment.getFileSystem(new HdfsContext(session), path) - .listLocatedStatus(path).hasNext(); + TrinoFileSystem fileSystem = fileSystemFactory.create(session); + return !fileSystem.listFiles(path).hasNext() && + fileSystem.listDirectories(path).isEmpty(); } - catch (IOException | RuntimeException e) { + catch (IOException e) { log.warn(e, "Could not check schema directory '%s'", path); return deleteSchemaLocationsFallback; } @@ -500,64 +496,36 @@ public boolean shouldDeleteDatabaseData(ConnectorSession session, String schemaN public synchronized void renameDatabase(String source, String target) { - setExclusive((delegate, hdfsEnvironment) -> delegate.renameDatabase(source, target)); + setExclusive(delegate -> delegate.renameDatabase(source, target)); } public synchronized void setDatabaseOwner(String source, HivePrincipal principal) { - setExclusive((delegate, hdfsEnvironment) -> delegate.setDatabaseOwner(source, principal)); + setExclusive(delegate -> delegate.setDatabaseOwner(source, principal)); } // TODO: Allow updating statistics for 2 tables in the same transaction public synchronized void setTableStatistics(Table table, PartitionStatistics tableStatistics) { - AcidTransaction transaction = currentHiveTransaction.isPresent() ? currentHiveTransaction.get().getTransaction() : NO_ACID_TRANSACTION; - setExclusive((delegate, hdfsEnvironment) -> - delegate.updateTableStatistics(table.getDatabaseName(), table.getTableName(), transaction, statistics -> updatePartitionStatistics(statistics, tableStatistics))); + OptionalLong acidWriteId = getOptionalAcidTransaction().getOptionalWriteId(); + setExclusive(delegate -> delegate.updateTableStatistics(table.getDatabaseName(), table.getTableName(), acidWriteId, OVERWRITE_SOME_COLUMNS, tableStatistics)); } // TODO: Allow updating statistics for 2 tables in the same transaction public synchronized void setPartitionStatistics(Table table, Map, PartitionStatistics> partitionStatisticsMap) { - Map> updates = partitionStatisticsMap.entrySet().stream().collect( + Map updates = partitionStatisticsMap.entrySet().stream().collect( toImmutableMap( entry -> getPartitionName(table, entry.getKey()), - entry -> oldPartitionStats -> updatePartitionStatistics(oldPartitionStats, entry.getValue()))); - setExclusive((delegate, hdfsEnvironment) -> + Entry::getValue)); + setExclusive(delegate -> delegate.updatePartitionStatistics( - table.getDatabaseName(), - table.getTableName(), + delegate.getTable(table.getDatabaseName(), table.getTableName()) + .orElseThrow(() -> new TableNotFoundException(table.getSchemaTableName())), + OVERWRITE_SOME_COLUMNS, updates)); } - // For HiveBasicStatistics, we only overwrite the original statistics if the new one is not empty. - // For HiveColumnStatistics, only overwrite the original statistics for columns present in the new ones and preserve the others. - private static PartitionStatistics updatePartitionStatistics(PartitionStatistics oldPartitionStats, PartitionStatistics newPartitionStats) - { - HiveBasicStatistics oldBasicStatistics = oldPartitionStats.getBasicStatistics(); - HiveBasicStatistics newBasicStatistics = newPartitionStats.getBasicStatistics(); - HiveBasicStatistics updatedBasicStatistics = new HiveBasicStatistics( - firstPresent(newBasicStatistics.getFileCount(), oldBasicStatistics.getFileCount()), - firstPresent(newBasicStatistics.getRowCount(), oldBasicStatistics.getRowCount()), - firstPresent(newBasicStatistics.getInMemoryDataSizeInBytes(), oldBasicStatistics.getInMemoryDataSizeInBytes()), - firstPresent(newBasicStatistics.getOnDiskDataSizeInBytes(), oldBasicStatistics.getOnDiskDataSizeInBytes())); - Map updatedColumnStatistics = - updateColumnStatistics(oldPartitionStats.getColumnStatistics(), newPartitionStats.getColumnStatistics()); - return new PartitionStatistics(updatedBasicStatistics, updatedColumnStatistics); - } - - private static Map updateColumnStatistics(Map oldColumnStats, Map newColumnStats) - { - Map result = new HashMap<>(oldColumnStats); - result.putAll(newColumnStats); - return ImmutableMap.copyOf(result); - } - - private static OptionalLong firstPresent(OptionalLong first, OptionalLong second) - { - return first.isPresent() ? first : second; - } - /** * {@code currentLocation} needs to be supplied if a writePath exists for the table. */ @@ -565,41 +533,31 @@ public synchronized void createTable( ConnectorSession session, Table table, PrincipalPrivileges principalPrivileges, - Optional currentPath, + Optional currentLocation, Optional> files, boolean ignoreExisting, PartitionStatistics statistics, boolean cleanExtraOutputFilesOnCommit) { setShared(); - // When creating a table, it should never have partition actions. This is just a sanity check. + // When creating a table, it should never have partition actions. This is just a validation check. checkNoPartitionAction(table.getDatabaseName(), table.getTableName()); Action oldTableAction = tableActions.get(table.getSchemaTableName()); - TableAndMore tableAndMore = new TableAndMore(table, Optional.of(principalPrivileges), currentPath, files, ignoreExisting, statistics, statistics, cleanExtraOutputFilesOnCommit); + TableAndMore tableAndMore = new TableAndMore(table, Optional.of(principalPrivileges), currentLocation, files, ignoreExisting, statistics, statistics, cleanExtraOutputFilesOnCommit); if (oldTableAction == null) { - HdfsContext hdfsContext = new HdfsContext(session); - tableActions.put(table.getSchemaTableName(), new Action<>(ActionType.ADD, tableAndMore, hdfsContext, session.getQueryId())); + tableActions.put(table.getSchemaTableName(), new Action<>(ActionType.ADD, tableAndMore, session.getIdentity(), session.getQueryId())); return; } - switch (oldTableAction.getType()) { - case DROP: - if (!oldTableAction.getHdfsContext().getIdentity().getUser().equals(session.getUser())) { + switch (oldTableAction.type()) { + case DROP -> { + if (!oldTableAction.identity().getUser().equals(session.getUser())) { throw new TrinoException(TRANSACTION_CONFLICT, "Operation on the same table with different user in the same transaction is not supported"); } - HdfsContext hdfsContext = new HdfsContext(session); - tableActions.put(table.getSchemaTableName(), new Action<>(ActionType.ALTER, tableAndMore, hdfsContext, session.getQueryId())); - return; - - case ADD: - case ALTER: - case INSERT_EXISTING: - case MERGE: - throw new TableAlreadyExistsException(table.getSchemaTableName()); - case DROP_PRESERVE_DATA: - // TODO - break; + tableActions.put(table.getSchemaTableName(), new Action<>(ActionType.ALTER, tableAndMore, session.getIdentity(), session.getQueryId())); + } + case ADD, ALTER, INSERT_EXISTING, MERGE -> throw new TableAlreadyExistsException(table.getSchemaTableName()); + case DROP_PRESERVE_DATA -> throw new IllegalStateException("Unsupported action type: " + oldTableAction.type()); } - throw new IllegalStateException("Unknown action type: " + oldTableAction.getType()); } public synchronized void dropTable(ConnectorSession session, String databaseName, String tableName) @@ -609,34 +567,25 @@ public synchronized void dropTable(ConnectorSession session, String databaseName checkNoPartitionAction(databaseName, tableName); SchemaTableName schemaTableName = new SchemaTableName(databaseName, tableName); Action oldTableAction = tableActions.get(schemaTableName); - if (oldTableAction == null || oldTableAction.getType() == ActionType.ALTER) { - HdfsContext hdfsContext = new HdfsContext(session); - tableActions.put(schemaTableName, new Action<>(ActionType.DROP, null, hdfsContext, session.getQueryId())); + if (oldTableAction == null || oldTableAction.type() == ActionType.ALTER) { + tableActions.put(schemaTableName, new Action<>(ActionType.DROP, null, session.getIdentity(), session.getQueryId())); return; } - switch (oldTableAction.getType()) { - case DROP: - throw new TableNotFoundException(schemaTableName); - case ADD: - case ALTER: - case INSERT_EXISTING: - case MERGE: - throw new UnsupportedOperationException("dropping a table added/modified in the same transaction is not supported"); - case DROP_PRESERVE_DATA: - // TODO - break; + switch (oldTableAction.type()) { + case DROP -> throw new TableNotFoundException(schemaTableName); + case ADD, ALTER, INSERT_EXISTING, MERGE -> throw new UnsupportedOperationException("dropping a table added/modified in the same transaction is not supported"); + case DROP_PRESERVE_DATA -> throw new IllegalStateException("Unsupported action type: " + oldTableAction.type()); } - throw new IllegalStateException("Unknown action type: " + oldTableAction.getType()); } public synchronized void replaceTable(String databaseName, String tableName, Table table, PrincipalPrivileges principalPrivileges) { - setExclusive((delegate, hdfsEnvironment) -> delegate.replaceTable(databaseName, tableName, table, principalPrivileges)); + setExclusive(delegate -> delegate.replaceTable(databaseName, tableName, table, principalPrivileges, ImmutableMap.of())); } public synchronized void renameTable(String databaseName, String tableName, String newDatabaseName, String newTableName) { - setExclusive((delegate, hdfsEnvironment) -> { + setExclusive(delegate -> { Optional
oldTable = delegate.getTable(databaseName, tableName); try { delegate.renameTable(databaseName, tableName, newDatabaseName, newTableName); @@ -650,32 +599,32 @@ public synchronized void renameTable(String databaseName, String tableName, Stri public synchronized void commentTable(String databaseName, String tableName, Optional comment) { - setExclusive((delegate, hdfsEnvironment) -> delegate.commentTable(databaseName, tableName, comment)); + setExclusive(delegate -> delegate.commentTable(databaseName, tableName, comment)); } public synchronized void setTableOwner(String schema, String table, HivePrincipal principal) { - setExclusive((delegate, hdfsEnvironment) -> delegate.setTableOwner(schema, table, principal)); + setExclusive(delegate -> delegate.setTableOwner(schema, table, principal)); } public synchronized void commentColumn(String databaseName, String tableName, String columnName, Optional comment) { - setExclusive((delegate, hdfsEnvironment) -> delegate.commentColumn(databaseName, tableName, columnName, comment)); + setExclusive(delegate -> delegate.commentColumn(databaseName, tableName, columnName, comment)); } public synchronized void addColumn(String databaseName, String tableName, String columnName, HiveType columnType, String columnComment) { - setExclusive((delegate, hdfsEnvironment) -> delegate.addColumn(databaseName, tableName, columnName, columnType, columnComment)); + setExclusive(delegate -> delegate.addColumn(databaseName, tableName, columnName, columnType, columnComment)); } public synchronized void renameColumn(String databaseName, String tableName, String oldColumnName, String newColumnName) { - setExclusive((delegate, hdfsEnvironment) -> delegate.renameColumn(databaseName, tableName, oldColumnName, newColumnName)); + setExclusive(delegate -> delegate.renameColumn(databaseName, tableName, oldColumnName, newColumnName)); } public synchronized void dropColumn(String databaseName, String tableName, String columnName) { - setExclusive((delegate, hdfsEnvironment) -> delegate.dropColumn(databaseName, tableName, columnName)); + setExclusive(delegate -> delegate.dropColumn(databaseName, tableName, columnName)); } public synchronized void finishChangingExistingTable( @@ -697,10 +646,9 @@ public synchronized void finishChangingExistingTable( if (oldTableAction == null) { Table table = getExistingTable(schemaTableName.getSchemaName(), schemaTableName.getTableName()); if (isAcidTransactionRunning()) { - table = Table.builder(table).setWriteId(OptionalLong.of(currentHiveTransaction.orElseThrow().getTransaction().getWriteId())).build(); + table = Table.builder(table).setWriteId(OptionalLong.of(getRequiredAcidTransaction().getWriteId())).build(); } PartitionStatistics currentStatistics = getTableStatistics(databaseName, tableName, Optional.empty()); - HdfsContext hdfsContext = new HdfsContext(session); tableActions.put( schemaTableName, new Action<>( @@ -708,33 +656,26 @@ public synchronized void finishChangingExistingTable( new TableAndMore( table, Optional.empty(), - Optional.of(new Path(currentLocation.toString())), + Optional.of(currentLocation), Optional.of(fileNames), false, - merge(currentStatistics, statisticsUpdate), + MERGE_INCREMENTAL.updatePartitionStatistics(currentStatistics, statisticsUpdate), statisticsUpdate, cleanExtraOutputFilesOnCommit), - hdfsContext, + session.getIdentity(), session.getQueryId())); return; } - switch (oldTableAction.getType()) { - case DROP: - throw new TableNotFoundException(schemaTableName); - case ADD: - case ALTER: - case INSERT_EXISTING: - case MERGE: - throw new UnsupportedOperationException("Inserting into an unpartitioned table that were added, altered, or inserted into in the same transaction is not supported"); - case DROP_PRESERVE_DATA: - // TODO - break; + switch (oldTableAction.type()) { + case DROP -> throw new TableNotFoundException(schemaTableName); + case ADD, ALTER, INSERT_EXISTING, MERGE -> + throw new UnsupportedOperationException("Inserting into an unpartitioned table that were added, altered, or inserted into in the same transaction is not supported"); + case DROP_PRESERVE_DATA -> throw new IllegalStateException("Unsupported action type: " + oldTableAction.type()); } - throw new IllegalStateException("Unknown action type: " + oldTableAction.getType()); } - private boolean isAcidTransactionRunning() + private synchronized boolean isAcidTransactionRunning() { return currentHiveTransaction.isPresent() && currentHiveTransaction.get().getTransaction().isAcidTransactionRunning(); } @@ -752,15 +693,15 @@ public synchronized void truncateUnpartitionedTable(ConnectorSession session, St throw new IllegalArgumentException("Table is partitioned"); } - Path path = new Path(table.getStorage().getLocation()); - HdfsContext context = new HdfsContext(session); - setExclusive((delegate, hdfsEnvironment) -> { - RecursiveDeleteResult recursiveDeleteResult = recursiveDeleteFiles(hdfsEnvironment, context, path, ImmutableSet.of(""), false); - if (!recursiveDeleteResult.getNotDeletedEligibleItems().isEmpty()) { + Location location = Location.of(table.getStorage().getLocation()); + TrinoFileSystem fileSystem = fileSystemFactory.create(session); + setExclusive(delegate -> { + RecursiveDeleteResult recursiveDeleteResult = recursiveDeleteFiles(fileSystem, location, ImmutableSet.of(""), false); + if (!recursiveDeleteResult.notDeletedEligibleItems().isEmpty()) { throw new TrinoException(HIVE_FILESYSTEM_ERROR, format( "Error deleting from unpartitioned table %s. These items cannot be deleted: %s", schemaTableName, - recursiveDeleteResult.getNotDeletedEligibleItems())); + recursiveDeleteResult.notDeletedEligibleItems())); } }); } @@ -785,7 +726,6 @@ public synchronized void finishMerge( Action oldTableAction = tableActions.get(schemaTableName); if (oldTableAction == null) { Table table = getExistingTable(schemaTableName.getSchemaName(), schemaTableName.getTableName()); - HdfsContext hdfsContext = new HdfsContext(session); PrincipalPrivileges principalPrivileges = table.getOwner().isEmpty() ? NO_PRIVILEGES : buildInitialPrivilegeSet(table.getOwner().get()); tableActions.put( @@ -795,27 +735,19 @@ public synchronized void finishMerge( new TableAndMergeResults( table, Optional.of(principalPrivileges), - Optional.of(new Path(currentLocation.toString())), - partitionUpdateAndMergeResults, - partitions), - hdfsContext, + Optional.of(currentLocation), + partitionUpdateAndMergeResults), + session.getIdentity(), session.getQueryId())); return; } - switch (oldTableAction.getType()) { - case DROP: - throw new TableNotFoundException(schemaTableName); - case ADD: - case ALTER: - case INSERT_EXISTING: - case MERGE: - throw new UnsupportedOperationException("Inserting, updating or deleting in a table that was added, altered, inserted into, updated or deleted from in the same transaction is not supported"); - case DROP_PRESERVE_DATA: - // TODO - break; + switch (oldTableAction.type()) { + case DROP -> throw new TableNotFoundException(schemaTableName); + case ADD, ALTER, INSERT_EXISTING, MERGE -> + throw new UnsupportedOperationException("Inserting, updating or deleting in a table that was added, altered, inserted into, updated or deleted from in the same transaction is not supported"); + case DROP_PRESERVE_DATA -> throw new IllegalStateException("Unsupported action type: " + oldTableAction.type()); } - throw new IllegalStateException("Unknown action type: " + oldTableAction.getType()); } public synchronized Optional> getPartitionNames(String databaseName, String tableName) @@ -863,20 +795,22 @@ private Optional> doGetPartitionNames( } List partitionNames; TableSource tableSource = getTableSource(databaseName, tableName); - switch (tableSource) { - case CREATED_IN_THIS_TRANSACTION: - partitionNames = ImmutableList.of(); - break; - case PRE_EXISTING_TABLE: - partitionNames = delegate.getPartitionNamesByFilter(databaseName, tableName, columnNames, partitionKeysFilter) - .orElseThrow(() -> new TrinoException(TRANSACTION_CONFLICT, format("Table '%s.%s' was dropped by another transaction", databaseName, tableName))); - break; - default: - throw new UnsupportedOperationException("Unknown table source"); + partitionNames = switch (tableSource) { + case CREATED_IN_THIS_TRANSACTION -> ImmutableList.of(); + case PRE_EXISTING_TABLE -> getOptionalPartitions(databaseName, tableName, columnNames, partitionKeysFilter) + .orElseThrow(() -> new TrinoException(TRANSACTION_CONFLICT, format("Table '%s.%s' was dropped by another transaction", databaseName, tableName))); + }; + Set duplicatePartitionNames = ImmutableMultiset.copyOf(partitionNames) + .entrySet().stream() + .filter(entry -> entry.getCount() > 1) + .map(Multiset.Entry::getElement) + .collect(toImmutableSet()); + if (!duplicatePartitionNames.isEmpty()) { + throw new TrinoException(HIVE_METASTORE_ERROR, format("Metastore returned duplicate partition names for %s", duplicatePartitionNames)); } Map, Action> partitionActionsOfTable = partitionActions.computeIfAbsent(table.get().getSchemaTableName(), k -> new HashMap<>()); ImmutableList.Builder resultBuilder = ImmutableList.builder(); - // alter/remove newly-altered/dropped partitions from the results from underlying metastore + // alter/remove newly altered/dropped partitions from the results from underlying metastore for (String partitionName : partitionNames) { List partitionValues = toPartitionValues(partitionName); Action partitionAction = partitionActionsOfTable.get(partitionValues); @@ -884,27 +818,19 @@ private Optional> doGetPartitionNames( resultBuilder.add(partitionName); continue; } - switch (partitionAction.getType()) { - case ADD: - throw new TrinoException(TRANSACTION_CONFLICT, format("Another transaction created partition %s in table %s.%s", partitionValues, databaseName, tableName)); - case DROP: - case DROP_PRESERVE_DATA: + switch (partitionAction.type()) { + case ADD -> throw new TrinoException(TRANSACTION_CONFLICT, format("Another transaction created partition %s in table %s.%s", partitionValues, databaseName, tableName)); + case DROP, DROP_PRESERVE_DATA -> { // do nothing - break; - case ALTER: - case INSERT_EXISTING: - case MERGE: - resultBuilder.add(partitionName); - break; - default: - throw new IllegalStateException("Unknown action type: " + partitionAction.getType()); + } + case ALTER, INSERT_EXISTING, MERGE -> resultBuilder.add(partitionName); } } - // add newly-added partitions to the results from underlying metastore. + // add newly added partitions to the results from underlying metastore. if (!partitionActionsOfTable.isEmpty()) { for (Action partitionAction : partitionActionsOfTable.values()) { - if (partitionAction.getType() == ActionType.ADD) { - List values = partitionAction.getData().getPartition().getValues(); + if (partitionAction.type() == ActionType.ADD) { + List values = partitionAction.data().partition().getValues(); resultBuilder.add(makePartName(columnNames, values)); } } @@ -924,14 +850,8 @@ public synchronized Map> getPartitionsByNames(String Action partitionAction = partitionActionsOfTable.get(partitionValues); if (partitionAction == null) { switch (tableSource) { - case PRE_EXISTING_TABLE: - partitionNamesToQueryBuilder.add(partitionName); - break; - case CREATED_IN_THIS_TRANSACTION: - resultBuilder.put(partitionName, Optional.empty()); - break; - default: - throw new UnsupportedOperationException("unknown table source"); + case PRE_EXISTING_TABLE -> partitionNamesToQueryBuilder.add(partitionName); + case CREATED_IN_THIS_TRANSACTION -> resultBuilder.put(partitionName, Optional.empty()); } } else { @@ -941,7 +861,7 @@ public synchronized Map> getPartitionsByNames(String List partitionNamesToQuery = partitionNamesToQueryBuilder.build(); if (!partitionNamesToQuery.isEmpty()) { - Map> delegateResult = delegate.getPartitionsByNames( + Map> delegateResult = getOptionalPartitions( databaseName, tableName, partitionNamesToQuery); @@ -953,17 +873,10 @@ public synchronized Map> getPartitionsByNames(String private static Optional getPartitionFromPartitionAction(Action partitionAction) { - switch (partitionAction.getType()) { - case ADD: - case ALTER: - case INSERT_EXISTING: - case MERGE: - return Optional.of(partitionAction.getData().getAugmentedPartitionForInTransactionRead()); - case DROP: - case DROP_PRESERVE_DATA: - return Optional.empty(); - } - throw new IllegalStateException("Unknown action type: " + partitionAction.getType()); + return switch (partitionAction.type()) { + case ADD, ALTER, INSERT_EXISTING, MERGE -> Optional.of(partitionAction.data().getAugmentedPartitionForInTransactionRead()); + case DROP, DROP_PRESERVE_DATA -> Optional.empty(); + }; } public synchronized void addPartition( @@ -980,30 +893,24 @@ public synchronized void addPartition( checkArgument(getQueryId(partition).isPresent()); Map, Action> partitionActionsOfTable = partitionActions.computeIfAbsent(new SchemaTableName(databaseName, tableName), k -> new HashMap<>()); Action oldPartitionAction = partitionActionsOfTable.get(partition.getValues()); - HdfsContext hdfsContext = new HdfsContext(session); if (oldPartitionAction == null) { partitionActionsOfTable.put( partition.getValues(), - new Action<>(ActionType.ADD, new PartitionAndMore(partition, currentLocation, files, statistics, statistics, cleanExtraOutputFilesOnCommit), hdfsContext, session.getQueryId())); + new Action<>(ActionType.ADD, new PartitionAndMore(partition, currentLocation, files, statistics, statistics, cleanExtraOutputFilesOnCommit), session.getIdentity(), session.getQueryId())); return; } - switch (oldPartitionAction.getType()) { - case DROP: - case DROP_PRESERVE_DATA: - if (!oldPartitionAction.getHdfsContext().getIdentity().getUser().equals(session.getUser())) { + switch (oldPartitionAction.type()) { + case DROP, DROP_PRESERVE_DATA -> { + if (!oldPartitionAction.identity().getUser().equals(session.getUser())) { throw new TrinoException(TRANSACTION_CONFLICT, "Operation on the same partition with different user in the same transaction is not supported"); } partitionActionsOfTable.put( partition.getValues(), - new Action<>(ActionType.ALTER, new PartitionAndMore(partition, currentLocation, files, statistics, statistics, cleanExtraOutputFilesOnCommit), hdfsContext, session.getQueryId())); - return; - case ADD: - case ALTER: - case INSERT_EXISTING: - case MERGE: - throw new TrinoException(ALREADY_EXISTS, format("Partition already exists for table '%s.%s': %s", databaseName, tableName, partition.getValues())); + new Action<>(ActionType.ALTER, new PartitionAndMore(partition, currentLocation, files, statistics, statistics, cleanExtraOutputFilesOnCommit), session.getIdentity(), session.getQueryId())); + } + case ADD, ALTER, INSERT_EXISTING, MERGE -> + throw new TrinoException(ALREADY_EXISTS, format("Partition already exists for table '%s.%s': %s", databaseName, tableName, partition.getValues())); } - throw new IllegalStateException("Unknown action type: " + oldPartitionAction.getType()); } public synchronized void dropPartition(ConnectorSession session, String databaseName, String tableName, List partitionValues, boolean deleteData) @@ -1012,28 +919,20 @@ public synchronized void dropPartition(ConnectorSession session, String database Map, Action> partitionActionsOfTable = partitionActions.computeIfAbsent(new SchemaTableName(databaseName, tableName), k -> new HashMap<>()); Action oldPartitionAction = partitionActionsOfTable.get(partitionValues); if (oldPartitionAction == null) { - HdfsContext hdfsContext = new HdfsContext(session); if (deleteData) { - partitionActionsOfTable.put(partitionValues, new Action<>(ActionType.DROP, null, hdfsContext, session.getQueryId())); + partitionActionsOfTable.put(partitionValues, new Action<>(ActionType.DROP, null, session.getIdentity(), session.getQueryId())); } else { - partitionActionsOfTable.put(partitionValues, new Action<>(ActionType.DROP_PRESERVE_DATA, null, hdfsContext, session.getQueryId())); + partitionActionsOfTable.put(partitionValues, new Action<>(ActionType.DROP_PRESERVE_DATA, null, session.getIdentity(), session.getQueryId())); } return; } - switch (oldPartitionAction.getType()) { - case DROP: - case DROP_PRESERVE_DATA: - throw new PartitionNotFoundException(new SchemaTableName(databaseName, tableName), partitionValues); - case ADD: - case ALTER: - case INSERT_EXISTING: - case MERGE: - throw new TrinoException( - NOT_SUPPORTED, - format("dropping a partition added in the same transaction is not supported: %s %s %s", databaseName, tableName, partitionValues)); + switch (oldPartitionAction.type()) { + case DROP, DROP_PRESERVE_DATA -> throw new PartitionNotFoundException(new SchemaTableName(databaseName, tableName), partitionValues); + case ADD, ALTER, INSERT_EXISTING, MERGE -> + throw new TrinoException(NOT_SUPPORTED, "dropping a partition added in the same transaction is not supported: %s %s %s" + .formatted(databaseName, tableName, partitionValues)); } - throw new IllegalStateException("Unknown action type: " + oldPartitionAction.getType()); } public synchronized void finishInsertIntoExistingPartitions( @@ -1044,73 +943,68 @@ public synchronized void finishInsertIntoExistingPartitions( boolean cleanExtraOutputFilesOnCommit) { setShared(); - SchemaTableName schemaTableName = new SchemaTableName(databaseName, tableName); - HdfsContext context = new HdfsContext(session); - Map, Action> partitionActionsOfTable = partitionActions.computeIfAbsent(schemaTableName, k -> new HashMap<>()); + Table table = getExistingTable(databaseName, tableName); + Map, Action> partitionActionsOfTable = partitionActions.computeIfAbsent(table.getSchemaTableName(), k -> new HashMap<>()); for (PartitionUpdateInfo partitionInfo : partitionUpdateInfos) { - Action oldPartitionAction = partitionActionsOfTable.get(partitionInfo.partitionValues); + Action oldPartitionAction = partitionActionsOfTable.get(partitionInfo.partitionValues()); if (oldPartitionAction != null) { - switch (oldPartitionAction.getType()) { - case DROP, DROP_PRESERVE_DATA -> - throw new PartitionNotFoundException(schemaTableName, partitionInfo.partitionValues); - case ADD, ALTER, INSERT_EXISTING, MERGE -> - throw new UnsupportedOperationException("Inserting into a partition that were added, altered, or inserted into in the same transaction is not supported"); - default -> throw new IllegalStateException("Unknown action type: " + oldPartitionAction.getType()); + switch (oldPartitionAction.type()) { + case DROP, DROP_PRESERVE_DATA -> throw new PartitionNotFoundException(table.getSchemaTableName(), partitionInfo.partitionValues()); + case ADD, ALTER, INSERT_EXISTING, MERGE -> throw new UnsupportedOperationException("Inserting into a partition that were added, altered, or inserted into in the same transaction is not supported"); + default -> throw new IllegalStateException("Unknown action type: " + oldPartitionAction.type()); } } } + // new data will on include current table columns + // partition column stats do not include the partition keys + Set columnNames = table.getDataColumns().stream() + .map(Column::getName) + .collect(toImmutableSet()); for (List partitionInfoBatch : Iterables.partition(partitionUpdateInfos, 100)) { List partitionNames = partitionInfoBatch.stream() .map(PartitionUpdateInfo::partitionValues) .map(partitionValues -> getPartitionName(databaseName, tableName, partitionValues)) .collect(toImmutableList()); - Map> partitionsByNames = delegate.getPartitionsByNames( - schemaTableName.getSchemaName(), - schemaTableName.getTableName(), - partitionNames); - Map partitionStatistics = delegate.getPartitionStatistics( - schemaTableName.getSchemaName(), - schemaTableName.getTableName(), - ImmutableSet.copyOf(partitionNames)); + Map partitionsByNames = getExistingPartitions(databaseName, tableName, partitionNames); + Map basicStats = partitionsByNames.entrySet().stream() + .collect(toImmutableMap(Entry::getKey, entry -> getHiveBasicStatistics(entry.getValue().getParameters()))); + Map> columnStats = delegate.getPartitionColumnStatistics(databaseName, tableName, basicStats.keySet(), columnNames); for (int i = 0; i < partitionInfoBatch.size(); i++) { PartitionUpdateInfo partitionInfo = partitionInfoBatch.get(i); String partitionName = partitionNames.get(i); - Optional partition = partitionsByNames.get(partitionName); - if (partition.isEmpty()) { - throw new PartitionNotFoundException(schemaTableName, partitionInfo.partitionValues); - } - - PartitionStatistics currentStatistics = partitionStatistics.get(partitionName); - if (currentStatistics == null) { - throw new TrinoException(HIVE_METASTORE_ERROR, "currentStatistics is null"); - } + Partition partition = partitionsByNames.get(partitionName); + PartitionStatistics currentStatistics = new PartitionStatistics(basicStats.get(partitionName), columnStats.get(partitionName)); partitionActionsOfTable.put( - partitionInfo.partitionValues, + partitionInfo.partitionValues(), new Action<>( ActionType.INSERT_EXISTING, new PartitionAndMore( - partition.get(), - partitionInfo.currentLocation, - Optional.of(partitionInfo.fileNames), - merge(currentStatistics, partitionInfo.statisticsUpdate), - partitionInfo.statisticsUpdate, + partition, + partitionInfo.currentLocation(), + Optional.of(partitionInfo.fileNames()), + MERGE_INCREMENTAL.updatePartitionStatistics(currentStatistics, partitionInfo.statisticsUpdate()), + partitionInfo.statisticsUpdate(), cleanExtraOutputFilesOnCommit), - context, + session.getIdentity(), session.getQueryId())); } } } - private synchronized AcidTransaction getCurrentAcidTransaction() + private synchronized AcidTransaction getRequiredAcidTransaction() + { + return currentHiveTransaction.orElseThrow(() -> new IllegalStateException("currentHiveTransaction not present")).getTransaction(); + } + + private synchronized AcidTransaction getOptionalAcidTransaction() { - return currentHiveTransaction.map(HiveTransaction::getTransaction) - .orElseThrow(() -> new IllegalStateException("currentHiveTransaction not present")); + return currentHiveTransaction.map(HiveTransaction::getTransaction).orElse(NO_ACID_TRANSACTION); } private String getPartitionName(String databaseName, String tableName, List partitionValues) @@ -1131,13 +1025,13 @@ private static String getPartitionName(Table table, List partitionValues @Override public synchronized void createRole(String role, String grantor) { - setExclusive((delegate, hdfsEnvironment) -> delegate.createRole(role, grantor)); + setExclusive(delegate -> delegate.createRole(role, grantor)); } @Override public synchronized void dropRole(String role) { - setExclusive((delegate, hdfsEnvironment) -> delegate.dropRole(role)); + setExclusive(delegate -> delegate.dropRole(role)); } @Override @@ -1150,20 +1044,13 @@ public synchronized Set listRoles() @Override public synchronized void grantRoles(Set roles, Set grantees, boolean adminOption, HivePrincipal grantor) { - setExclusive((delegate, hdfsEnvironment) -> delegate.grantRoles(roles, grantees, adminOption, grantor)); + setExclusive(delegate -> delegate.grantRoles(roles, grantees, adminOption, grantor)); } @Override public synchronized void revokeRoles(Set roles, Set grantees, boolean adminOption, HivePrincipal grantor) { - setExclusive((delegate, hdfsEnvironment) -> delegate.revokeRoles(roles, grantees, adminOption, grantor)); - } - - @Override - public synchronized Set listGrantedPrincipals(String role) - { - checkReadable(); - return delegate.listGrantedPrincipals(role); + setExclusive(delegate -> delegate.revokeRoles(roles, grantees, adminOption, grantor)); } @Override @@ -1191,36 +1078,30 @@ public synchronized Set listTablePrivileges(String databaseNa if (tableAction == null) { return delegate.listTablePrivileges(databaseName, tableName, getExistingTable(databaseName, tableName).getOwner(), principal); } - switch (tableAction.getType()) { - case ADD: - case ALTER: + return switch (tableAction.type()) { + case ADD, ALTER -> { if (principal.isPresent() && principal.get().getType() == PrincipalType.ROLE) { - return ImmutableSet.of(); + yield ImmutableSet.of(); } - Optional owner = tableAction.getData().getTable().getOwner(); + Optional owner = tableAction.data().getTable().getOwner(); if (owner.isEmpty()) { // todo the existing logic below seem off. Only permissions held by the table owner are returned - return ImmutableSet.of(); + yield ImmutableSet.of(); } String ownerUsername = owner.orElseThrow(); if (principal.isPresent() && !principal.get().getName().equals(ownerUsername)) { - return ImmutableSet.of(); + yield ImmutableSet.of(); } - Collection privileges = tableAction.getData().getPrincipalPrivileges().getUserPrivileges().get(ownerUsername); - return ImmutableSet.builder() + Collection privileges = tableAction.data().getPrincipalPrivileges().getUserPrivileges().get(ownerUsername); + yield ImmutableSet.builder() .addAll(privileges) .add(new HivePrivilegeInfo(OWNERSHIP, true, new HivePrincipal(USER, ownerUsername), new HivePrincipal(USER, ownerUsername))) .build(); - case INSERT_EXISTING: - case MERGE: - return delegate.listTablePrivileges(databaseName, tableName, getExistingTable(databaseName, tableName).getOwner(), principal); - case DROP: - throw new TableNotFoundException(schemaTableName); - case DROP_PRESERVE_DATA: - // TODO - break; - } - throw new IllegalStateException("Unknown action type: " + tableAction.getType()); + } + case INSERT_EXISTING, MERGE -> delegate.listTablePrivileges(databaseName, tableName, getExistingTable(databaseName, tableName).getOwner(), principal); + case DROP -> throw new TableNotFoundException(schemaTableName); + case DROP_PRESERVE_DATA -> throw new IllegalStateException("Unsupported action type: " + tableAction.type()); + }; } private synchronized String getRequiredTableOwner(String databaseName, String tableName) @@ -1234,16 +1115,60 @@ private Table getExistingTable(String databaseName, String tableName) .orElseThrow(() -> new TableNotFoundException(new SchemaTableName(databaseName, tableName))); } + private Map getExistingPartitions(String databaseName, String tableName, Collection partitionNames) + { + return getOptionalPartitions(databaseName, tableName, ImmutableList.copyOf(partitionNames)).entrySet().stream() + .collect(toImmutableMap(Entry::getKey, entry -> entry.getValue() + .orElseThrow(() -> new PartitionNotFoundException(new SchemaTableName(databaseName, tableName), extractPartitionValues(entry.getKey()))))); + } + + private Map> getOptionalPartitions(String databaseName, String tableName, List partitionNames) + { + return delegate.getTable(databaseName, tableName) + .map(table -> getOptionalPartitions(table, partitionNames)) + .orElseGet(() -> partitionNames.stream() + .collect(toImmutableMap(name -> name, ignore -> Optional.empty()))); + } + + private Map> getOptionalPartitions(Table table, List partitionNames) + { + if (partitionProjectionEnabled) { + Optional projection = getPartitionProjectionFromTable(table, typeManager); + if (projection.isPresent()) { + return projection.get().getProjectedPartitionsByNames(table, partitionNames); + } + } + return delegate.getPartitionsByNames(table, partitionNames); + } + + private Optional> getOptionalPartitions( + String databaseName, + String tableName, + List columnNames, + TupleDomain partitionKeysFilter) + { + if (partitionProjectionEnabled) { + Table table = getTable(databaseName, tableName) + .orElseThrow(() -> new TrinoException(HIVE_TABLE_DROPPED_DURING_QUERY, "Table does not exists: " + tableName)); + + Optional projection = getPartitionProjectionFromTable(table, typeManager); + if (projection.isPresent()) { + return projection.get().getProjectedPartitionNamesByFilter(columnNames, partitionKeysFilter); + } + } + return delegate.getPartitionNamesByFilter(databaseName, tableName, columnNames, partitionKeysFilter); + } + @Override public synchronized void grantTablePrivileges(String databaseName, String tableName, HivePrincipal grantee, HivePrincipal grantor, Set privileges, boolean grantOption) { - setExclusive((delegate, hdfsEnvironment) -> delegate.grantTablePrivileges(databaseName, tableName, getRequiredTableOwner(databaseName, tableName), grantee, grantor, privileges, grantOption)); + setExclusive(delegate -> delegate.grantTablePrivileges(databaseName, tableName, getRequiredTableOwner(databaseName, tableName), grantee, grantor, privileges, grantOption)); } @Override public synchronized void revokeTablePrivileges(String databaseName, String tableName, HivePrincipal grantee, HivePrincipal grantor, Set privileges, boolean grantOption) { - setExclusive((delegate, hdfsEnvironment) -> delegate.revokeTablePrivileges(databaseName, tableName, getRequiredTableOwner(databaseName, tableName), grantee, grantor, privileges, grantOption)); + setExclusive(delegate -> delegate.revokeTablePrivileges(databaseName, tableName, getRequiredTableOwner(databaseName, tableName), grantee, grantor, privileges, grantOption)); } public synchronized String declareIntentionToWrite(ConnectorSession session, WriteMode writeMode, Location stagingPathRoot, SchemaTableName schemaTableName) @@ -1255,23 +1180,23 @@ public synchronized String declareIntentionToWrite(ConnectorSession session, Wri throw new TrinoException(NOT_SUPPORTED, "Cannot insert into a table with a partition that has been modified in the same transaction when Trino is configured to skip temporary directories."); } } - HdfsContext hdfsContext = new HdfsContext(session); + ConnectorIdentity identity = session.getIdentity(); String queryId = session.getQueryId(); String declarationId = queryId + "_" + declaredIntentionsToWriteCounter; declaredIntentionsToWriteCounter++; - declaredIntentionsToWrite.add(new DeclaredIntentionToWrite(declarationId, writeMode, hdfsContext, queryId, new Path(stagingPathRoot.toString()), schemaTableName)); + declaredIntentionsToWrite.add(new DeclaredIntentionToWrite(declarationId, writeMode, identity, queryId, stagingPathRoot, schemaTableName)); return declarationId; } public synchronized void dropDeclaredIntentionToWrite(String declarationId) { - boolean removed = declaredIntentionsToWrite.removeIf(intention -> intention.getDeclarationId().equals(declarationId)); + boolean removed = declaredIntentionsToWrite.removeIf(intention -> intention.declarationId().equals(declarationId)); if (!removed) { throw new IllegalArgumentException("Declaration with id " + declarationId + " not found"); } } - public boolean isFinished() + public synchronized boolean isFinished() { return state == State.FINISHED; } @@ -1280,19 +1205,11 @@ public synchronized void commit() { try { switch (state) { - case EMPTY: - return; - case SHARED_OPERATION_BUFFERED: - commitShared(); - return; - case EXCLUSIVE_OPERATION_BUFFERED: - requireNonNull(bufferedExclusiveOperation, "bufferedExclusiveOperation is null"); - bufferedExclusiveOperation.execute(delegate, hdfsEnvironment); - return; - case FINISHED: - throw new IllegalStateException("Tried to commit buffered metastore operations after transaction has been committed/aborted"); + case EMPTY -> { /* nothing to do */ } + case SHARED_OPERATION_BUFFERED -> commitShared(); + case EXCLUSIVE_OPERATION_BUFFERED -> bufferedExclusiveOperation.execute(delegate); + case FINISHED -> throw new IllegalStateException("Tried to commit buffered metastore operations after transaction has been committed/aborted"); } - throw new IllegalStateException("Unknown state: " + state); } finally { state = State.FINISHED; @@ -1303,16 +1220,10 @@ public synchronized void rollback() { try { switch (state) { - case EMPTY: - case EXCLUSIVE_OPERATION_BUFFERED: - return; - case SHARED_OPERATION_BUFFERED: - rollbackShared(); - return; - case FINISHED: - throw new IllegalStateException("Tried to rollback buffered metastore operations after transaction has been committed/aborted"); + case EMPTY, EXCLUSIVE_OPERATION_BUFFERED -> { /* nothing to do */ } + case SHARED_OPERATION_BUFFERED -> rollbackShared(); + case FINISHED -> throw new IllegalStateException("Tried to rollback buffered metastore operations after transaction has been committed/aborted"); } - throw new IllegalStateException("Unknown state: " + state); } finally { state = State.FINISHED; @@ -1342,15 +1253,15 @@ public void beginQuery(ConnectorSession session) public AcidTransaction beginInsert(ConnectorSession session, Table table) { - return beginOperation(session, table, AcidOperation.INSERT, DataOperationType.INSERT); + return beginOperation(session, table, AcidOperation.INSERT); } public AcidTransaction beginMerge(ConnectorSession session, Table table) { - return beginOperation(session, table, AcidOperation.MERGE, DataOperationType.UPDATE); + return beginOperation(session, table, AcidOperation.MERGE); } - private AcidTransaction beginOperation(ConnectorSession session, Table table, AcidOperation operation, DataOperationType hiveOperation) + private AcidTransaction beginOperation(ConnectorSession session, Table table, AcidOperation operation) { String queryId = session.getQueryId(); @@ -1358,7 +1269,7 @@ private AcidTransaction beginOperation(ConnectorSession session, Table table, Ac currentQueryId = Optional.of(queryId); // We start the transaction immediately, and allocate the write lock and the writeId, - // because we need the writeId in order to write the delta files. + // because we need the writeId to write the delta files. HiveTransaction hiveTransaction = makeHiveTransaction(session, transactionId -> { acquireTableWriteLock( new AcidTransactionOwner(session.getUser()), @@ -1366,7 +1277,7 @@ private AcidTransaction beginOperation(ConnectorSession session, Table table, Ac transactionId, table.getDatabaseName(), table.getTableName(), - hiveOperation, + operation, !table.getPartitionColumns().isEmpty()); long writeId = allocateWriteId(table.getDatabaseName(), table.getTableName(), transactionId); return new AcidTransaction(operation, transactionId, writeId); @@ -1498,7 +1409,7 @@ private void postCommitCleanup(Optional transaction, boolean co heartbeatTask.cancel(true); if (commit) { - // Any failure around aborted transactions, etc would be handled by Hive Metastore commit and TrinoException will be thrown + // Any failure around aborted transactions, etc. would be handled by Hive Metastore commit, and TrinoException will be thrown delegate.commitTransaction(transactionId); } else { @@ -1506,7 +1417,6 @@ private void postCommitCleanup(Optional transaction, boolean co } } - @GuardedBy("this") private synchronized void clearCurrentTransaction() { currentQueryId = Optional.empty(); @@ -1519,59 +1429,33 @@ private void commitShared() { checkHoldsLock(); - AcidTransaction transaction = currentHiveTransaction.isEmpty() ? NO_ACID_TRANSACTION : currentHiveTransaction.get().getTransaction(); + AcidTransaction transaction = getOptionalAcidTransaction(); Committer committer = new Committer(transaction); try { - for (Map.Entry> entry : tableActions.entrySet()) { + for (Entry> entry : tableActions.entrySet()) { SchemaTableName schemaTableName = entry.getKey(); Action action = entry.getValue(); - switch (action.getType()) { - case DROP: - committer.prepareDropTable(schemaTableName); - break; - case ALTER: - committer.prepareAlterTable(action.getHdfsContext(), action.getQueryId(), action.getData()); - break; - case ADD: - committer.prepareAddTable(action.getHdfsContext(), action.getQueryId(), action.getData()); - break; - case INSERT_EXISTING: - committer.prepareInsertExistingTable(action.getHdfsContext(), action.getQueryId(), action.getData()); - break; - case MERGE: - committer.prepareMergeExistingTable(action.getHdfsContext(), action.getData()); - break; - default: - throw new IllegalStateException("Unknown action type: " + action.getType()); + switch (action.type()) { + case DROP -> committer.prepareDropTable(schemaTableName); + case ALTER -> committer.prepareAlterTable(action.identity(), action.queryId(), action.data()); + case ADD -> committer.prepareAddTable(action.identity(), action.queryId(), action.data()); + case INSERT_EXISTING -> committer.prepareInsertExistingTable(action.identity(), action.queryId(), action.data()); + case MERGE -> committer.prepareMergeExistingTable(action.identity(), action.data()); + case DROP_PRESERVE_DATA -> throw new IllegalArgumentException("Unsupported action type: " + action.type()); } } - for (Map.Entry, Action>> tableEntry : partitionActions.entrySet()) { + for (Entry, Action>> tableEntry : partitionActions.entrySet()) { SchemaTableName schemaTableName = tableEntry.getKey(); - for (Map.Entry, Action> partitionEntry : tableEntry.getValue().entrySet()) { + for (Entry, Action> partitionEntry : tableEntry.getValue().entrySet()) { List partitionValues = partitionEntry.getKey(); Action action = partitionEntry.getValue(); - switch (action.getType()) { - case DROP: - committer.prepareDropPartition(schemaTableName, partitionValues, true); - break; - case DROP_PRESERVE_DATA: - committer.prepareDropPartition(schemaTableName, partitionValues, false); - break; - case ALTER: - committer.prepareAlterPartition(action.getHdfsContext(), action.getQueryId(), action.getData()); - break; - case ADD: - committer.prepareAddPartition(action.getHdfsContext(), action.getQueryId(), action.getData()); - break; - case INSERT_EXISTING: - committer.prepareInsertExistingPartition(action.getHdfsContext(), action.getQueryId(), action.getData()); - break; - case MERGE: - committer.prepareInsertExistingPartition(action.getHdfsContext(), action.getQueryId(), action.getData()); - break; - default: - throw new IllegalStateException("Unknown action type: " + action.getType()); + switch (action.type()) { + case DROP -> committer.prepareDropPartition(schemaTableName, partitionValues, true); + case DROP_PRESERVE_DATA -> committer.prepareDropPartition(schemaTableName, partitionValues, false); + case ALTER -> committer.prepareAlterPartition(action.identity(), action.queryId(), action.data()); + case ADD -> committer.prepareAddPartition(action.identity(), action.queryId(), action.data()); + case INSERT_EXISTING, MERGE -> committer.prepareInsertExistingPartition(action.identity(), action.queryId(), action.data()); } } } @@ -1604,7 +1488,7 @@ private void commitShared() committer.executeRenameTasksForAbort(); - // Partition directory must be put back before relevant metastore operation can be undone + // Partition directory must be put back before the relevant metastore operation can be undone committer.undoAlterTableOperations(); committer.undoAlterPartitionOperations(); @@ -1623,21 +1507,21 @@ private void commitShared() try { // After this line, operations are no longer reversible. // The next section will deal with "dropping table/partition". Commit may still fail in - // this section. Even if commit fails, cleanups, instead of rollbacks, will be executed. + // this section. Even if the commit fails, cleanups, instead of rollbacks, will be executed. committer.executeIrreversibleMetastoreOperations(); // If control flow reached this point, this commit is considered successful no matter - // what happens later. The only kind of operations that haven't been carried out yet - // are cleanups. + // what happens later. The only operations that haven't been carried out yet + // are cleanup operations. // The program control flow will go to finally next. And cleanup will run because // moveForwardInFinally has been set to false. } finally { - // In this method, all operations are best-effort clean up operations. + // In this method, all operations are best-effort cleanup operations. // If any operation fails, the error will be logged and ignored. - // Additionally, other clean up operations should still be attempted. + // Additionally, other cleanup operations should still be attempted. // Execute deletion tasks committer.executeDeletionTasksForFinish(); @@ -1653,7 +1537,7 @@ private class Committer private final List> fileSystemOperationFutures = new ArrayList<>(); // File system - // For file system changes, only operations outside of writing paths (as specified in declared intentions to write) + // For file system changes, only operations outside the writing paths (as specified in declared intentions to write) // need to MOVE_BACKWARD tasks scheduled. Files in writing paths are handled by rollbackShared(). private final List deletionTasksForFinish = new ArrayList<>(); private final List renameTasksForAbort = new ArrayList<>(); @@ -1697,21 +1581,20 @@ private void prepareDropTable(SchemaTableName schemaTableName) })); } - private void prepareAlterTable(HdfsContext hdfsContext, String queryId, TableAndMore tableAndMore) + private void prepareAlterTable(ConnectorIdentity identity, String queryId, TableAndMore tableAndMore) { deleteOnly = false; Table table = tableAndMore.getTable(); - String targetLocation = table.getStorage().getLocation(); + Location targetLocation = Location.of(table.getStorage().getLocation()); Table oldTable = delegate.getTable(table.getDatabaseName(), table.getTableName()) .orElseThrow(() -> new TrinoException(TRANSACTION_CONFLICT, "The table that this transaction modified was deleted in another transaction. " + table.getSchemaTableName())); - String oldTableLocation = oldTable.getStorage().getLocation(); - Path oldTablePath = new Path(oldTableLocation); + Location oldTableLocation = Location.of(oldTable.getStorage().getLocation()); tablesToInvalidate.add(oldTable); - cleanExtraOutputFiles(hdfsContext, queryId, tableAndMore); + cleanExtraOutputFiles(identity, queryId, tableAndMore); - // Location of the old table and the new table can be different because we allow arbitrary directories through LocationService. + // The location of the old table and the new table can be different because we allow arbitrary directories through LocationService. // If the location of the old table is the same as the location of the new table: // * Rename the old data directory to a temporary path with a special suffix // * Remember we will need to delete that directory at the end if transaction successfully commits @@ -1719,36 +1602,34 @@ private void prepareAlterTable(HdfsContext hdfsContext, String queryId, TableAnd // Otherwise, // * Remember we will need to delete the location of the old partition at the end if transaction successfully commits if (targetLocation.equals(oldTableLocation)) { - Path oldTableStagingPath = new Path(oldTablePath.getParent(), "_temp_" + oldTablePath.getName() + "_" + queryId); + Location location = asFileLocation(oldTableLocation); + Location oldTableStagingPath = location.parentDirectory().appendPath("_temp_" + location.fileName() + "_" + queryId); renameDirectory( - hdfsContext, - hdfsEnvironment, - oldTablePath, + fileSystemFactory.create(identity), + oldTableLocation, oldTableStagingPath, - () -> renameTasksForAbort.add(new DirectoryRenameTask(hdfsContext, oldTableStagingPath, oldTablePath))); + () -> renameTasksForAbort.add(new DirectoryRenameTask(identity, oldTableStagingPath, oldTableLocation))); if (!skipDeletionForAlter) { - deletionTasksForFinish.add(new DirectoryDeletionTask(hdfsContext, oldTableStagingPath)); + deletionTasksForFinish.add(new DirectoryDeletionTask(identity, oldTableStagingPath)); } } else { if (!skipDeletionForAlter) { - deletionTasksForFinish.add(new DirectoryDeletionTask(hdfsContext, oldTablePath)); + deletionTasksForFinish.add(new DirectoryDeletionTask(identity, oldTableLocation)); } } - Path currentPath = tableAndMore.getCurrentLocation() + Location currentLocation = tableAndMore.getCurrentLocation() .orElseThrow(() -> new IllegalArgumentException("location should be present for alter table")); - Path targetPath = new Path(targetLocation); - if (!targetPath.equals(currentPath)) { + if (!targetLocation.equals(currentLocation)) { renameDirectory( - hdfsContext, - hdfsEnvironment, - currentPath, - targetPath, - () -> cleanUpTasksForAbort.add(new DirectoryCleanUpTask(hdfsContext, targetPath, true))); + fileSystemFactory.create(identity), + currentLocation, + targetLocation, + () -> cleanUpTasksForAbort.add(new DirectoryCleanUpTask(identity, targetLocation, true))); } - // Partition alter must happen regardless of whether original and current location is the same - // because metadata might change: e.g. storage format, column types, etc + // Partition alter must happen regardless of whether the original and current location is the same + // because metadata might change: e.g., storage format, column types, etc. alterTableOperations.add(new AlterTableOperation(tableAndMore.getTable(), oldTable, tableAndMore.getPrincipalPrivileges())); updateStatisticsOperations.add(new UpdateStatisticsOperation( @@ -1758,39 +1639,38 @@ private void prepareAlterTable(HdfsContext hdfsContext, String queryId, TableAnd false)); } - private void prepareAddTable(HdfsContext context, String queryId, TableAndMore tableAndMore) + private void prepareAddTable(ConnectorIdentity identity, String queryId, TableAndMore tableAndMore) { deleteOnly = false; - cleanExtraOutputFiles(context, queryId, tableAndMore); + cleanExtraOutputFiles(identity, queryId, tableAndMore); Table table = tableAndMore.getTable(); if (table.getTableType().equals(MANAGED_TABLE.name())) { - Optional targetLocation = table.getStorage().getOptionalLocation(); + Optional targetLocation = table.getStorage().getOptionalLocation().map(Location::of); if (targetLocation.isPresent()) { - checkArgument(!targetLocation.get().isEmpty(), "target location is empty"); - Optional currentPath = tableAndMore.getCurrentLocation(); - Path targetPath = new Path(targetLocation.get()); - if (table.getPartitionColumns().isEmpty() && currentPath.isPresent()) { + Optional currentLocation = tableAndMore.getCurrentLocation(); + Location targetPath = targetLocation.get(); + TrinoFileSystem fileSystem = fileSystemFactory.create(identity); + if (table.getPartitionColumns().isEmpty() && currentLocation.isPresent()) { // CREATE TABLE AS SELECT unpartitioned table - if (targetPath.equals(currentPath.get())) { + if (targetPath.equals(currentLocation.get())) { // Target path and current path are the same. Therefore, directory move is not needed. } else { renameDirectory( - context, - hdfsEnvironment, - currentPath.get(), + fileSystem, + currentLocation.get(), targetPath, - () -> cleanUpTasksForAbort.add(new DirectoryCleanUpTask(context, targetPath, true))); + () -> cleanUpTasksForAbort.add(new DirectoryCleanUpTask(identity, targetPath, true))); } } else { // CREATE TABLE AS SELECT partitioned table, or // CREATE TABLE partitioned/unpartitioned table (without data) - if (pathExists(context, hdfsEnvironment, targetPath)) { - if (currentPath.isPresent() && currentPath.get().equals(targetPath)) { - // It is okay to skip directory creation when currentPath is equal to targetPath + if (directoryExists(fileSystem, targetPath)) { + if (currentLocation.isPresent() && currentLocation.get().equals(targetPath)) { + // It is okay to skip directory creation when currentLocation is equal to targetPath // because the directory may have been created when creating partition directories. // However, it is important to note that the two being equal does not guarantee // a directory had been created. @@ -1802,32 +1682,33 @@ private void prepareAddTable(HdfsContext context, String queryId, TableAndMore t } } else { - cleanUpTasksForAbort.add(new DirectoryCleanUpTask(context, targetPath, true)); - createDirectory(context, hdfsEnvironment, targetPath); + cleanUpTasksForAbort.add(new DirectoryCleanUpTask(identity, targetPath, true)); + createDirectory(fileSystem, targetPath); } } } - // if targetLocation is not set in table we assume table directory is created by HMS + // if targetLocation is not set in table, we assume HMS creates table directory } addTableOperations.add(new CreateTableOperation(table, tableAndMore.getPrincipalPrivileges(), tableAndMore.isIgnoreExisting(), tableAndMore.getStatisticsUpdate())); } - private void prepareInsertExistingTable(HdfsContext context, String queryId, TableAndMore tableAndMore) + private void prepareInsertExistingTable(ConnectorIdentity identity, String queryId, TableAndMore tableAndMore) { deleteOnly = false; Table table = tableAndMore.getTable(); - Path targetPath = new Path(table.getStorage().getLocation()); + Location targetPath = Location.of(table.getStorage().getLocation()); tablesToInvalidate.add(table); - Path currentPath = tableAndMore.getCurrentLocation().orElseThrow(); - cleanUpTasksForAbort.add(new DirectoryCleanUpTask(context, targetPath, false)); + Location currentPath = tableAndMore.getCurrentLocation().orElseThrow(); + cleanUpTasksForAbort.add(new DirectoryCleanUpTask(identity, targetPath, false)); if (!targetPath.equals(currentPath)) { - // if staging directory is used we cherry-pick files to be moved - asyncRename(hdfsEnvironment, fileSystemExecutor, fileSystemOperationsCancelled, fileSystemOperationFutures, context, currentPath, targetPath, tableAndMore.getFileNames().orElseThrow()); + // if staging directory is used, we cherry-pick files to be moved + TrinoFileSystem fileSystem = fileSystemFactory.create(identity); + asyncRename(fileSystem, fileSystemExecutor, fileSystemOperationsCancelled, fileSystemOperationFutures, currentPath, targetPath, tableAndMore.getFileNames().orElseThrow()); } else { - // if we inserted directly into table directory we need to remove extra output files which should not be part of the table - cleanExtraOutputFiles(context, queryId, tableAndMore); + // if we inserted directly into table directory, we need to remove extra output files which should not be part of the table + cleanExtraOutputFiles(identity, queryId, tableAndMore); } updateStatisticsOperations.add(new UpdateStatisticsOperation( table.getSchemaTableName(), @@ -1836,24 +1717,24 @@ private void prepareInsertExistingTable(HdfsContext context, String queryId, Tab true)); if (isAcidTransactionRunning()) { - AcidTransaction transaction = getCurrentAcidTransaction(); + AcidTransaction transaction = getRequiredAcidTransaction(); updateTableWriteId(table.getDatabaseName(), table.getTableName(), transaction.getAcidTransactionId(), transaction.getWriteId(), OptionalLong.empty()); } } - private void prepareMergeExistingTable(HdfsContext context, TableAndMore tableAndMore) + private void prepareMergeExistingTable(ConnectorIdentity identity, TableAndMore tableAndMore) { - checkArgument(currentHiveTransaction.isPresent(), "currentHiveTransaction isn't present"); - AcidTransaction transaction = currentHiveTransaction.get().getTransaction(); + AcidTransaction transaction = getRequiredAcidTransaction(); checkArgument(transaction.isMerge(), "transaction should be merge, but is %s", transaction); deleteOnly = false; Table table = tableAndMore.getTable(); - Path targetPath = new Path(table.getStorage().getLocation()); - Path currentPath = tableAndMore.getCurrentLocation().get(); - cleanUpTasksForAbort.add(new DirectoryCleanUpTask(context, targetPath, false)); + Location targetPath = Location.of(table.getStorage().getLocation()); + Location currentPath = tableAndMore.getCurrentLocation().orElseThrow(); + cleanUpTasksForAbort.add(new DirectoryCleanUpTask(identity, targetPath, false)); if (!targetPath.equals(currentPath)) { - asyncRename(hdfsEnvironment, fileSystemExecutor, fileSystemOperationsCancelled, fileSystemOperationFutures, context, currentPath, targetPath, tableAndMore.getFileNames().get()); + TrinoFileSystem fileSystem = fileSystemFactory.create(identity); + asyncRename(fileSystem, fileSystemExecutor, fileSystemOperationsCancelled, fileSystemOperationFutures, currentPath, targetPath, tableAndMore.getFileNames().orElseThrow()); } updateStatisticsOperations.add(new UpdateStatisticsOperation( table.getSchemaTableName(), @@ -1869,7 +1750,7 @@ private void prepareDropPartition(SchemaTableName schemaTableName, List metastoreDeleteOperations.add(new IrreversibleMetastoreOperation( format("drop partition %s.%s %s", schemaTableName.getSchemaName(), schemaTableName.getTableName(), partitionValues), () -> { - Optional droppedPartition = delegate.getPartition(schemaTableName.getSchemaName(), schemaTableName.getTableName(), partitionValues); + Optional droppedPartition = getOptionalPartition(delegate, schemaTableName, partitionValues); try { delegate.dropPartition(schemaTableName.getSchemaName(), schemaTableName.getTableName(), partitionValues, deleteData); } @@ -1880,25 +1761,25 @@ private void prepareDropPartition(SchemaTableName schemaTableName, List })); } - private void prepareAlterPartition(HdfsContext hdfsContext, String queryId, PartitionAndMore partitionAndMore) + private void prepareAlterPartition(ConnectorIdentity identity, String queryId, PartitionAndMore partitionAndMore) { deleteOnly = false; - Partition partition = partitionAndMore.getPartition(); + Partition partition = partitionAndMore.partition(); partitionsToInvalidate.add(partition); String targetLocation = partition.getStorage().getLocation(); - Partition oldPartition = delegate.getPartition(partition.getDatabaseName(), partition.getTableName(), partition.getValues()) + Partition oldPartition = getOptionalPartition(delegate, partition.getSchemaTableName(), partition.getValues()) .orElseThrow(() -> new TrinoException( TRANSACTION_CONFLICT, format("The partition that this transaction modified was deleted in another transaction. %s %s", partition.getTableName(), partition.getValues()))); String partitionName = getPartitionName(partition.getDatabaseName(), partition.getTableName(), partition.getValues()); PartitionStatistics oldPartitionStatistics = getExistingPartitionStatistics(partition, partitionName); String oldPartitionLocation = oldPartition.getStorage().getLocation(); - Path oldPartitionPath = new Path(oldPartitionLocation); + Location oldPartitionPath = asFileLocation(Location.of(oldPartitionLocation)); - cleanExtraOutputFiles(hdfsContext, queryId, partitionAndMore); + cleanExtraOutputFiles(identity, queryId, partitionAndMore); - // Location of the old partition and the new partition can be different because we allow arbitrary directories through LocationService. + // The location of the old partition and the new partition can be different because we allow arbitrary directories through LocationService. // If the location of the old partition is the same as the location of the new partition: // * Rename the old data directory to a temporary path with a special suffix // * Remember we will need to delete that directory at the end if transaction successfully commits @@ -1906,71 +1787,77 @@ private void prepareAlterPartition(HdfsContext hdfsContext, String queryId, Part // Otherwise, // * Remember we will need to delete the location of the old partition at the end if transaction successfully commits if (targetLocation.equals(oldPartitionLocation)) { - Path oldPartitionStagingPath = new Path(oldPartitionPath.getParent(), "_temp_" + oldPartitionPath.getName() + "_" + queryId); + Location oldPartitionStagingPath = oldPartitionPath.sibling("_temp_" + oldPartitionPath.fileName() + "_" + queryId); renameDirectory( - hdfsContext, - hdfsEnvironment, + fileSystemFactory.create(identity), oldPartitionPath, oldPartitionStagingPath, - () -> renameTasksForAbort.add(new DirectoryRenameTask(hdfsContext, oldPartitionStagingPath, oldPartitionPath))); + () -> renameTasksForAbort.add(new DirectoryRenameTask(identity, oldPartitionStagingPath, oldPartitionPath))); if (!skipDeletionForAlter) { - deletionTasksForFinish.add(new DirectoryDeletionTask(hdfsContext, oldPartitionStagingPath)); + deletionTasksForFinish.add(new DirectoryDeletionTask(identity, oldPartitionStagingPath)); } } else { if (!skipDeletionForAlter) { - deletionTasksForFinish.add(new DirectoryDeletionTask(hdfsContext, oldPartitionPath)); + deletionTasksForFinish.add(new DirectoryDeletionTask(identity, oldPartitionPath)); } } - Path currentPath = new Path(partitionAndMore.getCurrentLocation().toString()); - Path targetPath = new Path(targetLocation); + Location currentPath = partitionAndMore.currentLocation(); + Location targetPath = Location.of(targetLocation); if (!targetPath.equals(currentPath)) { renameDirectory( - hdfsContext, - hdfsEnvironment, + fileSystemFactory.create(identity), currentPath, targetPath, - () -> cleanUpTasksForAbort.add(new DirectoryCleanUpTask(hdfsContext, targetPath, true))); + () -> cleanUpTasksForAbort.add(new DirectoryCleanUpTask(identity, targetPath, true))); } - // Partition alter must happen regardless of whether original and current location is the same - // because metadata might change: e.g. storage format, column types, etc + // Partition alter must happen regardless of whether the original and current location is the same + // because metadata might change: e.g., storage format, column types, etc. alterPartitionOperations.add(new AlterPartitionOperation( - new PartitionWithStatistics(partition, partitionName, partitionAndMore.getStatisticsUpdate()), + new PartitionWithStatistics(partition, partitionName, partitionAndMore.statisticsUpdate()), new PartitionWithStatistics(oldPartition, partitionName, oldPartitionStatistics))); } - private void cleanExtraOutputFiles(HdfsContext hdfsContext, String queryId, PartitionAndMore partitionAndMore) + private void cleanExtraOutputFiles(ConnectorIdentity identity, String queryId, PartitionAndMore partitionAndMore) { - if (!partitionAndMore.isCleanExtraOutputFilesOnCommit()) { + if (!partitionAndMore.cleanExtraOutputFilesOnCommit()) { return; } verify(partitionAndMore.hasFileNames(), "fileNames expected to be set if isCleanExtraOutputFilesOnCommit is true"); - SemiTransactionalHiveMetastore.cleanExtraOutputFiles(hdfsEnvironment, hdfsContext, queryId, partitionAndMore.getCurrentLocation(), ImmutableSet.copyOf(partitionAndMore.getFileNames())); + TrinoFileSystem fileSystem = fileSystemFactory.create(identity); + SemiTransactionalHiveMetastore.cleanExtraOutputFiles(fileSystem, queryId, partitionAndMore.currentLocation(), ImmutableSet.copyOf(partitionAndMore.getFileNames())); } - private void cleanExtraOutputFiles(HdfsContext hdfsContext, String queryId, TableAndMore tableAndMore) + private void cleanExtraOutputFiles(ConnectorIdentity identity, String queryId, TableAndMore tableAndMore) { if (!tableAndMore.isCleanExtraOutputFilesOnCommit()) { return; } - Path tableLocation = tableAndMore.getCurrentLocation().orElseThrow(() -> new IllegalArgumentException("currentLocation expected to be set if isCleanExtraOutputFilesOnCommit is true")); + TrinoFileSystem fileSystem = fileSystemFactory.create(identity); + Location tableLocation = tableAndMore.getCurrentLocation().orElseThrow(() -> + new IllegalArgumentException("currentLocation expected to be set if isCleanExtraOutputFilesOnCommit is true")); List files = tableAndMore.getFileNames().orElseThrow(() -> new IllegalArgumentException("fileNames expected to be set if isCleanExtraOutputFilesOnCommit is true")); - SemiTransactionalHiveMetastore.cleanExtraOutputFiles(hdfsEnvironment, hdfsContext, queryId, Location.of(tableLocation.toString()), ImmutableSet.copyOf(files)); + SemiTransactionalHiveMetastore.cleanExtraOutputFiles(fileSystem, queryId, tableLocation, ImmutableSet.copyOf(files)); } private PartitionStatistics getExistingPartitionStatistics(Partition partition, String partitionName) { try { - PartitionStatistics statistics = delegate.getPartitionStatistics(partition.getDatabaseName(), partition.getTableName(), ImmutableSet.of(partitionName)) + HiveBasicStatistics basicStatistics = getHiveBasicStatistics(partition.getParameters()); + Map columnStatistics = delegate.getPartitionColumnStatistics( + partition.getDatabaseName(), + partition.getTableName(), + ImmutableSet.of(partitionName), + partition.getColumns().stream().map(Column::getName).collect(toImmutableSet())) .get(partitionName); - if (statistics == null) { + if (columnStatistics == null) { throw new TrinoException( TRANSACTION_CONFLICT, format("The partition that this transaction modified was deleted in another transaction. %s %s", partition.getTableName(), partition.getValues())); } - return statistics; + return new PartitionStatistics(basicStatistics, columnStatistics); } catch (TrinoException e) { if (e.getErrorCode().equals(HIVE_CORRUPTED_COLUMN_STATISTICS.toErrorCode())) { @@ -1986,86 +1873,91 @@ private PartitionStatistics getExistingPartitionStatistics(Partition partition, } } - private void prepareAddPartition(HdfsContext hdfsContext, String queryId, PartitionAndMore partitionAndMore) + private void prepareAddPartition(ConnectorIdentity identity, String queryId, PartitionAndMore partitionAndMore) { deleteOnly = false; - Partition partition = partitionAndMore.getPartition(); + Partition partition = partitionAndMore.partition(); String targetLocation = partition.getStorage().getLocation(); - Path currentPath = new Path(partitionAndMore.getCurrentLocation().toString()); - Path targetPath = new Path(targetLocation); + Location currentPath = partitionAndMore.currentLocation(); + Location targetPath = Location.of(targetLocation); - cleanExtraOutputFiles(hdfsContext, queryId, partitionAndMore); + cleanExtraOutputFiles(identity, queryId, partitionAndMore); PartitionAdder partitionAdder = partitionAdders.computeIfAbsent( partition.getSchemaTableName(), - ignored -> new PartitionAdder(partition.getDatabaseName(), partition.getTableName(), delegate, PARTITION_COMMIT_BATCH_SIZE)); + ignore -> new PartitionAdder(partition.getDatabaseName(), partition.getTableName(), delegate, PARTITION_COMMIT_BATCH_SIZE)); fileSystemOperationFutures.add(CompletableFuture.runAsync(() -> { if (fileSystemOperationsCancelled.get()) { return; } - if (pathExists(hdfsContext, hdfsEnvironment, currentPath)) { + TrinoFileSystem fileSystem = fileSystemFactory.create(identity); + if (directoryExists(fileSystem, currentPath)) { if (!targetPath.equals(currentPath)) { renameDirectory( - hdfsContext, - hdfsEnvironment, + fileSystem, currentPath, targetPath, - () -> cleanUpTasksForAbort.add(new DirectoryCleanUpTask(hdfsContext, targetPath, true))); + () -> cleanUpTasksForAbort.add(new DirectoryCleanUpTask(identity, targetPath, true))); } } else { - cleanUpTasksForAbort.add(new DirectoryCleanUpTask(hdfsContext, targetPath, true)); - createDirectory(hdfsContext, hdfsEnvironment, targetPath); + cleanUpTasksForAbort.add(new DirectoryCleanUpTask(identity, targetPath, true)); + createDirectory(fileSystem, targetPath); } }, fileSystemExecutor)); String partitionName = getPartitionName(partition.getDatabaseName(), partition.getTableName(), partition.getValues()); - partitionAdder.addPartition(new PartitionWithStatistics(partition, partitionName, partitionAndMore.getStatisticsUpdate())); + partitionAdder.addPartition(new PartitionWithStatistics(partition, partitionName, partitionAndMore.statisticsUpdate())); } - private void prepareInsertExistingPartition(HdfsContext hdfsContext, String queryId, PartitionAndMore partitionAndMore) + private void prepareInsertExistingPartition(ConnectorIdentity identity, String queryId, PartitionAndMore partitionAndMore) { deleteOnly = false; - Partition partition = partitionAndMore.getPartition(); + Partition partition = partitionAndMore.partition(); partitionsToInvalidate.add(partition); - Path targetPath = new Path(partition.getStorage().getLocation()); - Path currentPath = new Path(partitionAndMore.getCurrentLocation().toString()); - cleanUpTasksForAbort.add(new DirectoryCleanUpTask(hdfsContext, targetPath, false)); + Location targetPath = Location.of(partition.getStorage().getLocation()); + Location currentPath = partitionAndMore.currentLocation(); + cleanUpTasksForAbort.add(new DirectoryCleanUpTask(identity, targetPath, false)); if (!targetPath.equals(currentPath)) { - // if staging directory is used we cherry-pick files to be moved - asyncRename(hdfsEnvironment, fileSystemExecutor, fileSystemOperationsCancelled, fileSystemOperationFutures, hdfsContext, currentPath, targetPath, partitionAndMore.getFileNames()); + // if staging directory is used, we cherry-pick files to be moved + TrinoFileSystem fileSystem = fileSystemFactory.create(identity); + asyncRename(fileSystem, fileSystemExecutor, fileSystemOperationsCancelled, fileSystemOperationFutures, currentPath, targetPath, partitionAndMore.getFileNames()); } else { - // if we inserted directly into partition directory we need to remove extra output files which should not be part of the table - cleanExtraOutputFiles(hdfsContext, queryId, partitionAndMore); + // if we inserted directly into partition directory, we need to remove extra output files which should not be part of the table + cleanExtraOutputFiles(identity, queryId, partitionAndMore); } updateStatisticsOperations.add(new UpdateStatisticsOperation( partition.getSchemaTableName(), Optional.of(getPartitionName(partition.getDatabaseName(), partition.getTableName(), partition.getValues())), - partitionAndMore.getStatisticsUpdate(), + partitionAndMore.statisticsUpdate(), true)); } private void executeCleanupTasksForAbort(Collection declaredIntentionsToWrite) { Set queryIds = declaredIntentionsToWrite.stream() - .map(DeclaredIntentionToWrite::getQueryId) + .map(DeclaredIntentionToWrite::queryId) .collect(toImmutableSet()); for (DirectoryCleanUpTask cleanUpTask : cleanUpTasksForAbort) { - recursiveDeleteFilesAndLog(cleanUpTask.getContext(), cleanUpTask.getPath(), queryIds, cleanUpTask.isDeleteEmptyDirectory(), "temporary directory commit abort"); + recursiveDeleteFilesAndLog(cleanUpTask.identity(), cleanUpTask.location(), queryIds, cleanUpTask.deleteEmptyDirectory(), "temporary directory commit abort"); } } private void executeDeletionTasksForFinish() { for (DirectoryDeletionTask deletionTask : deletionTasksForFinish) { - if (!deleteRecursivelyIfExists(deletionTask.getContext(), hdfsEnvironment, deletionTask.getPath())) { - logCleanupFailure("Error deleting directory %s", deletionTask.getPath()); + TrinoFileSystem fileSystem = fileSystemFactory.create(deletionTask.identity()); + try { + fileSystem.deleteDirectory(deletionTask.location()); + } + catch (IOException e) { + logCleanupFailure(e, "Error deleting directory: %s", deletionTask.location()); } } } @@ -2076,12 +1968,13 @@ private void executeRenameTasksForAbort() try { // Ignore the task if the source directory doesn't exist. // This is probably because the original rename that we are trying to undo here never succeeded. - if (pathExists(directoryRenameTask.getContext(), hdfsEnvironment, directoryRenameTask.getRenameFrom())) { - renameDirectory(directoryRenameTask.getContext(), hdfsEnvironment, directoryRenameTask.getRenameFrom(), directoryRenameTask.getRenameTo(), () -> {}); + TrinoFileSystem fileSystem = fileSystemFactory.create(directoryRenameTask.identity()); + if (directoryExists(fileSystem, directoryRenameTask.renameFrom())) { + renameDirectory(fileSystem, directoryRenameTask.renameFrom(), directoryRenameTask.renameTo(), () -> {}); } } catch (Throwable throwable) { - logCleanupFailure(throwable, "failed to undo rename of partition directory: %s to %s", directoryRenameTask.getRenameFrom(), directoryRenameTask.getRenameTo()); + logCleanupFailure(throwable, "failed to undo rename of partition directory: %s to %s", directoryRenameTask.renameFrom(), directoryRenameTask.renameTo()); } } } @@ -2089,16 +1982,16 @@ private void executeRenameTasksForAbort() private void pruneAndDeleteStagingDirectories(List declaredIntentionsToWrite) { for (DeclaredIntentionToWrite declaredIntentionToWrite : declaredIntentionsToWrite) { - if (declaredIntentionToWrite.getMode() != WriteMode.STAGE_AND_MOVE_TO_TARGET_DIRECTORY) { + if (declaredIntentionToWrite.mode() != WriteMode.STAGE_AND_MOVE_TO_TARGET_DIRECTORY) { continue; } Set queryIds = declaredIntentionsToWrite.stream() - .map(DeclaredIntentionToWrite::getQueryId) + .map(DeclaredIntentionToWrite::queryId) .collect(toImmutableSet()); - Path path = declaredIntentionToWrite.getRootPath(); - recursiveDeleteFilesAndLog(declaredIntentionToWrite.getHdfsContext(), path, queryIds, true, "staging directory cleanup"); + Location path = declaredIntentionToWrite.rootPath(); + recursiveDeleteFilesAndLog(declaredIntentionToWrite.identity(), path, queryIds, true, "staging directory cleanup"); } } @@ -2276,7 +2169,7 @@ private void executeIrreversibleMetastoreOperations() } catch (Throwable t) { synchronized (failedIrreversibleOperationDescriptions) { - addSuppressedExceptions(suppressedExceptions, t, failedIrreversibleOperationDescriptions, irreversibleMetastoreOperation.getDescription()); + addSuppressedExceptions(suppressedExceptions, t, failedIrreversibleOperationDescriptions, irreversibleMetastoreOperation.description()); } } }, dropExecutor)); @@ -2303,45 +2196,50 @@ private void executeIrreversibleMetastoreOperations() } } + private static Optional getOptionalPartition(HiveMetastore metastore, SchemaTableName schemaTableName, List partitionValues) + { + return metastore.getTable(schemaTableName.getSchemaName(), schemaTableName.getTableName()) + .flatMap(table -> metastore.getPartition(table, partitionValues)); + } + @GuardedBy("this") private void rollbackShared() { checkHoldsLock(); for (DeclaredIntentionToWrite declaredIntentionToWrite : declaredIntentionsToWrite) { - switch (declaredIntentionToWrite.getMode()) { - case STAGE_AND_MOVE_TO_TARGET_DIRECTORY: - case DIRECT_TO_TARGET_NEW_DIRECTORY: - // For STAGE_AND_MOVE_TO_TARGET_DIRECTORY, there is no need to cleanup the target directory as + switch (declaredIntentionToWrite.mode()) { + case STAGE_AND_MOVE_TO_TARGET_DIRECTORY, DIRECT_TO_TARGET_NEW_DIRECTORY -> { + // For STAGE_AND_MOVE_TO_TARGET_DIRECTORY, there is no need to clean up the target directory as // it will only be written to during the commit call and the commit call cleans up after failures. - if ((declaredIntentionToWrite.getMode() == DIRECT_TO_TARGET_NEW_DIRECTORY) && skipTargetCleanupOnRollback) { + if ((declaredIntentionToWrite.mode() == DIRECT_TO_TARGET_NEW_DIRECTORY) && skipTargetCleanupOnRollback) { break; } - Path rootPath = declaredIntentionToWrite.getRootPath(); + Location rootPath = declaredIntentionToWrite.rootPath(); // In the case of DIRECT_TO_TARGET_NEW_DIRECTORY, if the directory is not guaranteed to be unique // for the query, it is possible that another query or compute engine may see the directory, wrote - // data to it, and exported it through metastore. Therefore it may be argued that cleanup of staging + // data to it, and exported it through metastore. Therefore, it may be argued that cleanup of staging // directories must be carried out conservatively. To be safe, we only delete files that start or // end with the query IDs in this transaction. recursiveDeleteFilesAndLog( - declaredIntentionToWrite.getHdfsContext(), + declaredIntentionToWrite.identity(), rootPath, - ImmutableSet.of(declaredIntentionToWrite.getQueryId()), + ImmutableSet.of(declaredIntentionToWrite.queryId()), true, - format("staging/target_new directory rollback for table %s", declaredIntentionToWrite.getSchemaTableName())); - break; - case DIRECT_TO_TARGET_EXISTING_DIRECTORY: - Set pathsToClean = new HashSet<>(); + format("staging/target_new directory rollback for table %s", declaredIntentionToWrite.schemaTableName())); + } + case DIRECT_TO_TARGET_EXISTING_DIRECTORY -> { + Set pathsToClean = new HashSet<>(); // Check the base directory of the declared intention // * existing partition may also be in this directory // * this is where new partitions are created - Path baseDirectory = declaredIntentionToWrite.getRootPath(); + Location baseDirectory = declaredIntentionToWrite.rootPath(); pathsToClean.add(baseDirectory); - SchemaTableName schemaTableName = declaredIntentionToWrite.getSchemaTableName(); + SchemaTableName schemaTableName = declaredIntentionToWrite.schemaTableName(); Optional
table = delegate.getTable(schemaTableName.getSchemaName(), schemaTableName.getTableName()); if (table.isPresent()) { // check every existing partition that is outside for the base directory @@ -2350,20 +2248,19 @@ private void rollbackShared() List partitionColumnNames = partitionColumns.stream() .map(Column::getName) .collect(toImmutableList()); - List partitionNames = delegate.getPartitionNamesByFilter( + List partitionNames = getOptionalPartitions( schemaTableName.getSchemaName(), schemaTableName.getTableName(), partitionColumnNames, TupleDomain.all()) .orElse(ImmutableList.of()); for (List partitionNameBatch : Iterables.partition(partitionNames, 10)) { - Collection> partitions = delegate.getPartitionsByNames(schemaTableName.getSchemaName(), schemaTableName.getTableName(), partitionNameBatch).values(); + Collection> partitions = getOptionalPartitions(schemaTableName.getSchemaName(), schemaTableName.getTableName(), partitionNameBatch).values(); partitions.stream() - .filter(Optional::isPresent) - .map(Optional::get) + .flatMap(Optional::stream) .map(partition -> partition.getStorage().getLocation()) - .map(Path::new) - .filter(path -> !isSameOrParent(baseDirectory, path)) + .filter(path -> !path.startsWith(baseDirectory.toString())) + .map(Location::of) .forEach(pathsToClean::add); } } @@ -2376,20 +2273,16 @@ private void rollbackShared() } // delete any file that starts or ends with the query ID - for (Path path : pathsToClean) { - // TODO: It is a known deficiency that some empty directory does not get cleaned up in S3. + for (Location path : pathsToClean) { // We cannot delete any of the directories here since we do not know who created them. recursiveDeleteFilesAndLog( - declaredIntentionToWrite.getHdfsContext(), + declaredIntentionToWrite.identity(), path, - ImmutableSet.of(declaredIntentionToWrite.getQueryId()), + ImmutableSet.of(declaredIntentionToWrite.queryId()), false, format("target_existing directory rollback for table %s", schemaTableName)); } - - break; - default: - throw new UnsupportedOperationException("Unknown write mode"); + } } } } @@ -2402,30 +2295,20 @@ public synchronized void testOnlyCheckIsReadOnly() } } - @VisibleForTesting - public void testOnlyThrowOnCleanupFailures() - { - throwOnCleanupFailure = true; - } - @GuardedBy("this") - private void checkReadable() + private synchronized void checkReadable() { checkHoldsLock(); switch (state) { - case EMPTY: - case SHARED_OPERATION_BUFFERED: - return; - case EXCLUSIVE_OPERATION_BUFFERED: - throw new TrinoException(NOT_SUPPORTED, "Unsupported combination of operations in a single transaction"); - case FINISHED: - throw new IllegalStateException("Tried to access metastore after transaction has been committed/aborted"); + case EMPTY, SHARED_OPERATION_BUFFERED -> { /* nothing to do */ } + case EXCLUSIVE_OPERATION_BUFFERED -> throw new TrinoException(NOT_SUPPORTED, "Unsupported combination of operations in a single transaction"); + case FINISHED -> throw new IllegalStateException("Tried to access metastore after transaction has been committed/aborted"); } } @GuardedBy("this") - private void setShared() + private synchronized void setShared() { checkHoldsLock(); @@ -2434,7 +2317,7 @@ private void setShared() } @GuardedBy("this") - private void setExclusive(ExclusiveOperation exclusiveOperation) + private synchronized void setExclusive(ExclusiveOperation exclusiveOperation) { checkHoldsLock(); @@ -2456,34 +2339,15 @@ private void checkNoPartitionAction(String databaseName, String tableName) } } - private static boolean isSameOrParent(Path parent, Path child) - { - int parentDepth = parent.depth(); - int childDepth = child.depth(); - if (parentDepth > childDepth) { - return false; - } - for (int i = childDepth; i > parentDepth; i--) { - child = child.getParent(); - } - return parent.equals(child); - } - @FormatMethod - private void logCleanupFailure(String format, Object... args) + private static void logCleanupFailure(String format, Object... args) { - if (throwOnCleanupFailure) { - throw new RuntimeException(format(format, args)); - } log.warn(format, args); } @FormatMethod - private void logCleanupFailure(Throwable t, String format, Object... args) + private static void logCleanupFailure(Throwable t, String format, Object... args) { - if (throwOnCleanupFailure) { - throw new RuntimeException(format(format, args), t); - } log.warn(t, format, args); } @@ -2497,37 +2361,23 @@ private static void addSuppressedExceptions(List suppressedExceptions } private static void asyncRename( - HdfsEnvironment hdfsEnvironment, + TrinoFileSystem fileSystem, Executor executor, AtomicBoolean cancelled, List> fileRenameFutures, - HdfsContext context, - Path currentPath, - Path targetPath, + Location currentPath, + Location targetPath, List fileNames) { - FileSystem fileSystem; - try { - fileSystem = hdfsEnvironment.getFileSystem(context, currentPath); - } - catch (IOException e) { - throw new TrinoException(HIVE_FILESYSTEM_ERROR, format("Error moving data files to final location. Error listing directory %s", currentPath), e); - } - for (String fileName : fileNames) { - Path source = new Path(currentPath, fileName); - Path target = new Path(targetPath, fileName); + Location source = currentPath.appendPath(fileName); + Location target = targetPath.appendPath(fileName); fileRenameFutures.add(CompletableFuture.runAsync(() -> { if (cancelled.get()) { return; } try { - if (fileSystem.exists(target)) { - throw new TrinoException(HIVE_FILESYSTEM_ERROR, format("Error moving data files from %s to final location %s: target location already exists", source, target)); - } - if (!fileSystem.rename(source, target)) { - throw new TrinoException(HIVE_FILESYSTEM_ERROR, format("Error moving data files from %s to final location %s: rename not successful", source, target)); - } + fileSystem.renameFile(source, target); } catch (IOException e) { throw new TrinoException(HIVE_FILESYSTEM_ERROR, format("Error moving data files from %s to final location %s", source, target), e); @@ -2536,22 +2386,21 @@ private static void asyncRename( } } - private void recursiveDeleteFilesAndLog(HdfsContext context, Path directory, Set queryIds, boolean deleteEmptyDirectories, String reason) + private void recursiveDeleteFilesAndLog(ConnectorIdentity identity, Location directory, Set queryIds, boolean deleteEmptyDirectories, String reason) { RecursiveDeleteResult recursiveDeleteResult = recursiveDeleteFiles( - hdfsEnvironment, - context, + fileSystemFactory.create(identity), directory, queryIds, deleteEmptyDirectories); - if (!recursiveDeleteResult.getNotDeletedEligibleItems().isEmpty()) { + if (!recursiveDeleteResult.notDeletedEligibleItems().isEmpty()) { logCleanupFailure( "Error deleting directory %s for %s. Some eligible items cannot be deleted: %s.", directory.toString(), reason, - recursiveDeleteResult.getNotDeletedEligibleItems()); + recursiveDeleteResult.notDeletedEligibleItems()); } - else if (deleteEmptyDirectories && !recursiveDeleteResult.isDirectoryNoLongerExists()) { + else if (deleteEmptyDirectories && !recursiveDeleteResult.directoryNoLongerExists()) { logCleanupFailure( "Error deleting directory %s for %s. Cannot delete the directory.", directory.toString(), @@ -2574,13 +2423,10 @@ else if (deleteEmptyDirectories && !recursiveDeleteResult.isDirectoryNoLongerExi * @param queryIds prefix or suffix of files that should be deleted * @param deleteEmptyDirectories whether empty directories should be deleted */ - private static RecursiveDeleteResult recursiveDeleteFiles(HdfsEnvironment hdfsEnvironment, HdfsContext context, Path directory, Set queryIds, boolean deleteEmptyDirectories) + private static RecursiveDeleteResult recursiveDeleteFiles(TrinoFileSystem fileSystem, Location directory, Set queryIds, boolean deleteEmptyDirectories) { - FileSystem fileSystem; try { - fileSystem = hdfsEnvironment.getFileSystem(context, directory); - - if (!fileSystem.exists(directory)) { + if (!fileSystem.directoryExists(directory).orElse(false)) { return new RecursiveDeleteResult(true, ImmutableList.of()); } } @@ -2593,16 +2439,30 @@ private static RecursiveDeleteResult recursiveDeleteFiles(HdfsEnvironment hdfsEn return doRecursiveDeleteFiles(fileSystem, directory, queryIds, deleteEmptyDirectories); } - private static RecursiveDeleteResult doRecursiveDeleteFiles(FileSystem fileSystem, Path directory, Set queryIds, boolean deleteEmptyDirectories) + private static RecursiveDeleteResult doRecursiveDeleteFiles(TrinoFileSystem fileSystem, Location directory, Set queryIds, boolean deleteEmptyDirectories) { // don't delete hidden Trino directories use by FileHiveMetastore - if (directory.getName().startsWith(".trino")) { + directory = asFileLocation(directory); + if (directory.fileName().startsWith(".trino")) { return new RecursiveDeleteResult(false, ImmutableList.of()); } - FileStatus[] allFiles; + // TODO: this lists recursively but only uses the first level + List allFiles = new ArrayList<>(); + Set allDirectories; try { - allFiles = fileSystem.listStatus(directory); + FileIterator iterator = fileSystem.listFiles(directory); + while (iterator.hasNext()) { + Location location = iterator.next().location(); + String child = location.toString().substring(directory.toString().length()); + while (child.startsWith("/")) { + child = child.substring(1); + } + if (!child.contains("/")) { + allFiles.add(location); + } + } + allDirectories = fileSystem.listDirectories(directory); } catch (IOException e) { ImmutableList.Builder notDeletedItems = ImmutableList.builder(); @@ -2612,44 +2472,37 @@ private static RecursiveDeleteResult doRecursiveDeleteFiles(FileSystem fileSyste boolean allDescendentsDeleted = true; ImmutableList.Builder notDeletedEligibleItems = ImmutableList.builder(); - for (FileStatus fileStatus : allFiles) { - if (fileStatus.isFile()) { - Path filePath = fileStatus.getPath(); - String fileName = filePath.getName(); - boolean eligible = false; - // don't delete hidden Trino directories use by FileHiveMetastore - if (!fileName.startsWith(".trino")) { - eligible = queryIds.stream().anyMatch(id -> isFileCreatedByQuery(fileName, id)); - } - if (eligible) { - if (!deleteIfExists(fileSystem, filePath, false)) { - allDescendentsDeleted = false; - notDeletedEligibleItems.add(filePath.toString()); - } - } - else { + for (Location file : allFiles) { + String fileName = file.fileName(); + boolean eligible = false; + // don't delete hidden Trino directories use by FileHiveMetastore + if (!fileName.startsWith(".trino")) { + eligible = queryIds.stream().anyMatch(id -> isFileCreatedByQuery(fileName, id)); + } + if (eligible) { + if (!deleteFileIfExists(fileSystem, file)) { allDescendentsDeleted = false; - } - } - else if (fileStatus.isDirectory()) { - RecursiveDeleteResult subResult = doRecursiveDeleteFiles(fileSystem, fileStatus.getPath(), queryIds, deleteEmptyDirectories); - if (!subResult.isDirectoryNoLongerExists()) { - allDescendentsDeleted = false; - } - if (!subResult.getNotDeletedEligibleItems().isEmpty()) { - notDeletedEligibleItems.addAll(subResult.getNotDeletedEligibleItems()); + notDeletedEligibleItems.add(file.toString()); } } else { allDescendentsDeleted = false; - notDeletedEligibleItems.add(fileStatus.getPath().toString()); + } + } + for (Location file : allDirectories) { + RecursiveDeleteResult subResult = doRecursiveDeleteFiles(fileSystem, file, queryIds, deleteEmptyDirectories); + if (!subResult.directoryNoLongerExists()) { + allDescendentsDeleted = false; + } + if (!subResult.notDeletedEligibleItems().isEmpty()) { + notDeletedEligibleItems.addAll(subResult.notDeletedEligibleItems()); } } // Unconditionally delete empty delta_ and delete_delta_ directories, because that's // what Hive does, and leaving them in place confuses delta file readers. - if (allDescendentsDeleted && (deleteEmptyDirectories || DELTA_DIRECTORY_MATCHER.matcher(directory.getName()).matches())) { + if (allDescendentsDeleted && (deleteEmptyDirectories || isDeltaDirectory(directory))) { verify(notDeletedEligibleItems.build().isEmpty()); - if (!deleteIfExists(fileSystem, directory, false)) { + if (!deleteEmptyDirectoryIfExists(fileSystem, directory)) { return new RecursiveDeleteResult(false, ImmutableList.of(directory + "/")); } return new RecursiveDeleteResult(true, ImmutableList.of()); @@ -2657,59 +2510,54 @@ else if (fileStatus.isDirectory()) { return new RecursiveDeleteResult(false, notDeletedEligibleItems.build()); } - /** - * Attempts to remove the file or empty directory. - * - * @return true if the location no longer exists - */ - private static boolean deleteIfExists(FileSystem fileSystem, Path path, boolean recursive) + private static boolean isDeltaDirectory(Location directory) { - try { - // attempt to delete the path - if (fileSystem.delete(path, recursive)) { - return true; - } + return DELTA_DIRECTORY_MATCHER.matcher(asFileLocation(directory).fileName()).matches(); + } - // delete failed - // check if path still exists - return !fileSystem.exists(path); + private static boolean deleteFileIfExists(TrinoFileSystem fileSystem, Location location) + { + try { + fileSystem.deleteFile(location); + return true; } - catch (FileNotFoundException ignored) { - // path was already removed or never existed + catch (FileNotFoundException e) { return true; } - catch (IOException ignored) { + catch (IOException e) { + return false; } - return false; } - /** - * Attempts to remove the file or empty directory. - * - * @return true if the location no longer exists - */ - private static boolean deleteRecursivelyIfExists(HdfsContext context, HdfsEnvironment hdfsEnvironment, Path path) + private static boolean deleteEmptyDirectoryIfExists(TrinoFileSystem fileSystem, Location location) { - FileSystem fileSystem; try { - fileSystem = hdfsEnvironment.getFileSystem(context, path); + if (fileSystem.listFiles(location).hasNext()) { + log.warn("Not deleting non-empty directory: %s", location); + return false; + } + fileSystem.deleteDirectory(location); + return true; } - catch (IOException ignored) { - return false; + catch (IOException e) { + try { + return !fileSystem.directoryExists(location).orElse(false); + } + catch (IOException ex) { + return false; + } } - - return deleteIfExists(fileSystem, path, true); } - private static void renameDirectory(HdfsContext context, HdfsEnvironment hdfsEnvironment, Path source, Path target, Runnable runWhenPathDoesntExist) + private static void renameDirectory(TrinoFileSystem fileSystem, Location source, Location target, Runnable runWhenPathDoesntExist) { - if (pathExists(context, hdfsEnvironment, target)) { - throw new TrinoException(HIVE_PATH_ALREADY_EXISTS, - format("Unable to rename from %s to %s: target directory already exists", source, target)); + if (directoryExists(fileSystem, target)) { + throw new TrinoException(HIVE_PATH_ALREADY_EXISTS, format("Unable to rename from %s to %s: target directory already exists", source, target)); } - if (!pathExists(context, hdfsEnvironment, target.getParent())) { - createDirectory(context, hdfsEnvironment, target.getParent()); + Location parent = asFileLocation(target).parentDirectory(); + if (!directoryExists(fileSystem, parent)) { + createDirectory(fileSystem, parent); } // The runnable will assume that if rename fails, it will be okay to delete the directory (if the directory is empty). @@ -2717,34 +2565,62 @@ private static void renameDirectory(HdfsContext context, HdfsEnvironment hdfsEnv runWhenPathDoesntExist.run(); try { - if (!hdfsEnvironment.getFileSystem(context, source).rename(source, target)) { - throw new TrinoException(HIVE_FILESYSTEM_ERROR, format("Failed to rename %s to %s: rename returned false", source, target)); - } + fileSystem.renameDirectory(source, target); } catch (IOException e) { throw new TrinoException(HIVE_FILESYSTEM_ERROR, format("Failed to rename %s to %s", source, target), e); } } + private static void createDirectory(TrinoFileSystem fileSystem, Location directory) + { + try { + fileSystem.createDirectory(directory); + } + catch (IOException e) { + throw new TrinoException(HIVE_FILESYSTEM_ERROR, e); + } + } + + private static boolean directoryExists(TrinoFileSystem fileSystem, Location directory) + { + try { + return fileSystem.directoryExists(directory).orElse(false); + } + catch (IOException e) { + throw new TrinoException(HIVE_FILESYSTEM_ERROR, e); + } + } + private static Optional getQueryId(Database database) { - return Optional.ofNullable(database.getParameters().get(PRESTO_QUERY_ID_NAME)); + return Optional.ofNullable(database.getParameters().get(TRINO_QUERY_ID_NAME)); } private static Optional getQueryId(Table table) { - return Optional.ofNullable(table.getParameters().get(PRESTO_QUERY_ID_NAME)); + return Optional.ofNullable(table.getParameters().get(TRINO_QUERY_ID_NAME)); } private static Optional getQueryId(Partition partition) { - return Optional.ofNullable(partition.getParameters().get(PRESTO_QUERY_ID_NAME)); + return Optional.ofNullable(partition.getParameters().get(TRINO_QUERY_ID_NAME)); + } + + private static Location asFileLocation(Location location) + { + // TODO: this is to work around the file-only restriction of Location methods + String value = location.toString(); + while (value.endsWith("/")) { + value = value.substring(0, value.length() - 1); + } + return Location.of(value); } private void checkHoldsLock() { // This method serves a similar purpose at runtime as GuardedBy on method serves during static analysis. - // This method should not have significant performance impact. If it does, it may be reasonably to remove this method. + // This method should not have a significant performance impact. If it does, it may be reasonably to remove this method. // This intentionally does not use checkState. if (!Thread.holdsLock(this)) { throw new IllegalStateException(format("Thread must hold a lock on the %s", getClass().getSimpleName())); @@ -2776,64 +2652,34 @@ private enum TableSource // RECREATED_IN_THIS_TRANSACTION is a possible case, but it is not supported with the current implementation } - public static class Action + private record Action(ActionType type, T data, ConnectorIdentity identity, String queryId) { - private final ActionType type; - private final T data; - private final HdfsContext hdfsContext; - private final String queryId; - - public Action(ActionType type, T data, HdfsContext hdfsContext, String queryId) + private Action { - this.type = requireNonNull(type, "type is null"); + requireNonNull(type, "type is null"); if (type == ActionType.DROP || type == ActionType.DROP_PRESERVE_DATA) { checkArgument(data == null, "data is not null"); } else { requireNonNull(data, "data is null"); } - this.data = data; - this.hdfsContext = requireNonNull(hdfsContext, "hdfsContext is null"); - this.queryId = requireNonNull(queryId, "queryId is null"); + requireNonNull(identity, "identity is null"); + requireNonNull(queryId, "queryId is null"); } - public ActionType getType() - { - return type; - } - - public T getData() + @Override + public T data() { checkState(type != ActionType.DROP); return data; } - - public HdfsContext getHdfsContext() - { - return hdfsContext; - } - - public String getQueryId() - { - return queryId; - } - - @Override - public String toString() - { - return toStringHelper(this) - .add("type", type) - .add("queryId", queryId) - .add("data", data) - .toString(); - } } private static class TableAndMore { private final Table table; private final Optional principalPrivileges; - private final Optional currentLocation; // unpartitioned table only + private final Optional currentLocation; // unpartitioned table only private final Optional> fileNames; private final boolean ignoreExisting; private final PartitionStatistics statistics; @@ -2843,7 +2689,7 @@ private static class TableAndMore public TableAndMore( Table table, Optional principalPrivileges, - Optional currentLocation, + Optional currentLocation, Optional> fileNames, boolean ignoreExisting, PartitionStatistics statistics, @@ -2879,7 +2725,7 @@ public PrincipalPrivileges getPrincipalPrivileges() return principalPrivileges.get(); } - public Optional getCurrentLocation() + public Optional getCurrentLocation() { return currentLocation; } @@ -2924,18 +2770,11 @@ private static class TableAndMergeResults extends TableAndMore { private final List partitionMergeResults; - private final List partitions; - public TableAndMergeResults(Table table, Optional principalPrivileges, Optional currentLocation, List partitionMergeResults, List partitions) + public TableAndMergeResults(Table table, Optional principalPrivileges, Optional currentLocation, List partitionMergeResults) { super(table, principalPrivileges, currentLocation, Optional.empty(), false, PartitionStatistics.empty(), PartitionStatistics.empty(), false); // retries are not supported for transactional tables this.partitionMergeResults = requireNonNull(partitionMergeResults, "partitionMergeResults is null"); - this.partitions = requireNonNull(partitions, "partitions is nul"); - } - - public List getPartitions() - { - return partitions; } @Override @@ -2944,40 +2783,27 @@ public String toString() return toStringHelper(this) .add("table", getTable()) .add("partitionMergeResults", partitionMergeResults) - .add("partitions", partitions) .add("principalPrivileges", getPrincipalPrivileges()) .add("currentLocation", getCurrentLocation()) .toString(); } } - private static class PartitionAndMore + private record PartitionAndMore( + Partition partition, + Location currentLocation, + Optional> fileNames, + PartitionStatistics statistics, + PartitionStatistics statisticsUpdate, + boolean cleanExtraOutputFilesOnCommit) { - private final Partition partition; - private final Location currentLocation; - private final Optional> fileNames; - private final PartitionStatistics statistics; - private final PartitionStatistics statisticsUpdate; - private final boolean cleanExtraOutputFilesOnCommit; - - public PartitionAndMore(Partition partition, Location currentLocation, Optional> fileNames, PartitionStatistics statistics, PartitionStatistics statisticsUpdate, boolean cleanExtraOutputFilesOnCommit) + private PartitionAndMore { - this.partition = requireNonNull(partition, "partition is null"); - this.currentLocation = requireNonNull(currentLocation, "currentLocation is null"); - this.fileNames = requireNonNull(fileNames, "fileNames is null"); - this.statistics = requireNonNull(statistics, "statistics is null"); - this.statisticsUpdate = requireNonNull(statisticsUpdate, "statisticsUpdate is null"); - this.cleanExtraOutputFilesOnCommit = cleanExtraOutputFilesOnCommit; - } - - public Partition getPartition() - { - return partition; - } - - public Location getCurrentLocation() - { - return currentLocation; + requireNonNull(partition, "partition is null"); + requireNonNull(currentLocation, "currentLocation is null"); + requireNonNull(fileNames, "fileNames is null"); + requireNonNull(statistics, "statistics is null"); + requireNonNull(statisticsUpdate, "statisticsUpdate is null"); } public List getFileNames() @@ -2991,21 +2817,6 @@ public boolean hasFileNames() return fileNames.isPresent(); } - public PartitionStatistics getStatistics() - { - return statistics; - } - - public PartitionStatistics getStatisticsUpdate() - { - return statisticsUpdate; - } - - public boolean isCleanExtraOutputFilesOnCommit() - { - return cleanExtraOutputFilesOnCommit; - } - public Partition getAugmentedPartitionForInTransactionRead() { // This method augments the location field of the partition to the staging location. @@ -3020,204 +2831,55 @@ public Partition getAugmentedPartitionForInTransactionRead() } return partition; } - - @Override - public String toString() - { - return toStringHelper(this) - .add("partition", partition) - .add("currentLocation", currentLocation) - .add("fileNames", fileNames) - .add("cleanExtraOutputFilesOnCommit", cleanExtraOutputFilesOnCommit) - .toString(); - } } - private static class DeclaredIntentionToWrite + public record DeclaredIntentionToWrite(String declarationId, WriteMode mode, ConnectorIdentity identity, String queryId, Location rootPath, SchemaTableName schemaTableName) { - private final String declarationId; - private final WriteMode mode; - private final HdfsContext hdfsContext; - private final String queryId; - private final Path rootPath; - private final SchemaTableName schemaTableName; - - public DeclaredIntentionToWrite(String declarationId, WriteMode mode, HdfsContext hdfsContext, String queryId, Path stagingPathRoot, SchemaTableName schemaTableName) - { - this.declarationId = requireNonNull(declarationId, "declarationId is null"); - this.mode = requireNonNull(mode, "mode is null"); - this.hdfsContext = requireNonNull(hdfsContext, "hdfsContext is null"); - this.queryId = requireNonNull(queryId, "queryId is null"); - this.rootPath = requireNonNull(stagingPathRoot, "stagingPathRoot is null"); - this.schemaTableName = requireNonNull(schemaTableName, "schemaTableName is null"); - } - - public String getDeclarationId() - { - return declarationId; - } - - public WriteMode getMode() - { - return mode; - } - - public HdfsContext getHdfsContext() - { - return hdfsContext; - } - - public String getQueryId() - { - return queryId; - } - - public Path getRootPath() - { - return rootPath; - } - - public SchemaTableName getSchemaTableName() - { - return schemaTableName; - } - - @Override - public String toString() + public DeclaredIntentionToWrite { - return toStringHelper(this) - .add("mode", mode) - .add("hdfsContext", hdfsContext) - .add("queryId", queryId) - .add("rootPath", rootPath) - .add("schemaTableName", schemaTableName) - .toString(); + requireNonNull(declarationId, "declarationId is null"); + requireNonNull(mode, "mode is null"); + requireNonNull(identity, "identity is null"); + requireNonNull(queryId, "queryId is null"); + requireNonNull(rootPath, "rootPath is null"); + requireNonNull(schemaTableName, "schemaTableName is null"); } } - private static class DirectoryCleanUpTask + private record DirectoryCleanUpTask(ConnectorIdentity identity, Location location, boolean deleteEmptyDirectory) { - private final HdfsContext context; - private final Path path; - private final boolean deleteEmptyDirectory; - - public DirectoryCleanUpTask(HdfsContext context, Path path, boolean deleteEmptyDirectory) - { - this.context = context; - this.path = path; - this.deleteEmptyDirectory = deleteEmptyDirectory; - } - - public HdfsContext getContext() - { - return context; - } - - public Path getPath() - { - return path; - } - - public boolean isDeleteEmptyDirectory() - { - return deleteEmptyDirectory; - } - - @Override - public String toString() + private DirectoryCleanUpTask { - return toStringHelper(this) - .add("context", context) - .add("path", path) - .add("deleteEmptyDirectory", deleteEmptyDirectory) - .toString(); + requireNonNull(identity, "identity is null"); + requireNonNull(location, "location is null"); } } - private static class DirectoryDeletionTask + private record DirectoryDeletionTask(ConnectorIdentity identity, Location location) { - private final HdfsContext context; - private final Path path; - - public DirectoryDeletionTask(HdfsContext context, Path path) - { - this.context = context; - this.path = path; - } - - public HdfsContext getContext() - { - return context; - } - - public Path getPath() + private DirectoryDeletionTask { - return path; - } - - @Override - public String toString() - { - return toStringHelper(this) - .add("context", context) - .add("path", path) - .toString(); + requireNonNull(identity, "identity is null"); + requireNonNull(location, "location is null"); } } - private static class DirectoryRenameTask + private record DirectoryRenameTask(ConnectorIdentity identity, Location renameFrom, Location renameTo) { - private final HdfsContext context; - private final Path renameFrom; - private final Path renameTo; - - public DirectoryRenameTask(HdfsContext context, Path renameFrom, Path renameTo) - { - this.context = requireNonNull(context, "context is null"); - this.renameFrom = requireNonNull(renameFrom, "renameFrom is null"); - this.renameTo = requireNonNull(renameTo, "renameTo is null"); - } - - public HdfsContext getContext() + private DirectoryRenameTask { - return context; - } - - public Path getRenameFrom() - { - return renameFrom; - } - - public Path getRenameTo() - { - return renameTo; - } - - @Override - public String toString() - { - return toStringHelper(this) - .add("context", context) - .add("renameFrom", renameFrom) - .add("renameTo", renameTo) - .toString(); + requireNonNull(identity, "identity is null"); + requireNonNull(renameFrom, "renameFrom is null"); + requireNonNull(renameTo, "renameTo is null"); } } - private static class IrreversibleMetastoreOperation + private record IrreversibleMetastoreOperation(String description, Runnable action) { - private final String description; - private final Runnable action; - - public IrreversibleMetastoreOperation(String description, Runnable action) - { - this.description = requireNonNull(description, "description is null"); - this.action = requireNonNull(action, "action is null"); - } - - public String getDescription() + private IrreversibleMetastoreOperation { - return description; + requireNonNull(description, "description is null"); + requireNonNull(action, "action is null"); } public void run() @@ -3250,7 +2912,7 @@ public String getDescription() return format("add table %s.%s", newTable.getDatabaseName(), newTable.getTableName()); } - public void run(HiveMetastoreClosure metastore, AcidTransaction transaction) + public void run(HiveMetastore metastore, AcidTransaction transaction) { boolean created = false; try { @@ -3258,7 +2920,7 @@ public void run(HiveMetastoreClosure metastore, AcidTransaction transaction) created = true; } catch (RuntimeException e) { - boolean done = false; + RuntimeException failure = e; try { Optional
existingTable = metastore.getTable(newTable.getDatabaseName(), newTable.getTableName()); if (existingTable.isPresent()) { @@ -3266,38 +2928,40 @@ public void run(HiveMetastoreClosure metastore, AcidTransaction transaction) Optional existingTableQueryId = getQueryId(table); if (existingTableQueryId.isPresent() && existingTableQueryId.get().equals(queryId)) { // ignore table if it was already created by the same query during retries - done = true; + failure = null; created = true; } else { - // If the table definition in the metastore is different than what this tx wants to create + // If the table definition in the metastore is different from what this tx wants to create, // then there is a conflict (e.g., current tx wants to create T(a: bigint), // but another tx already created T(a: varchar)). // This may be a problem if there is an insert after this step. if (!hasTheSameSchema(newTable, table)) { - e = new TrinoException(TRANSACTION_CONFLICT, format("Table already exists with a different schema: '%s'", newTable.getTableName())); + // produce an understandable error message + failure = new TrinoException(TRANSACTION_CONFLICT, format("Table already exists with a different schema: '%s'", newTable.getTableName())); } - else { - done = ignoreExisting; + else if (ignoreExisting) { + // if the statement is "CREATE TABLE IF NOT EXISTS", then ignore the exception + failure = null; } } } } - catch (RuntimeException ignored) { + catch (RuntimeException ignore) { // When table could not be fetched from metastore, it is not known whether the table was added. - // Deleting the table when aborting commit has the risk of deleting table not added in this transaction. + // Deleting the table when aborting commit has the risk of deleting a table not added in this transaction. // Not deleting the table may leave garbage behind. The former is much more dangerous than the latter. // Therefore, the table is not considered added. } - if (!done) { - throw e; + if (failure != null) { + throw failure; } } tableCreated = true; - if (created && !isPrestoView(newTable)) { - metastore.updateTableStatistics(newTable.getDatabaseName(), newTable.getTableName(), transaction, ignored -> statistics); + if (created && !isTrinoView(newTable) && !isTrinoMaterializedView(newTable)) { + metastore.updateTableStatistics(newTable.getDatabaseName(), newTable.getTableName(), transaction.getOptionalWriteId(), OVERWRITE_ALL, statistics); } } @@ -3320,7 +2984,7 @@ private static boolean hasTheSameSchema(Table newTable, Table existingTable) return true; } - public void undo(HiveMetastoreClosure metastore) + public void undo(HiveMetastore metastore) { if (!tableCreated) { return; @@ -3353,18 +3017,18 @@ public String getDescription() newTable.getTableName()); } - public void run(HiveMetastoreClosure metastore, AcidTransaction transaction) + public void run(HiveMetastore metastore, AcidTransaction transaction) { undo = true; if (transaction.isTransactional()) { metastore.alterTransactionalTable(newTable, transaction.getAcidTransactionId(), transaction.getWriteId(), principalPrivileges); } else { - metastore.replaceTable(newTable.getDatabaseName(), newTable.getTableName(), newTable, principalPrivileges); + metastore.replaceTable(newTable.getDatabaseName(), newTable.getTableName(), newTable, principalPrivileges, ImmutableMap.of()); } } - public void undo(HiveMetastoreClosure metastore, AcidTransaction transaction) + public void undo(HiveMetastore metastore, AcidTransaction transaction) { if (!undo) { return; @@ -3374,7 +3038,7 @@ public void undo(HiveMetastoreClosure metastore, AcidTransaction transaction) metastore.alterTransactionalTable(oldTable, transaction.getAcidTransactionId(), transaction.getWriteId(), principalPrivileges); } else { - metastore.replaceTable(oldTable.getDatabaseName(), oldTable.getTableName(), oldTable, principalPrivileges); + metastore.replaceTable(oldTable.getDatabaseName(), oldTable.getTableName(), oldTable, principalPrivileges, ImmutableMap.of()); } } } @@ -3403,13 +3067,13 @@ public String getDescription() newPartition.getPartition().getValues()); } - public void run(HiveMetastoreClosure metastore) + public void run(HiveMetastore metastore) { undo = true; metastore.alterPartition(newPartition.getPartition().getDatabaseName(), newPartition.getPartition().getTableName(), newPartition); } - public void undo(HiveMetastoreClosure metastore) + public void undo(HiveMetastore metastore) { if (!undo) { return; @@ -3435,27 +3099,36 @@ public UpdateStatisticsOperation(SchemaTableName tableName, Optional par this.merge = merge; } - public void run(HiveMetastoreClosure metastore, AcidTransaction transaction) + public void run(HiveMetastore metastore, AcidTransaction transaction) { + StatisticsUpdateMode mode = merge ? MERGE_INCREMENTAL : OVERWRITE_ALL; if (partitionName.isPresent()) { - metastore.updatePartitionStatistics(tableName.getSchemaName(), tableName.getTableName(), partitionName.get(), this::updateStatistics); + metastore.updatePartitionStatistics( + metastore.getTable(tableName.getSchemaName(), tableName.getTableName()) + .orElseThrow(() -> new TableNotFoundException(tableName)), + mode, + ImmutableMap.of(partitionName.get(), statistics)); } else { - metastore.updateTableStatistics(tableName.getSchemaName(), tableName.getTableName(), transaction, this::updateStatistics); + metastore.updateTableStatistics(tableName.getSchemaName(), tableName.getTableName(), transaction.getOptionalWriteId(), mode, statistics); } done = true; } - public void undo(HiveMetastoreClosure metastore, AcidTransaction transaction) + public void undo(HiveMetastore metastore, AcidTransaction transaction) { if (!done) { return; } if (partitionName.isPresent()) { - metastore.updatePartitionStatistics(tableName.getSchemaName(), tableName.getTableName(), partitionName.get(), this::resetStatistics); + metastore.updatePartitionStatistics( + metastore.getTable(tableName.getSchemaName(), tableName.getTableName()) + .orElseThrow(() -> new TableNotFoundException(tableName)), + UNDO_MERGE_INCREMENTAL, + ImmutableMap.of(partitionName.get(), statistics)); } else { - metastore.updateTableStatistics(tableName.getSchemaName(), tableName.getTableName(), transaction, this::resetStatistics); + metastore.updateTableStatistics(tableName.getSchemaName(), tableName.getTableName(), transaction.getOptionalWriteId(), UNDO_MERGE_INCREMENTAL, statistics); } } @@ -3466,28 +3139,18 @@ public String getDescription() } return format("replace table parameters %s", tableName); } - - private PartitionStatistics updateStatistics(PartitionStatistics currentStatistics) - { - return merge ? merge(currentStatistics, statistics) : statistics; - } - - private PartitionStatistics resetStatistics(PartitionStatistics currentStatistics) - { - return new PartitionStatistics(reduce(currentStatistics.getBasicStatistics(), statistics.getBasicStatistics(), SUBTRACT), ImmutableMap.of()); - } } private static class PartitionAdder { private final String schemaName; private final String tableName; - private final HiveMetastoreClosure metastore; + private final HiveMetastore metastore; private final int batchSize; private final List partitions; private List> createdPartitionValues = new ArrayList<>(); - public PartitionAdder(String schemaName, String tableName, HiveMetastoreClosure metastore, int batchSize) + public PartitionAdder(String schemaName, String tableName, HiveMetastore metastore, int batchSize) { this.schemaName = schemaName; this.tableName = tableName; @@ -3528,7 +3191,7 @@ public void execute(AcidTransaction transaction) boolean batchCompletelyAdded = true; for (PartitionWithStatistics partition : batch) { try { - Optional remotePartition = metastore.getPartition(schemaName, tableName, partition.getPartition().getValues()); + Optional remotePartition = getOptionalPartition(metastore, new SchemaTableName(schemaName, tableName), partition.getPartition().getValues()); // getQueryId(partition) is guaranteed to be non-empty. It is asserted in PartitionAdder.addPartition. if (remotePartition.isPresent() && getQueryId(remotePartition.get()).equals(getQueryId(partition.getPartition()))) { createdPartitionValues.add(partition.getPartition().getValues()); @@ -3537,7 +3200,7 @@ public void execute(AcidTransaction transaction) batchCompletelyAdded = false; } } - catch (Throwable ignored) { + catch (Throwable ignore) { // When partition could not be fetched from metastore, it is not known whether the partition was added. // Deleting the partition when aborting commit has the risk of deleting partition not added in this transaction. // Not deleting the partition may leave garbage behind. The former is much more dangerous than the latter. @@ -3546,8 +3209,8 @@ public void execute(AcidTransaction transaction) } } // If all the partitions were added successfully, the add_partition operation was actually successful. - // For some reason, it threw an exception (communication failure, retry failure after communication failure, etc). - // But we would consider it successful anyways. + // For some reason, it threw an exception (communication failure, retry failure after communication failure, etc.). + // But we would consider it successful regardless. if (!batchCompletelyAdded) { if (t instanceof TableNotFoundException) { throw new TrinoException(HIVE_TABLE_DROPPED_DURING_QUERY, t); @@ -3557,7 +3220,7 @@ public void execute(AcidTransaction transaction) } } if (transaction.isAcidTransactionRunning()) { - List partitionNames = partitions.stream().map(PartitionWithStatistics::getPartitionName).collect(Collectors.toUnmodifiableList()); + List partitionNames = partitions.stream().map(PartitionWithStatistics::getPartitionName).toList(); metastore.addDynamicPartitions(schemaName, tableName, partitionNames, transaction.getAcidTransactionId(), transaction.getWriteId(), transaction.getOperation()); } partitions.clear(); @@ -3572,8 +3235,8 @@ public List> rollback() metastore.dropPartition(schemaName, tableName, createdPartitionValue, false); } catch (PartitionNotFoundException e) { - // Maybe some one deleted the partition we added. - // Anyways, we are good because the partition is not there anymore. + // Maybe someone deleted the partition we added. + // Anyway, we are good because the partition is not there anymore. } catch (Throwable t) { partitionsFailedToRollback.add(createdPartitionValue); @@ -3584,31 +3247,11 @@ public List> rollback() } } - private static class RecursiveDeleteResult - { - private final boolean directoryNoLongerExists; - private final List notDeletedEligibleItems; - - public RecursiveDeleteResult(boolean directoryNoLongerExists, List notDeletedEligibleItems) - { - this.directoryNoLongerExists = directoryNoLongerExists; - this.notDeletedEligibleItems = notDeletedEligibleItems; - } - - public boolean isDirectoryNoLongerExists() - { - return directoryNoLongerExists; - } - - public List getNotDeletedEligibleItems() - { - return notDeletedEligibleItems; - } - } + private record RecursiveDeleteResult(boolean directoryNoLongerExists, List notDeletedEligibleItems) {} private interface ExclusiveOperation { - void execute(HiveMetastoreClosure delegate, HdfsEnvironment hdfsEnvironment); + void execute(HiveMetastore delegate); } private long allocateWriteId(String dbName, String tableName, long transactionId) @@ -3622,85 +3265,51 @@ private void acquireTableWriteLock( long transactionId, String dbName, String tableName, - DataOperationType operation, + AcidOperation operation, boolean isPartitioned) { delegate.acquireTableWriteLock(transactionOwner, queryId, transactionId, dbName, tableName, operation, isPartitioned); } - public void updateTableWriteId(String dbName, String tableName, long transactionId, long writeId, OptionalLong rowCountChange) + private void updateTableWriteId(String dbName, String tableName, long transactionId, long writeId, OptionalLong rowCountChange) { delegate.updateTableWriteId(dbName, tableName, transactionId, writeId, rowCountChange); } - public void alterPartitions(String dbName, String tableName, List partitions, long writeId) - { - delegate.alterPartitions(dbName, tableName, partitions, writeId); - } - public void addDynamicPartitions(String dbName, String tableName, List partitionNames, long transactionId, long writeId, AcidOperation operation) { delegate.addDynamicPartitions(dbName, tableName, partitionNames, transactionId, writeId, operation); } - public void commitTransaction(long transactionId) + public static void cleanExtraOutputFiles(TrinoFileSystem fileSystem, String queryId, Location path, Set filesToKeep) { - delegate.commitTransaction(transactionId); - } - - public static void cleanExtraOutputFiles(HdfsEnvironment hdfsEnvironment, HdfsContext hdfsContext, String queryId, Location location, Set filesToKeep) - { - Path path = new Path(location.toString()); - List filesToDelete = new LinkedList<>(); + List filesToDelete = new ArrayList<>(); try { - log.debug("Deleting failed attempt files from %s for query %s", path, queryId); - FileSystem fileSystem = hdfsEnvironment.getFileSystem(hdfsContext, path); - if (!fileSystem.exists(path)) { - // directory may nat exit if no files were actually written - return; - } - - // files are written flat in a single directory so we do not need to list recursively - RemoteIterator iterator = fileSystem.listFiles(path, false); - while (iterator.hasNext()) { - Path file = iterator.next().getPath(); - if (isFileCreatedByQuery(file.getName(), queryId) && !filesToKeep.contains(file.getName())) { - filesToDelete.add(file.getName()); + Failsafe.with(DELETE_RETRY_POLICY).run(() -> { + log.debug("Deleting failed attempt files from %s for query %s", path, queryId); + + filesToDelete.clear(); + FileIterator iterator = fileSystem.listFiles(path); + while (iterator.hasNext()) { + Location file = iterator.next().location(); + if (isFileCreatedByQuery(file.fileName(), queryId) && !filesToKeep.contains(file.fileName())) { + filesToDelete.add(file); + } } - } - - ImmutableList.Builder deletedFilesBuilder = ImmutableList.builder(); - Iterator filesToDeleteIterator = filesToDelete.iterator(); - while (filesToDeleteIterator.hasNext()) { - String fileName = filesToDeleteIterator.next(); - Path filePath = new Path(path, fileName); - log.debug("Deleting failed attempt file %s for query %s", filePath, queryId); - DELETE_RETRY.run("delete " + filePath, () -> { - checkedDelete(fileSystem, filePath, false); - return null; - }); - deletedFilesBuilder.add(fileName); - filesToDeleteIterator.remove(); - } - List deletedFiles = deletedFilesBuilder.build(); - if (!deletedFiles.isEmpty()) { - log.info("Deleted failed attempt files %s from %s for query %s", deletedFiles, path, queryId); - } + log.debug("Found %s failed attempt file(s) to delete for query %s", filesToDelete.size(), queryId); + fileSystem.deleteFiles(filesToDelete); + }); } - catch (Exception e) { - if (e instanceof InterruptedException) { - Thread.currentThread().interrupt(); - } - // If we fail here query will be rolled back. The optimal outcome would be for rollback to complete successfully and clean up everything for query. - // Yet if we have problem here, probably rollback will also fail. + catch (FailsafeException e) { + // If we fail here, the query will be rolled back. The optimal outcome would be for rollback to complete successfully and clean up everything for the query. + // Yet if we have a problem here, probably rollback will also fail. // - // Thrown exception is listing files which we could not delete. So those can be cleaned up later by user manually. + // Thrown exception is listing files which we could not delete. So the user can clean up those later manually. // Note it is not a bullet-proof solution. - // The rollback routine will still fire and try to cleanup the changes query made. It will cleanup some, leave some behind probably. - // It is not obvious that if at this point user cleans up the failed attempt files the table would be in the expected state. + // The rollback routine will still fire and try to clean up any changes the query made. It will clean up some, and probably leave some behind. // - // Still we cannot do much better for non-transactional Hive tables. + // Still, we cannot do much better for non-transactional Hive tables. throw new TrinoException( HIVE_FILESYSTEM_ERROR, format("Error deleting failed retry attempt files from %s; remaining files %s; manual cleanup may be required", path, filesToDelete), diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/StatisticsUpdateMode.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/StatisticsUpdateMode.java new file mode 100644 index 000000000000..fcbc891d1389 --- /dev/null +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/StatisticsUpdateMode.java @@ -0,0 +1,322 @@ +/* + * 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.hive.metastore; + +import com.google.common.collect.ImmutableMap; +import io.trino.plugin.hive.HiveBasicStatistics; +import io.trino.plugin.hive.PartitionStatistics; + +import java.util.HashMap; +import java.util.Map; +import java.util.Optional; +import java.util.OptionalDouble; +import java.util.OptionalLong; + +import static com.google.common.collect.ImmutableMap.toImmutableMap; +import static com.google.common.collect.Sets.intersection; + +public enum StatisticsUpdateMode +{ + /** + * Remove all existing table and column statistics, and replace them with the new statistics. + */ + OVERWRITE_ALL { + @Override + public PartitionStatistics updatePartitionStatistics(PartitionStatistics oldPartitionStats, PartitionStatistics newPartitionStats) + { + return newPartitionStats; + } + }, + /** + * Replace table statistics and present columns statics, but retain the statistics for columns that are not present in the new statistics. + */ + OVERWRITE_SOME_COLUMNS { + @Override + public PartitionStatistics updatePartitionStatistics(PartitionStatistics oldPartitionStats, PartitionStatistics newPartitionStats) + { + Map allColumnStatists = new HashMap<>(oldPartitionStats.getColumnStatistics()); + allColumnStatists.putAll(newPartitionStats.getColumnStatistics()); + + return new PartitionStatistics(newPartitionStats.getBasicStatistics(), allColumnStatists); + } + }, + /** + * Merges the new incremental data from an INSERT operation into the existing statistics. + */ + MERGE_INCREMENTAL { + @Override + public PartitionStatistics updatePartitionStatistics(PartitionStatistics oldPartitionStats, PartitionStatistics newPartitionStats) + { + return addIncrementalStatistics(oldPartitionStats, newPartitionStats); + } + }, + /** + * Undo the effect of a previous MERGE_INCREMENTAL operation. This will clear all column statistics because min/max calculations are not reversible. + */ + UNDO_MERGE_INCREMENTAL { + @Override + public PartitionStatistics updatePartitionStatistics(PartitionStatistics oldPartitionStats, PartitionStatistics newPartitionStats) + { + HiveBasicStatistics newTableStatistics = reduce(oldPartitionStats.getBasicStatistics(), newPartitionStats.getBasicStatistics(), Operator.SUBTRACT); + return new PartitionStatistics(newTableStatistics, ImmutableMap.of()); + } + }, + /** + * Removes all statics from the table and columns. + */ + CLEAR_ALL { + @Override + public PartitionStatistics updatePartitionStatistics(PartitionStatistics oldPartitionStats, PartitionStatistics newPartitionStats) + { + return PartitionStatistics.empty(); + } + }; + + public abstract PartitionStatistics updatePartitionStatistics(PartitionStatistics oldPartitionStats, PartitionStatistics newPartitionStats); + + private static PartitionStatistics addIncrementalStatistics(PartitionStatistics existingStatistics, PartitionStatistics incrementalStatistics) + { + if (existingStatistics.getBasicStatistics().getRowCount().isPresent() && existingStatistics.getBasicStatistics().getRowCount().getAsLong() == 0) { + return incrementalStatistics; + } + + if (incrementalStatistics.getBasicStatistics().getRowCount().isPresent() && incrementalStatistics.getBasicStatistics().getRowCount().getAsLong() == 0) { + return existingStatistics; + } + + var mergedTableStatistics = reduce(existingStatistics.getBasicStatistics(), incrementalStatistics.getBasicStatistics(), Operator.ADD); + + // only keep columns that have statistics in old and new + var mergedColumnStatistics = intersection(existingStatistics.getColumnStatistics().keySet(), incrementalStatistics.getColumnStatistics().keySet()).stream() + .collect(toImmutableMap( + column -> column, + column -> merge(column, existingStatistics, incrementalStatistics))); + + return new PartitionStatistics(mergedTableStatistics, mergedColumnStatistics); + } + + private static HiveBasicStatistics reduce(HiveBasicStatistics first, HiveBasicStatistics second, Operator operator) + { + return new HiveBasicStatistics( + reduce(first.getFileCount(), second.getFileCount(), operator, false), + reduce(first.getRowCount(), second.getRowCount(), operator, false), + reduce(first.getInMemoryDataSizeInBytes(), second.getInMemoryDataSizeInBytes(), operator, false), + reduce(first.getOnDiskDataSizeInBytes(), second.getOnDiskDataSizeInBytes(), operator, false)); + } + + private static HiveColumnStatistics merge(String column, PartitionStatistics firstStats, PartitionStatistics secondStats) + { + HiveColumnStatistics first = firstStats.getColumnStatistics().get(column); + HiveColumnStatistics second = secondStats.getColumnStatistics().get(column); + + return new HiveColumnStatistics( + mergeIntegerStatistics(first.getIntegerStatistics(), second.getIntegerStatistics()), + mergeDoubleStatistics(first.getDoubleStatistics(), second.getDoubleStatistics()), + mergeDecimalStatistics(first.getDecimalStatistics(), second.getDecimalStatistics()), + mergeDateStatistics(first.getDateStatistics(), second.getDateStatistics()), + mergeBooleanStatistics(first.getBooleanStatistics(), second.getBooleanStatistics()), + reduce(first.getMaxValueSizeInBytes(), second.getMaxValueSizeInBytes(), Operator.MAX, true), + mergeAverageColumnLength(column, firstStats, secondStats), + reduce(first.getNullsCount(), second.getNullsCount(), Operator.ADD, false), + mergeDistinctValueCount(column, firstStats, secondStats)); + } + + private static OptionalLong mergeDistinctValueCount(String column, PartitionStatistics first, PartitionStatistics second) + { + HiveColumnStatistics firstColumn = first.getColumnStatistics().get(column); + HiveColumnStatistics secondColumn = second.getColumnStatistics().get(column); + + OptionalLong firstDistinct = firstColumn.getDistinctValuesWithNullCount(); + OptionalLong secondDistinct = secondColumn.getDistinctValuesWithNullCount(); + + // if one column is entirely non-null and the other is entirely null + if (firstDistinct.isPresent() && noNulls(firstColumn) && isAllNull(second, secondColumn)) { + return OptionalLong.of(firstDistinct.getAsLong() + 1); + } + if (secondDistinct.isPresent() && noNulls(secondColumn) && isAllNull(first, firstColumn)) { + return OptionalLong.of(secondDistinct.getAsLong() + 1); + } + + if (firstDistinct.isPresent() && secondDistinct.isPresent()) { + return OptionalLong.of(max(firstDistinct.getAsLong(), secondDistinct.getAsLong())); + } + return OptionalLong.empty(); + } + + private static boolean noNulls(HiveColumnStatistics columnStats) + { + if (columnStats.getNullsCount().isEmpty()) { + return false; + } + return columnStats.getNullsCount().orElse(-1) == 0; + } + + private static boolean isAllNull(PartitionStatistics stats, HiveColumnStatistics columnStats) + { + if (stats.getBasicStatistics().getRowCount().isEmpty() || columnStats.getNullsCount().isEmpty()) { + return false; + } + return stats.getBasicStatistics().getRowCount().getAsLong() == columnStats.getNullsCount().getAsLong(); + } + + private static OptionalDouble mergeAverageColumnLength(String column, PartitionStatistics first, PartitionStatistics second) + { + // row count is required to merge average column length + if (first.getBasicStatistics().getRowCount().isEmpty() || second.getBasicStatistics().getRowCount().isEmpty()) { + return OptionalDouble.empty(); + } + long firstRowCount = first.getBasicStatistics().getRowCount().getAsLong(); + long secondRowCount = second.getBasicStatistics().getRowCount().getAsLong(); + + HiveColumnStatistics firstColumn = first.getColumnStatistics().get(column); + HiveColumnStatistics secondColumn = second.getColumnStatistics().get(column); + + // if one column is entirely null, return the average column length of the other column + if (firstRowCount == firstColumn.getNullsCount().orElse(0)) { + return secondColumn.getAverageColumnLength(); + } + if (secondRowCount == secondColumn.getNullsCount().orElse(0)) { + return firstColumn.getAverageColumnLength(); + } + + if (firstColumn.getAverageColumnLength().isEmpty() || secondColumn.getAverageColumnLength().isEmpty()) { + return OptionalDouble.empty(); + } + + long firstNonNullRowCount = firstRowCount - firstColumn.getNullsCount().orElse(0); + long secondNonNullRowCount = secondRowCount - secondColumn.getNullsCount().orElse(0); + + double firstTotalSize = firstColumn.getAverageColumnLength().getAsDouble() * firstNonNullRowCount; + double secondTotalSize = secondColumn.getAverageColumnLength().getAsDouble() * secondNonNullRowCount; + + return OptionalDouble.of((firstTotalSize + secondTotalSize) / (firstNonNullRowCount + secondNonNullRowCount)); + } + + private static Optional mergeIntegerStatistics(Optional first, Optional second) + { + // normally, either both or none is present + if (first.isPresent() && second.isPresent()) { + return Optional.of(new IntegerStatistics( + reduce(first.get().getMin(), second.get().getMin(), Operator.MIN, true), + reduce(first.get().getMax(), second.get().getMax(), Operator.MAX, true))); + } + return Optional.empty(); + } + + private static Optional mergeDoubleStatistics(Optional first, Optional second) + { + // normally, either both or none is present + if (first.isPresent() && second.isPresent()) { + return Optional.of(new DoubleStatistics( + reduce(first.get().getMin(), second.get().getMin(), Operator.MIN, true), + reduce(first.get().getMax(), second.get().getMax(), Operator.MAX, true))); + } + return Optional.empty(); + } + + private static Optional mergeDecimalStatistics(Optional first, Optional second) + { + // normally, either both or none is present + if (first.isPresent() && second.isPresent()) { + return Optional.of(new DecimalStatistics( + mergeComparable(first.get().getMin(), second.get().getMin(), Operator.MIN), + mergeComparable(first.get().getMax(), second.get().getMax(), Operator.MAX))); + } + return Optional.empty(); + } + + private static Optional mergeDateStatistics(Optional first, Optional second) + { + // normally, either both or none is present + if (first.isPresent() && second.isPresent()) { + return Optional.of(new DateStatistics( + mergeComparable(first.get().getMin(), second.get().getMin(), Operator.MIN), + mergeComparable(first.get().getMax(), second.get().getMax(), Operator.MAX))); + } + return Optional.empty(); + } + + private static Optional mergeBooleanStatistics(Optional first, Optional second) + { + // normally, either both or none is present + if (first.isPresent() && second.isPresent()) { + return Optional.of(new BooleanStatistics( + reduce(first.get().getTrueCount(), second.get().getTrueCount(), Operator.ADD, false), + reduce(first.get().getFalseCount(), second.get().getFalseCount(), Operator.ADD, false))); + } + return Optional.empty(); + } + + private static OptionalLong reduce(OptionalLong first, OptionalLong second, Operator operator, boolean returnFirstNonEmpty) + { + if (first.isPresent() && second.isPresent()) { + return switch (operator) { + case ADD -> OptionalLong.of(first.getAsLong() + second.getAsLong()); + case SUBTRACT -> OptionalLong.of(first.getAsLong() - second.getAsLong()); + case MAX -> OptionalLong.of(max(first.getAsLong(), second.getAsLong())); + case MIN -> OptionalLong.of(min(first.getAsLong(), second.getAsLong())); + }; + } + if (returnFirstNonEmpty) { + return first.isPresent() ? first : second; + } + return OptionalLong.empty(); + } + + private static OptionalDouble reduce(OptionalDouble first, OptionalDouble second, Operator operator, boolean returnFirstNonEmpty) + { + if (first.isPresent() && second.isPresent()) { + return switch (operator) { + case ADD -> OptionalDouble.of(first.getAsDouble() + second.getAsDouble()); + case SUBTRACT -> OptionalDouble.of(first.getAsDouble() - second.getAsDouble()); + case MAX -> OptionalDouble.of(max(first.getAsDouble(), second.getAsDouble())); + case MIN -> OptionalDouble.of(min(first.getAsDouble(), second.getAsDouble())); + }; + } + if (returnFirstNonEmpty) { + return first.isPresent() ? first : second; + } + return OptionalDouble.empty(); + } + + private static > Optional mergeComparable(Optional first, Optional second, Operator operator) + { + if (first.isPresent() && second.isPresent()) { + return switch (operator) { + case MAX -> Optional.of(max(first.get(), second.get())); + case MIN -> Optional.of(min(first.get(), second.get())); + default -> throw new IllegalArgumentException("Unexpected operator: " + operator); + }; + } + return first.isPresent() ? first : second; + } + + private static > T max(T first, T second) + { + return first.compareTo(second) >= 0 ? first : second; + } + + private static > T min(T first, T second) + { + return first.compareTo(second) <= 0 ? first : second; + } + + private enum Operator + { + ADD, + SUBTRACT, + MIN, + MAX, + } +} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/Storage.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/Storage.java index 824bc13ec766..c7fb7dee00ec 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/Storage.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/Storage.java @@ -174,6 +174,12 @@ public Builder setLocation(String location) return this; } + public Builder setBucketProperty(HiveBucketProperty bucketProperty) + { + this.bucketProperty = Optional.of(bucketProperty); + return this; + } + public Builder setBucketProperty(Optional bucketProperty) { this.bucketProperty = bucketProperty; diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/Table.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/Table.java index ba48867d8f7c..9ecccd042466 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/Table.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/Table.java @@ -18,8 +18,10 @@ import com.fasterxml.jackson.annotation.JsonProperty; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; -import com.google.common.collect.Maps; import com.google.errorprone.annotations.Immutable; +import io.trino.plugin.hive.HiveBucketProperty; +import io.trino.spi.TrinoException; +import io.trino.spi.connector.ColumnNotFoundException; import io.trino.spi.connector.SchemaTableName; import java.util.ArrayList; @@ -29,18 +31,22 @@ import java.util.Objects; import java.util.Optional; import java.util.OptionalLong; -import java.util.Set; import java.util.function.Consumer; import java.util.function.Function; import java.util.stream.Stream; import static com.google.common.base.MoreObjects.toStringHelper; -import static com.google.common.collect.ImmutableList.toImmutableList; +import static io.trino.plugin.hive.HiveErrorCode.HIVE_INVALID_METADATA; +import static io.trino.spi.StandardErrorCode.ALREADY_EXISTS; +import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED; +import static io.trino.spi.security.PrincipalType.USER; import static java.util.Objects.requireNonNull; @Immutable public class Table { + public static final String TABLE_COMMENT = "comment"; + private final String databaseName; private final String tableName; private final Optional owner; @@ -159,13 +165,171 @@ public OptionalLong getWriteId() return writeId; } - public Table withSelectedDataColumnsOnly(Set columns) + public Table withComment(Optional comment) + { + Map parameters = new LinkedHashMap<>(this.parameters); + comment.ifPresentOrElse( + value -> parameters.put(TABLE_COMMENT, value), + () -> parameters.remove(TABLE_COMMENT)); + + return builder(this) + .setParameters(parameters) + .build(); + } + + public Table withOwner(HivePrincipal principal) + { + // TODO Add role support https://github.com/trinodb/trino/issues/5706 + if (principal.getType() != USER) { + throw new TrinoException(NOT_SUPPORTED, "Setting table owner type as a role is not supported"); + } + + return builder(this) + .setOwner(Optional.of(principal.getName())) + .build(); + } + + public Table withColumnComment(String columnName, Optional columnComment) + { + boolean found = false; + ImmutableList.Builder newDataColumns = ImmutableList.builder(); + for (Column dataColumn : dataColumns) { + if (dataColumn.getName().equals(columnName)) { + if (found) { + throw new TrinoException(HIVE_INVALID_METADATA, "Table %s.%s has multiple columns named %s".formatted(databaseName, tableName, columnName)); + } + dataColumn = dataColumn.withComment(columnComment); + found = true; + } + newDataColumns.add(dataColumn); + } + + ImmutableList.Builder newPartitionColumns = ImmutableList.builder(); + for (Column partitionColumn : partitionColumns) { + if (partitionColumn.getName().equals(columnName)) { + if (found) { + throw new TrinoException(HIVE_INVALID_METADATA, "Table %s.%s has multiple columns named %s".formatted(databaseName, tableName, columnName)); + } + partitionColumn = partitionColumn.withComment(columnComment); + found = true; + } + newPartitionColumns.add(partitionColumn); + } + + if (!found) { + throw new ColumnNotFoundException(getSchemaTableName(), columnName); + } + + return builder(this) + .setDataColumns(newDataColumns.build()) + .setPartitionColumns(newPartitionColumns.build()) + .build(); + } + + public Table withAddColumn(Column newColumn) + { + if (dataColumns.stream().map(Column::getName).anyMatch(columnName -> columnName.equals(newColumn.getName())) || + partitionColumns.stream().map(Column::getName).anyMatch(columnName -> columnName.equals(newColumn.getName()))) { + throw new TrinoException(ALREADY_EXISTS, "Table %s.%s already has a column named %s".formatted(databaseName, tableName, newColumn.getName())); + } + return builder(this) + .addDataColumn(newColumn) + .build(); + } + + public Table withDropColumn(String columnName) + { + if (partitionColumns.stream().map(Column::getName).anyMatch(columnName::equals)) { + throw new TrinoException(NOT_SUPPORTED, "Cannot drop partition columns"); + } + if (storage.getBucketProperty().stream().flatMap(bucket -> bucket.getBucketedBy().stream()).anyMatch(columnName::equals) || + storage.getBucketProperty().stream().flatMap(bucket -> bucket.getSortedBy().stream()).map(SortingColumn::getColumnName).anyMatch(columnName::equals)) { + throw new TrinoException(NOT_SUPPORTED, "Cannot drop bucket column %s.%s.%s".formatted(databaseName, tableName, columnName)); + } + + boolean found = false; + ImmutableList.Builder newDataColumns = ImmutableList.builder(); + for (Column dataColumn : dataColumns) { + if (dataColumn.getName().equals(columnName)) { + found = true; + } + else { + newDataColumns.add(dataColumn); + } + } + if (!found) { + throw new ColumnNotFoundException(getSchemaTableName(), columnName); + } + List dataColumns = newDataColumns.build(); + if (dataColumns.isEmpty()) { + throw new TrinoException(NOT_SUPPORTED, "Cannot drop the only non-partition column in a table"); + } + return builder(this) + .setDataColumns(dataColumns) + .build(); + } + + public Table withRenameColumn(String oldColumnName, String newColumnName) + { + if (partitionColumns.stream().map(Column::getName).anyMatch(oldColumnName::equals)) { + throw new TrinoException(NOT_SUPPORTED, "Renaming partition columns is not supported"); + } + + if (dataColumns.stream().map(Column::getName).anyMatch(newColumnName::equals) || + partitionColumns.stream().map(Column::getName).anyMatch(newColumnName::equals)) { + throw new TrinoException(ALREADY_EXISTS, "Table %s.%s already has a column named %s".formatted(databaseName, tableName, newColumnName)); + } + + boolean found = false; + ImmutableList.Builder newDataColumns = ImmutableList.builder(); + for (Column dataColumn : dataColumns) { + if (dataColumn.getName().equals(oldColumnName)) { + newDataColumns.add(dataColumn.withName(newColumnName)); + found = true; + } + else { + newDataColumns.add(dataColumn); + } + } + if (!found) { + throw new ColumnNotFoundException(getSchemaTableName(), oldColumnName); + } + + Builder builder = builder(this) + .setDataColumns(newDataColumns.build()); + + if (storage.getBucketProperty().isPresent()) { + HiveBucketProperty bucketProperty = storage.getBucketProperty().get(); + + ImmutableList.Builder newBucketColumns = ImmutableList.builder(); + for (String bucketColumn : bucketProperty.getBucketedBy()) { + if (bucketColumn.equals(oldColumnName)) { + newBucketColumns.add(newColumnName); + } + else { + newBucketColumns.add(bucketColumn); + } + } + + ImmutableList.Builder newSortedColumns = ImmutableList.builder(); + for (SortingColumn sortingColumn : bucketProperty.getSortedBy()) { + if (sortingColumn.getColumnName().equals(oldColumnName)) { + newSortedColumns.add(new SortingColumn(newColumnName, sortingColumn.getOrder())); + } + else { + newSortedColumns.add(sortingColumn); + } + } + + builder.getStorageBuilder().setBucketProperty(new HiveBucketProperty(newBucketColumns.build(), bucketProperty.getBucketCount(), newSortedColumns.build())); + } + return builder.build(); + } + + public Table withParameters(Map newParameters) { - Map columnNameToColumn = Maps.uniqueIndex(getDataColumns(), Column::getName); - return Table.builder(this) - .setDataColumns(columns.stream() - .map(column -> requireNonNull(columnNameToColumn.get(column), "column " + column + " not found in table: " + this)) - .collect(toImmutableList())) + return builder(this) + .setParameters(newParameters) .build(); } diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/TableInfo.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/TableInfo.java new file mode 100644 index 000000000000..387e1f225ded --- /dev/null +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/TableInfo.java @@ -0,0 +1,69 @@ +/* + * 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.hive.metastore; + +import io.trino.spi.connector.RelationType; +import io.trino.spi.connector.SchemaTableName; + +import static java.util.Objects.requireNonNull; + +public record TableInfo(SchemaTableName tableName, ExtendedRelationType extendedRelationType) +{ + public static final String PRESTO_VIEW_COMMENT = "Presto View"; + public static final String ICEBERG_MATERIALIZED_VIEW_COMMENT = "Presto Materialized View"; + + public TableInfo + { + requireNonNull(tableName, "tableName is null"); + requireNonNull(extendedRelationType, "extendedRelationType is null"); + } + + public enum ExtendedRelationType + { + TABLE(RelationType.TABLE), + OTHER_VIEW(RelationType.VIEW), + OTHER_MATERIALIZED_VIEW(RelationType.MATERIALIZED_VIEW), + TRINO_VIEW(RelationType.VIEW), + TRINO_MATERIALIZED_VIEW(RelationType.MATERIALIZED_VIEW); + + private final RelationType relationType; + + ExtendedRelationType(RelationType relationType) + { + this.relationType = relationType; + } + + public RelationType toRelationType() + { + return relationType; + } + + public static ExtendedRelationType fromTableTypeAndComment(String tableType, String comment) + { + return switch (tableType) { + case "VIRTUAL_VIEW" -> { + if (PRESTO_VIEW_COMMENT.equals(comment)) { + yield TRINO_VIEW; + } + if (ICEBERG_MATERIALIZED_VIEW_COMMENT.equals(comment)) { + yield TRINO_MATERIALIZED_VIEW; + } + yield OTHER_VIEW; + } + case "MATERIALIZED_VIEW" -> ICEBERG_MATERIALIZED_VIEW_COMMENT.equals(comment) ? TRINO_MATERIALIZED_VIEW : OTHER_MATERIALIZED_VIEW; + default -> TABLE; + }; + } + } +} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/cache/CachingHiveMetastore.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/cache/CachingHiveMetastore.java index 29bbe003bbdf..eb02197195d0 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/cache/CachingHiveMetastore.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/cache/CachingHiveMetastore.java @@ -18,6 +18,7 @@ import com.google.common.cache.CacheLoader.InvalidCacheLoadException; import com.google.common.cache.LoadingCache; import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableSet; import com.google.common.collect.Iterables; import com.google.common.collect.Sets.SetView; import com.google.common.util.concurrent.ListenableFuture; @@ -28,15 +29,11 @@ import io.airlift.jmx.CacheStatsMBean; import io.airlift.units.Duration; import io.trino.cache.EvictableCacheBuilder; -import io.trino.hive.thrift.metastore.DataOperationType; -import io.trino.plugin.hive.HiveColumnStatisticType; import io.trino.plugin.hive.HivePartition; import io.trino.plugin.hive.HiveType; import io.trino.plugin.hive.PartitionStatistics; import io.trino.plugin.hive.acid.AcidOperation; -import io.trino.plugin.hive.acid.AcidTransaction; import io.trino.plugin.hive.metastore.AcidTransactionOwner; -import io.trino.plugin.hive.metastore.Column; import io.trino.plugin.hive.metastore.Database; import io.trino.plugin.hive.metastore.HiveColumnStatistics; import io.trino.plugin.hive.metastore.HiveMetastore; @@ -49,18 +46,17 @@ import io.trino.plugin.hive.metastore.PartitionFilter; import io.trino.plugin.hive.metastore.PartitionWithStatistics; import io.trino.plugin.hive.metastore.PrincipalPrivileges; +import io.trino.plugin.hive.metastore.StatisticsUpdateMode; import io.trino.plugin.hive.metastore.Table; -import io.trino.plugin.hive.metastore.TablesWithParameterCacheKey; -import io.trino.plugin.hive.metastore.UserTableKey; +import io.trino.plugin.hive.metastore.TableInfo; import io.trino.spi.TrinoException; import io.trino.spi.connector.SchemaTableName; -import io.trino.spi.connector.TableNotFoundException; import io.trino.spi.predicate.TupleDomain; import io.trino.spi.security.RoleGrant; -import io.trino.spi.type.Type; import org.weakref.jmx.Managed; import org.weakref.jmx.Nested; +import java.util.AbstractMap; import java.util.ArrayList; import java.util.Collection; import java.util.HashMap; @@ -88,23 +84,26 @@ import static com.google.common.collect.ImmutableSet.toImmutableSet; import static com.google.common.collect.Sets.difference; import static com.google.common.util.concurrent.Futures.immediateFuture; +import static io.airlift.concurrent.Threads.daemonThreadsNamed; import static io.trino.cache.CacheUtils.invalidateAllIf; import static io.trino.cache.CacheUtils.uncheckedCacheGet; import static io.trino.plugin.hive.metastore.HivePartitionName.hivePartitionName; import static io.trino.plugin.hive.metastore.HiveTableName.hiveTableName; -import static io.trino.plugin.hive.metastore.MetastoreUtil.makePartitionName; import static io.trino.plugin.hive.metastore.PartitionFilter.partitionFilter; +import static io.trino.plugin.hive.metastore.cache.CachingHiveMetastore.ObjectType.OTHER; +import static io.trino.plugin.hive.metastore.cache.CachingHiveMetastore.ObjectType.PARTITION; +import static io.trino.plugin.hive.metastore.cache.CachingHiveMetastore.ObjectType.STATS; import static io.trino.plugin.hive.util.HiveUtil.makePartName; import static java.util.Collections.unmodifiableSet; import static java.util.Objects.requireNonNull; +import static java.util.concurrent.Executors.newCachedThreadPool; import static java.util.concurrent.TimeUnit.MILLISECONDS; -import static java.util.function.UnaryOperator.identity; /** * Hive Metastore Cache */ @ThreadSafe -public class CachingHiveMetastore +public final class CachingHiveMetastore implements HiveMetastore { public enum StatsRecording @@ -113,24 +112,27 @@ public enum StatsRecording DISABLED } - protected final HiveMetastore delegate; - private final boolean cacheMissing; + public enum ObjectType + { + PARTITION, + STATS, + OTHER, + } + + private final HiveMetastore delegate; + private final Set cacheMissing; private final LoadingCache> databaseCache; private final LoadingCache> databaseNamesCache; private final LoadingCache> tableCache; - private final LoadingCache> tableNamesCache; - private final LoadingCache>> allTableNamesCache; - private final LoadingCache> tablesWithParameterCache; - private final Cache> tableStatisticsCache; - private final Cache> partitionStatisticsCache; - private final LoadingCache> viewNamesCache; - private final LoadingCache>> allViewNamesCache; + private final LoadingCache> tablesCacheNew; + private final Cache>> tableColumnStatisticsCache; + private final LoadingCache> tableNamesWithParametersCache; + private final Cache>> partitionStatisticsCache; private final Cache>> partitionCache; private final LoadingCache>> partitionFilterCache; private final LoadingCache> tablePrivilegesCache; private final LoadingCache> rolesCache; private final LoadingCache> roleGrantsCache; - private final LoadingCache> grantedPrincipalsCache; private final LoadingCache> configValuesCache; public static CachingHiveMetastoreBuilder builder() @@ -304,74 +306,109 @@ public CachingHiveMetastore build() requireNonNull(maximumSize, "maximumSize not set"); requireNonNull(cacheMissing, "cacheMissing not set"); requireNonNull(partitionCacheEnabled, "partitionCacheEnabled not set"); - return new CachingHiveMetastore( + Executor refreshExecutor = newCachedThreadPool(daemonThreadsNamed("hive-metastore--%s")); + ImmutableSet.Builder builder = ImmutableSet.builder(); + if (cacheMissing) { + builder.add(ObjectType.OTHER); + } + if (partitionCacheEnabled) { + builder.add(ObjectType.PARTITION); + } + if (statsCacheEnabled) { + builder.add(ObjectType.STATS); + } + return createCachingHiveMetastore( delegate, - metadataCacheEnabled, - statsCacheEnabled, - expiresAfterWriteMillis, - statsExpiresAfterWriteMillis, - refreshMills, - executor, + expiresAfterWriteMillis.isPresent() ? new Duration(expiresAfterWriteMillis.getAsLong(), MILLISECONDS) : Duration.ZERO, + statsExpiresAfterWriteMillis.isPresent() ? new Duration(statsExpiresAfterWriteMillis.getAsLong(), MILLISECONDS) : Duration.ZERO, + refreshMills.isPresent() ? Optional.of(new Duration(refreshMills.getAsLong(), MILLISECONDS)) : Optional.empty(), + refreshExecutor, maximumSize, - statsRecording, - cacheMissing, - partitionCacheEnabled); + StatsRecording.ENABLED, + partitionCacheEnabled, + builder.build()); } } - protected CachingHiveMetastore( + public static CachingHiveMetastore createPerTransactionCache(HiveMetastore delegate, long maximumSize) + { + return new CachingHiveMetastore( + delegate, + ImmutableSet.copyOf(ObjectType.values()), + new CacheFactory(maximumSize), + new CacheFactory(maximumSize), + new CacheFactory(maximumSize), + new CacheFactory(maximumSize)); + } + + public static CachingHiveMetastore createCachingHiveMetastore( HiveMetastore delegate, - boolean metadataCacheEnabled, - boolean statsCacheEnabled, - OptionalLong expiresAfterWriteMillis, - OptionalLong statsExpiresAfterWriteMillis, - OptionalLong refreshMills, - Optional executor, + Duration metadataCacheTtl, + Duration statsCacheTtl, + Optional refreshInterval, + Executor refreshExecutor, long maximumSize, StatsRecording statsRecording, - boolean cacheMissing, - boolean partitionCacheEnabled) + boolean partitionCacheEnabled, + Set cacheMissing) { - checkArgument(metadataCacheEnabled || statsCacheEnabled, "Cache not enabled"); - this.delegate = requireNonNull(delegate, "delegate is null"); - this.cacheMissing = cacheMissing; - requireNonNull(executor, "executor is null"); + // refresh executor is only required when the refresh interval is set, but the executor is + // always set, so it is simpler to just enforce that + requireNonNull(refreshExecutor, "refreshExecutor is null"); - CacheFactory cacheFactory; - CacheFactory partitionCacheFactory; - if (metadataCacheEnabled) { - cacheFactory = cacheFactory(expiresAfterWriteMillis, refreshMills, executor, maximumSize, statsRecording); - partitionCacheFactory = partitionCacheEnabled ? cacheFactory : neverCacheFactory(); - } - else { - cacheFactory = neverCacheFactory(); - partitionCacheFactory = neverCacheFactory(); - } + long metadataCacheMillis = metadataCacheTtl.toMillis(); + long statsCacheMillis = statsCacheTtl.toMillis(); + checkArgument(metadataCacheMillis > 0 || statsCacheMillis > 0, "Cache not enabled"); - CacheFactory statsCacheFactory; - CacheFactory partitionStatsCacheFactory; - if (statsCacheEnabled) { - statsCacheFactory = cacheFactory(statsExpiresAfterWriteMillis, refreshMills, executor, maximumSize, statsRecording); - partitionStatsCacheFactory = partitionCacheEnabled ? statsCacheFactory : neverCacheFactory(); + OptionalLong refreshMillis = refreshInterval.stream().mapToLong(Duration::toMillis).findAny(); + + CacheFactory cacheFactory = CacheFactory.NEVER_CACHE; + CacheFactory partitionCacheFactory = CacheFactory.NEVER_CACHE; + if (metadataCacheMillis > 0) { + cacheFactory = new CacheFactory(OptionalLong.of(metadataCacheMillis), refreshMillis, Optional.of(refreshExecutor), maximumSize, statsRecording); + if (partitionCacheEnabled) { + partitionCacheFactory = cacheFactory; + } } - else { - statsCacheFactory = neverCacheFactory(); - partitionStatsCacheFactory = neverCacheFactory(); + + CacheFactory statsCacheFactory = CacheFactory.NEVER_CACHE; + CacheFactory partitionStatsCacheFactory = CacheFactory.NEVER_CACHE; + if (statsCacheMillis > 0) { + statsCacheFactory = new CacheFactory(OptionalLong.of(statsCacheMillis), refreshMillis, Optional.of(refreshExecutor), maximumSize, statsRecording); + if (partitionCacheEnabled) { + partitionStatsCacheFactory = statsCacheFactory; + } } - databaseNamesCache = cacheFactory.buildCache(ignored -> loadAllDatabases()); + return new CachingHiveMetastore( + delegate, + cacheMissing, + cacheFactory, + partitionCacheFactory, + statsCacheFactory, + partitionStatsCacheFactory); + } + + private CachingHiveMetastore( + HiveMetastore delegate, + Set cacheMissing, + CacheFactory cacheFactory, + CacheFactory partitionCacheFactory, + CacheFactory statsCacheFactory, + CacheFactory partitionStatsCacheFactory) + { + this.delegate = requireNonNull(delegate, "delegate is null"); + this.cacheMissing = cacheMissing; + + databaseNamesCache = cacheFactory.buildCache(ignore -> loadAllDatabases()); databaseCache = cacheFactory.buildCache(this::loadDatabase); - tableNamesCache = cacheFactory.buildCache(this::loadAllTables); - allTableNamesCache = cacheFactory.buildCache(ignore -> loadAllTables()); - tablesWithParameterCache = cacheFactory.buildCache(this::loadTablesMatchingParameter); - tableStatisticsCache = statsCacheFactory.buildCache(this::refreshTableStatistics); + tablesCacheNew = cacheFactory.buildCache(this::loadTablesNew); + tableColumnStatisticsCache = statsCacheFactory.buildCache(this::refreshTableColumnStatistics); tableCache = cacheFactory.buildCache(this::loadTable); - viewNamesCache = cacheFactory.buildCache(this::loadAllViews); - allViewNamesCache = cacheFactory.buildCache(ignore -> loadAllViews()); - tablePrivilegesCache = cacheFactory.buildCache(key -> loadTablePrivileges(key.getDatabase(), key.getTable(), key.getOwner(), key.getPrincipal())); - rolesCache = cacheFactory.buildCache(ignored -> loadRoles()); + tableNamesWithParametersCache = cacheFactory.buildCache(this::loadTablesMatchingParameter); + tablePrivilegesCache = cacheFactory.buildCache(key -> loadTablePrivileges(key.database(), key.table(), key.owner(), key.principal())); + rolesCache = cacheFactory.buildCache(ignore -> loadRoles()); roleGrantsCache = cacheFactory.buildCache(this::loadRoleGrants); - grantedPrincipalsCache = cacheFactory.buildCache(this::loadPrincipals); configValuesCache = cacheFactory.buildCache(this::loadConfigValue); partitionStatisticsCache = partitionStatsCacheFactory.buildBulkCache(); @@ -383,16 +420,14 @@ protected CachingHiveMetastore( public void flushCache() { databaseNamesCache.invalidateAll(); - tableNamesCache.invalidateAll(); - allTableNamesCache.invalidateAll(); - viewNamesCache.invalidateAll(); - allViewNamesCache.invalidateAll(); + tablesCacheNew.invalidateAll(); databaseCache.invalidateAll(); tableCache.invalidateAll(); + tableNamesWithParametersCache.invalidateAll(); partitionCache.invalidateAll(); partitionFilterCache.invalidateAll(); tablePrivilegesCache.invalidateAll(); - tableStatisticsCache.invalidateAll(); + tableColumnStatisticsCache.invalidateAll(); partitionStatisticsCache.invalidateAll(); rolesCache.invalidateAll(); } @@ -408,20 +443,19 @@ public void flushPartitionCache(String schemaName, String tableName, List partitionNameToCheck.map(value -> value.equals(providedPartitionName)).orElse(false)); } - private AtomicReference refreshTableStatistics(HiveTableName tableName, AtomicReference currentValueHolder) + private AtomicReference> refreshTableColumnStatistics(HiveTableName tableName, AtomicReference> currentValueHolder) { - PartitionStatistics currentValue = currentValueHolder.get(); + Map currentValue = currentValueHolder.get(); if (currentValue == null) { // do not refresh empty value return currentValueHolder; } - PartitionStatistics reloaded = getTable(tableName.getDatabaseName(), tableName.getTableName()) - .map(table -> table.withSelectedDataColumnsOnly(currentValue.getColumnStatistics().keySet())) - .map(delegate::getTableStatistics) - .orElseThrow(); + + // only refresh currently loaded columns + Map columnStatistics = delegate.getTableColumnStatistics(tableName.getDatabaseName(), tableName.getTableName(), currentValue.keySet()); // return new value holder to have only fresh data in case of concurrent loads - return new AtomicReference<>(reloaded); + return new AtomicReference<>(columnStatistics); } private static V get(LoadingCache cache, K key) @@ -437,14 +471,14 @@ private static V get(LoadingCache cache, K key) } } - private Optional getOptional(LoadingCache> cache, K key) + private Optional getOptional(ObjectType objectType, LoadingCache> cache, K key) { try { Optional value = cache.getIfPresent(key); @SuppressWarnings("OptionalAssignedToNull") boolean valueIsPresent = value != null; if (valueIsPresent) { - if (value.isPresent() || cacheMissing) { + if (value.isPresent() || cacheMissing.contains(objectType)) { return value; } cache.invalidate(key); @@ -501,7 +535,7 @@ private static V getIncrementally( private static Map getAll(Cache> cache, Iterable keys, Function, Map> bulkLoader) { ImmutableMap.Builder result = ImmutableMap.builder(); - Map> toLoad = new HashMap<>(); + ImmutableMap.Builder> toLoadBuilder = ImmutableMap.builder(); for (K key : keys) { AtomicReference valueHolder = uncheckedCacheGet(cache, key, AtomicReference::new); @@ -510,10 +544,11 @@ private static Map getAll(Cache> cache, Itera result.put(key, value); } else { - toLoad.put(key, valueHolder); + toLoadBuilder.put(key, valueHolder); } } + Map> toLoad = toLoadBuilder.buildOrThrow(); if (toLoad.isEmpty()) { return result.buildOrThrow(); } @@ -543,7 +578,7 @@ private static Map getAll( keys.forEach(key -> { // make sure the value holder is retrieved before the new values are loaded - // so that in case of invalidation we will not set the stale value + // so that in case of invalidation, we will not set the stale value AtomicReference currentValueHolder = uncheckedCacheGet(cache, key, AtomicReference::new); V currentValue = currentValueHolder.get(); if (currentValue != null && isSufficient.test(currentValue)) { @@ -572,7 +607,7 @@ private static Map getAll( @Override public Optional getDatabase(String databaseName) { - return getOptional(databaseCache, databaseName); + return getOptional(OTHER, databaseCache, databaseName); } private Optional loadDatabase(String databaseName) @@ -591,22 +626,10 @@ private List loadAllDatabases() return delegate.getAllDatabases(); } - private Table getExistingTable(String databaseName, String tableName) - { - return getTable(databaseName, tableName) - .orElseThrow(() -> new TableNotFoundException(new SchemaTableName(databaseName, tableName))); - } - @Override public Optional
getTable(String databaseName, String tableName) { - return getOptional(tableCache, hiveTableName(databaseName, tableName)); - } - - @Override - public Set getSupportedColumnStatistics(Type type) - { - return delegate.getSupportedColumnStatistics(type); + return getOptional(OTHER, tableCache, hiveTableName(databaseName, tableName)); } private Optional
loadTable(HiveTableName hiveTableName) @@ -614,127 +637,116 @@ private Optional
loadTable(HiveTableName hiveTableName) return delegate.getTable(hiveTableName.getDatabaseName(), hiveTableName.getTableName()); } - /** - * The method will cache and return columns specified in the {@link Table#getDataColumns()} - * but may return more if other columns are already cached. - */ @Override - public PartitionStatistics getTableStatistics(Table table) - { - Set dataColumns = table.getDataColumns().stream().map(Column::getName).collect(toImmutableSet()); - - return getIncrementally( - tableStatisticsCache, - hiveTableName(table.getDatabaseName(), table.getTableName()), - currentStatistics -> currentStatistics.getColumnStatistics().keySet().containsAll(dataColumns), - () -> delegate.getTableStatistics(table), + public Map getTableColumnStatistics(String databaseName, String tableName, Set columnNames) + { + checkArgument(!columnNames.isEmpty(), "columnNames is empty"); + Map columnStatistics = getIncrementally( + tableColumnStatisticsCache, + hiveTableName(databaseName, tableName), + currentStatistics -> currentStatistics.keySet().containsAll(columnNames), + () -> delegate.getTableColumnStatistics(databaseName, tableName, columnNames), currentStatistics -> { - SetView missingColumns = difference(dataColumns, currentStatistics.getColumnStatistics().keySet()); - Table tableWithOnlyMissingColumns = table.withSelectedDataColumnsOnly(missingColumns); - return delegate.getTableStatistics(tableWithOnlyMissingColumns); + SetView missingColumns = difference(columnNames, currentStatistics.keySet()); + return delegate.getTableColumnStatistics(databaseName, tableName, missingColumns); }, - CachingHiveMetastore::mergePartitionColumnStatistics); + (currentStats, newStats) -> mergeColumnStatistics(currentStats, newStats, columnNames)); + // HiveColumnStatistics.empty() are removed to make output consistent with non-cached metastore which simplifies testing + return removeEmptyColumnStatistics(columnNames, columnStatistics); } - private PartitionStatistics loadTableColumnStatistics(HiveTableName tableName) + @Override + public Map> getPartitionColumnStatistics(String databaseName, String tableName, Set partitionNames, Set columnNames) { - Table table = getExistingTable(tableName.getDatabaseName(), tableName.getTableName()); - return delegate.getTableStatistics(table); + checkArgument(!columnNames.isEmpty(), "columnNames is empty"); + HiveTableName hiveTableName = hiveTableName(databaseName, tableName); + List hivePartitionNames = partitionNames.stream().map(partitionName -> hivePartitionName(hiveTableName, partitionName)).toList(); + Map> statistics = getAll( + partitionStatisticsCache, + hivePartitionNames, + missingPartitions -> loadPartitionsColumnStatistics(databaseName, tableName, columnNames, missingPartitions), + currentStats -> currentStats.keySet().containsAll(columnNames), + (currentStats, newStats) -> mergeColumnStatistics(currentStats, newStats, columnNames)); + // HiveColumnStatistics.empty() are removed to make output consistent with non-cached metastore which simplifies testing + return statistics.entrySet().stream() + .collect(toImmutableMap( + entry -> entry.getKey().getPartitionName().orElseThrow(), + entry -> removeEmptyColumnStatistics(columnNames, entry.getValue()))); } - /** - * The method will cache and return columns specified in the {@link Table#getDataColumns()} - * but may return more if other columns are already cached for a given partition. - */ @Override - public Map getPartitionStatistics(Table table, List partitions) + public boolean useSparkTableStatistics() { - HiveTableName hiveTableName = hiveTableName(table.getDatabaseName(), table.getTableName()); - Map partitionsByName = partitions.stream() - .collect(toImmutableMap(partition -> hivePartitionName(hiveTableName, makePartitionName(table, partition)), identity())); - - Set dataColumns = table.getDataColumns().stream().map(Column::getName).collect(toImmutableSet()); + return delegate.useSparkTableStatistics(); + } - Map statistics = getAll( - partitionStatisticsCache, - partitionsByName.keySet(), - missingPartitions -> loadPartitionsColumnStatistics(table, partitionsByName, missingPartitions), - currentStats -> currentStats.getColumnStatistics().keySet().containsAll(dataColumns), - CachingHiveMetastore::mergePartitionColumnStatistics); - return statistics.entrySet().stream() - .collect(toImmutableMap(entry -> entry.getKey().getPartitionName().orElseThrow(), Entry::getValue)); + private static ImmutableMap removeEmptyColumnStatistics(Set columnNames, Map columnStatistics) + { + return columnStatistics.entrySet().stream() + .filter(entry -> columnNames.contains(entry.getKey()) && !entry.getValue().equals(HiveColumnStatistics.empty())) + .collect(toImmutableMap(Entry::getKey, Entry::getValue)); } - private static PartitionStatistics mergePartitionColumnStatistics(PartitionStatistics currentStats, PartitionStatistics newStats) + private Map mergeColumnStatistics(Map currentStats, Map newStats, Set dataColumns) { requireNonNull(newStats, "newStats is null"); - if (currentStats == null) { - return newStats; + ImmutableMap.Builder columnStatisticsBuilder = ImmutableMap.builder(); + // Populate empty statistics for all requested columns to cache absence of column statistics for future requests. + if (cacheMissing.contains(STATS)) { + columnStatisticsBuilder.putAll(Iterables.transform( + dataColumns, + column -> new AbstractMap.SimpleEntry<>(column, HiveColumnStatistics.empty()))); + } + if (currentStats != null) { + columnStatisticsBuilder.putAll(currentStats); } - return new PartitionStatistics( - newStats.getBasicStatistics(), - ImmutableMap.builder() - .putAll(currentStats.getColumnStatistics()) - .putAll(newStats.getColumnStatistics()) - .buildKeepingLast()); + columnStatisticsBuilder.putAll(newStats); + return columnStatisticsBuilder.buildKeepingLast(); } - private Map loadPartitionsColumnStatistics(Table table, Map partitionsByName, Collection partitionNamesToLoad) + private Map> loadPartitionsColumnStatistics( + String databaseName, + String tableName, + Set columnNames, + Collection partitionNamesToLoad) { if (partitionNamesToLoad.isEmpty()) { return ImmutableMap.of(); } - ImmutableMap.Builder result = ImmutableMap.builder(); - List partitionsToLoad = partitionNamesToLoad.stream() - .map(partitionsByName::get) - .collect(toImmutableList()); - Map statisticsByPartitionName = delegate.getPartitionStatistics(table, partitionsToLoad); + Set partitionsToLoad = partitionNamesToLoad.stream() + .map(partitionName -> partitionName.getPartitionName().orElseThrow()) + .collect(toImmutableSet()); + Map> columnStatistics = delegate.getPartitionColumnStatistics(databaseName, tableName, partitionsToLoad, columnNames); + + ImmutableMap.Builder> result = ImmutableMap.builder(); for (HivePartitionName partitionName : partitionNamesToLoad) { - String stringNameForPartition = partitionName.getPartitionName().orElseThrow(); - result.put(partitionName, statisticsByPartitionName.get(stringNameForPartition)); + result.put(partitionName, columnStatistics.getOrDefault(partitionName.getPartitionName().orElseThrow(), ImmutableMap.of())); } return result.buildOrThrow(); } @Override - public void updateTableStatistics(String databaseName, - String tableName, - AcidTransaction transaction, - Function update) + public void updateTableStatistics(String databaseName, String tableName, OptionalLong acidWriteId, StatisticsUpdateMode mode, PartitionStatistics statisticsUpdate) { try { - delegate.updateTableStatistics(databaseName, tableName, transaction, update); + delegate.updateTableStatistics(databaseName, tableName, acidWriteId, mode, statisticsUpdate); } finally { HiveTableName hiveTableName = hiveTableName(databaseName, tableName); - tableStatisticsCache.invalidate(hiveTableName); + tableColumnStatisticsCache.invalidate(hiveTableName); // basic stats are stored as table properties tableCache.invalidate(hiveTableName); } } @Override - public void updatePartitionStatistics(Table table, String partitionName, Function update) + public void updatePartitionStatistics(Table table, StatisticsUpdateMode mode, Map partitionUpdates) { try { - delegate.updatePartitionStatistics(table, partitionName, update); + delegate.updatePartitionStatistics(table, mode, partitionUpdates); } finally { - HivePartitionName hivePartitionName = hivePartitionName(hiveTableName(table.getDatabaseName(), table.getTableName()), partitionName); - partitionStatisticsCache.invalidate(hivePartitionName); - // basic stats are stored as partition properties - partitionCache.invalidate(hivePartitionName); - } - } - - @Override - public void updatePartitionStatistics(Table table, Map> updates) - { - try { - delegate.updatePartitionStatistics(table, updates); - } - finally { - updates.forEach((partitionName, update) -> { + partitionUpdates.keySet().forEach(partitionName -> { HivePartitionName hivePartitionName = hivePartitionName(hiveTableName(table.getDatabaseName(), table.getTableName()), partitionName); partitionStatisticsCache.invalidate(hivePartitionName); // basic stats are stored as partition properties @@ -744,59 +756,26 @@ public void updatePartitionStatistics(Table table, Map getAllTables(String databaseName) + public List getTables(String databaseName) { - return get(tableNamesCache, databaseName); + return get(tablesCacheNew, databaseName); } - private List loadAllTables(String databaseName) + private List loadTablesNew(String databaseName) { - return delegate.getAllTables(databaseName); + return delegate.getTables(databaseName); } @Override - public Optional> getAllTables() - { - return getOptional(allTableNamesCache, SingletonCacheKey.INSTANCE); - } - - private Optional> loadAllTables() + public List getTableNamesWithParameters(String databaseName, String parameterKey, ImmutableSet parameterValues) { - return delegate.getAllTables(); - } - - @Override - public List getTablesWithParameter(String databaseName, String parameterKey, String parameterValue) - { - TablesWithParameterCacheKey key = new TablesWithParameterCacheKey(databaseName, parameterKey, parameterValue); - return get(tablesWithParameterCache, key); + TablesWithParameterCacheKey key = new TablesWithParameterCacheKey(databaseName, parameterKey, parameterValues); + return get(tableNamesWithParametersCache, key); } private List loadTablesMatchingParameter(TablesWithParameterCacheKey key) { - return delegate.getTablesWithParameter(key.getDatabaseName(), key.getParameterKey(), key.getParameterValue()); - } - - @Override - public List getAllViews(String databaseName) - { - return get(viewNamesCache, databaseName); - } - - private List loadAllViews(String databaseName) - { - return delegate.getAllViews(databaseName); - } - - @Override - public Optional> getAllViews() - { - return getOptional(allViewNamesCache, SingletonCacheKey.INSTANCE); - } - - private Optional> loadAllViews() - { - return delegate.getAllViews(); + return delegate.getTableNamesWithParameters(key.databaseName(), key.parameterKey(), key.parameterValues()); } @Override @@ -844,7 +823,7 @@ public void setDatabaseOwner(String databaseName, HivePrincipal principal) } } - protected void invalidateDatabase(String databaseName) + private void invalidateDatabase(String databaseName) { databaseCache.invalidate(databaseName); databaseNamesCache.invalidateAll(); @@ -873,10 +852,10 @@ public void dropTable(String databaseName, String tableName, boolean deleteData) } @Override - public void replaceTable(String databaseName, String tableName, Table newTable, PrincipalPrivileges principalPrivileges) + public void replaceTable(String databaseName, String tableName, Table newTable, PrincipalPrivileges principalPrivileges, Map environmentContext) { try { - delegate.replaceTable(databaseName, tableName, newTable, principalPrivileges); + delegate.replaceTable(databaseName, tableName, newTable, principalPrivileges, environmentContext); } finally { invalidateTable(databaseName, tableName); @@ -966,27 +945,13 @@ public void invalidateTable(String databaseName, String tableName) { HiveTableName hiveTableName = new HiveTableName(databaseName, tableName); tableCache.invalidate(hiveTableName); - tableNamesCache.invalidate(databaseName); - allTableNamesCache.invalidateAll(); - viewNamesCache.invalidate(databaseName); - allViewNamesCache.invalidateAll(); + tablesCacheNew.invalidate(databaseName); + tableNamesWithParametersCache.invalidateAll(); invalidateAllIf(tablePrivilegesCache, userTableKey -> userTableKey.matches(databaseName, tableName)); - tableStatisticsCache.invalidate(hiveTableName); - invalidateTablesWithParameterCache(databaseName, tableName); + tableColumnStatisticsCache.invalidate(hiveTableName); invalidatePartitionCache(databaseName, tableName); } - private void invalidateTablesWithParameterCache(String databaseName, String tableName) - { - tablesWithParameterCache.asMap().keySet().stream() - .filter(cacheKey -> cacheKey.getDatabaseName().equals(databaseName)) - .filter(cacheKey -> { - List cacheValue = tablesWithParameterCache.getIfPresent(cacheKey); - return cacheValue != null && cacheValue.contains(tableName); - }) - .forEach(tablesWithParameterCache::invalidate); - } - @Override public Optional getPartition(Table table, List partitionValues) { @@ -1000,7 +965,7 @@ public Optional> getPartitionNamesByFilter( List columnNames, TupleDomain partitionKeysFilter) { - return getOptional(partitionFilterCache, partitionFilter(databaseName, tableName, columnNames, partitionKeysFilter)); + return getOptional(PARTITION, partitionFilterCache, partitionFilter(databaseName, tableName, columnNames, partitionKeysFilter)); } private Optional> loadPartitionNamesByFilter(PartitionFilter partitionFilter) @@ -1142,12 +1107,6 @@ public void revokeRoles(Set roles, Set grantees, boolean } } - @Override - public Set listGrantedPrincipals(String role) - { - return get(grantedPrincipalsCache, role); - } - @Override public Set listRoleGrants(HivePrincipal principal) { @@ -1159,11 +1118,6 @@ private Set loadRoleGrants(HivePrincipal principal) return delegate.listRoleGrants(principal); } - private Set loadPrincipals(String role) - { - return delegate.listGrantedPrincipals(role); - } - private void invalidatePartitionCache(String databaseName, String tableName) { invalidatePartitionCache(databaseName, tableName, partitionName -> true); @@ -1219,7 +1173,7 @@ public Set listTablePrivileges(String databaseName, String ta @Override public Optional getConfigValue(String name) { - return getOptional(configValuesCache, name); + return getOptional(OTHER, configValuesCache, name); } private Optional loadConfigValue(String name) @@ -1292,7 +1246,7 @@ public void acquireTableWriteLock( long transactionId, String dbName, String tableName, - DataOperationType operation, + AcidOperation operation, boolean isDynamicPartitionWrite) { delegate.acquireTableWriteLock(transactionOwner, queryId, transactionId, dbName, tableName, operation, isDynamicPartitionWrite); @@ -1309,17 +1263,6 @@ public void updateTableWriteId(String dbName, String tableName, long transaction } } - @Override - public void alterPartitions(String dbName, String tableName, List partitions, long writeId) - { - try { - delegate.alterPartitions(dbName, tableName, partitions, writeId); - } - finally { - invalidatePartitionCache(dbName, tableName); - } - } - @Override public void addDynamicPartitions(String dbName, String tableName, List partitionNames, long transactionId, long writeId, AcidOperation operation) { @@ -1342,26 +1285,6 @@ public void alterTransactionalTable(Table table, long transactionId, long writeI } } - private static CacheFactory cacheFactory( - OptionalLong expiresAfterWriteMillis, - OptionalLong refreshMillis, - Optional refreshExecutor, - long maximumSize, - StatsRecording statsRecording) - { - return new CacheFactory(expiresAfterWriteMillis, refreshMillis, refreshExecutor, maximumSize, statsRecording); - } - - private static CacheFactory neverCacheFactory() - { - return cacheFactory( - OptionalLong.of(0), - OptionalLong.empty(), - Optional.empty(), - 0, - StatsRecording.DISABLED); - } - private static LoadingCache buildCache( OptionalLong expiresAfterWriteMillis, OptionalLong refreshMillis, @@ -1408,15 +1331,35 @@ private static Cache> buildBulkCache( return cacheBuilder.build(); } - private enum SingletonCacheKey + record TablesWithParameterCacheKey(String databaseName, String parameterKey, ImmutableSet parameterValues) + { + TablesWithParameterCacheKey + { + requireNonNull(databaseName, "databaseName is null"); + requireNonNull(parameterKey, "parameterKey is null"); + requireNonNull(parameterValues, "parameterValues is null"); + } + } + + record UserTableKey(Optional principal, String database, String table, Optional owner) { - INSTANCE + UserTableKey + { + requireNonNull(principal, "principal is null"); + requireNonNull(database, "database is null"); + requireNonNull(table, "table is null"); + requireNonNull(owner, "owner is null"); + } + + public boolean matches(String databaseName, String tableName) + { + return this.database.equals(databaseName) && this.table.equals(tableName); + } } // // Stats used for non-impersonation shared caching // - @Managed @Nested public CacheStatsMBean getDatabaseStats() @@ -1442,28 +1385,21 @@ public CacheStatsMBean getTableStats() @Nested public CacheStatsMBean getTableNamesStats() { - return new CacheStatsMBean(tableNamesCache); - } - - @Managed - @Nested - public CacheStatsMBean getAllTableNamesStats() - { - return new CacheStatsMBean(allTableNamesCache); + return new CacheStatsMBean(tablesCacheNew); } @Managed @Nested public CacheStatsMBean getTableWithParameterStats() { - return new CacheStatsMBean(tablesWithParameterCache); + return new CacheStatsMBean(tableNamesWithParametersCache); } @Managed @Nested - public CacheStatsMBean getTableStatisticsStats() + public CacheStatsMBean getTableColumnStatisticsStats() { - return new CacheStatsMBean(tableStatisticsCache); + return new CacheStatsMBean(tableColumnStatisticsCache); } @Managed @@ -1473,20 +1409,6 @@ public CacheStatsMBean getPartitionStatisticsStats() return new CacheStatsMBean(partitionStatisticsCache); } - @Managed - @Nested - public CacheStatsMBean getViewNamesStats() - { - return new CacheStatsMBean(viewNamesCache); - } - - @Managed - @Nested - public CacheStatsMBean getAllViewNamesStats() - { - return new CacheStatsMBean(allViewNamesCache); - } - @Managed @Nested public CacheStatsMBean getPartitionStats() @@ -1522,13 +1444,6 @@ public CacheStatsMBean getRoleGrantsStats() return new CacheStatsMBean(roleGrantsCache); } - @Managed - @Nested - public CacheStatsMBean getGrantedPrincipalsStats() - { - return new CacheStatsMBean(grantedPrincipalsCache); - } - @Managed @Nested public CacheStatsMBean getConfigValuesStats() @@ -1554,41 +1469,26 @@ LoadingCache> getTableCache() return tableCache; } - LoadingCache> getTableNamesCache() - { - return tableNamesCache; - } - - LoadingCache>> getAllTableNamesCache() + LoadingCache> getTableNamesWithParametersCache() { - return allTableNamesCache; + return tableNamesWithParametersCache; } - LoadingCache> getTablesWithParameterCache() + public LoadingCache> getTablesCacheNew() { - return tablesWithParameterCache; + return tablesCacheNew; } - Cache> getTableStatisticsCache() + Cache>> getTableColumnStatisticsCache() { - return tableStatisticsCache; + return tableColumnStatisticsCache; } - Cache> getPartitionStatisticsCache() + Cache>> getPartitionStatisticsCache() { return partitionStatisticsCache; } - LoadingCache> getViewNamesCache() - { - return viewNamesCache; - } - - LoadingCache>> getAllViewNamesCache() - { - return allViewNamesCache; - } - Cache>> getPartitionCache() { return partitionCache; @@ -1614,41 +1514,43 @@ LoadingCache> getRoleGrantsCache() return roleGrantsCache; } - LoadingCache> getGrantedPrincipalsCache() - { - return grantedPrincipalsCache; - } - LoadingCache> getConfigValuesCache() { return configValuesCache; } - private static class CacheFactory + private record CacheFactory( + OptionalLong expiresAfterWriteMillis, + OptionalLong refreshMillis, + Optional refreshExecutor, + long maximumSize, + StatsRecording statsRecording) { - private final OptionalLong expiresAfterWriteMillis; - private final OptionalLong refreshMillis; - private final Optional refreshExecutor; - private final long maximumSize; - private final StatsRecording statsRecording; + private static final CacheFactory NEVER_CACHE = new CacheFactory(OptionalLong.empty(), OptionalLong.empty(), Optional.empty(), 0, StatsRecording.DISABLED); - public CacheFactory(OptionalLong expiresAfterWriteMillis, OptionalLong refreshMillis, Optional refreshExecutor, long maximumSize, StatsRecording statsRecording) + private CacheFactory(long maximumSize) { - this.expiresAfterWriteMillis = requireNonNull(expiresAfterWriteMillis, "expiresAfterWriteMillis is null"); - this.refreshMillis = requireNonNull(refreshMillis, "refreshMillis is null"); - this.refreshExecutor = requireNonNull(refreshExecutor, "refreshExecutor is null"); - this.maximumSize = maximumSize; - this.statsRecording = requireNonNull(statsRecording, "statsRecording is null"); + this(OptionalLong.empty(), OptionalLong.empty(), Optional.empty(), maximumSize, StatsRecording.DISABLED); + } + + private CacheFactory + { + requireNonNull(expiresAfterWriteMillis, "expiresAfterWriteMillis is null"); + checkArgument(expiresAfterWriteMillis.isEmpty() || expiresAfterWriteMillis.getAsLong() > 0, "expiresAfterWriteMillis must be empty or at least 1 millisecond"); + requireNonNull(refreshMillis, "refreshMillis is null"); + checkArgument(refreshMillis.isEmpty() || refreshMillis.getAsLong() > 0, "refreshMillis must be empty or at least 1 millisecond"); + requireNonNull(refreshExecutor, "refreshExecutor is null"); + requireNonNull(statsRecording, "statsRecording is null"); } - public LoadingCache buildCache(com.google.common.base.Function loader) + public LoadingCache buildCache(Function loader) { - return CachingHiveMetastore.buildCache(expiresAfterWriteMillis, refreshMillis, refreshExecutor, maximumSize, statsRecording, CacheLoader.from(loader)); + return CachingHiveMetastore.buildCache(expiresAfterWriteMillis, refreshMillis, refreshExecutor, maximumSize, statsRecording, CacheLoader.from(loader::apply)); } - public Cache buildCache(BiFunction reloader) + public Cache buildCache(BiFunction loader) { - CacheLoader onlyReloader = new CacheLoader<>() + CacheLoader cacheLoader = new CacheLoader<>() { @Override public V load(K key) @@ -1662,10 +1564,10 @@ public ListenableFuture reload(K key, V oldValue) requireNonNull(key); requireNonNull(oldValue); // async reloading is configured in CachingHiveMetastore.buildCache if refreshMillis is present - return immediateFuture(reloader.apply(key, oldValue)); + return immediateFuture(loader.apply(key, oldValue)); } }; - return CachingHiveMetastore.buildCache(expiresAfterWriteMillis, refreshMillis, refreshExecutor, maximumSize, statsRecording, onlyReloader); + return CachingHiveMetastore.buildCache(expiresAfterWriteMillis, refreshMillis, refreshExecutor, maximumSize, statsRecording, cacheLoader); } public Cache> buildBulkCache() diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/cache/CachingHiveMetastoreConfig.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/cache/CachingHiveMetastoreConfig.java index eaae97eed9b5..853bceec16d8 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/cache/CachingHiveMetastoreConfig.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/cache/CachingHiveMetastoreConfig.java @@ -22,6 +22,8 @@ import java.util.Optional; import java.util.concurrent.TimeUnit; +import static com.google.common.base.MoreObjects.firstNonNull; + public class CachingHiveMetastoreConfig { private Duration metastoreCacheTtl = new Duration(0, TimeUnit.SECONDS); @@ -35,6 +37,8 @@ public class CachingHiveMetastoreConfig private int maxMetastoreRefreshThreads = 10; private boolean cacheMissing = true; private boolean partitionCacheEnabled = true; + private Boolean cacheMissingPartitions; + private Boolean cacheMissingStats; @NotNull public Duration getMetastoreCacheTtl() @@ -124,4 +128,28 @@ public CachingHiveMetastoreConfig setPartitionCacheEnabled(boolean enabled) this.partitionCacheEnabled = enabled; return this; } + + public boolean isCacheMissingPartitions() + { + return firstNonNull(cacheMissingPartitions, cacheMissing); + } + + @Config("hive.metastore-cache.cache-missing-partitions") + public CachingHiveMetastoreConfig setCacheMissingPartitions(boolean cacheMissingPartitions) + { + this.cacheMissingPartitions = cacheMissingPartitions; + return this; + } + + public boolean isCacheMissingStats() + { + return firstNonNull(cacheMissingStats, cacheMissing); + } + + @Config("hive.metastore-cache.cache-missing-stats") + public CachingHiveMetastoreConfig setCacheMissingStats(boolean cacheMissingStats) + { + this.cacheMissingStats = cacheMissingStats; + return this; + } } diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/cache/ReentrantBoundedExecutor.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/cache/ReentrantBoundedExecutor.java index 7cdd4aa6ef5e..6005d3c77d6d 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/cache/ReentrantBoundedExecutor.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/cache/ReentrantBoundedExecutor.java @@ -30,7 +30,7 @@ public class ReentrantBoundedExecutor private final Executor boundedExecutor; private final Executor coreExecutor; - ReentrantBoundedExecutor(Executor coreExecutor, int maxThreads) + public ReentrantBoundedExecutor(Executor coreExecutor, int maxThreads) { this.boundedExecutor = new BoundedExecutor(requireNonNull(coreExecutor, "coreExecutor is null"), maxThreads); this.coreExecutor = coreExecutor; diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/cache/SharedHiveMetastoreCache.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/cache/SharedHiveMetastoreCache.java index 7c2ab0ece44f..88ef7d87cce9 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/cache/SharedHiveMetastoreCache.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/cache/SharedHiveMetastoreCache.java @@ -18,6 +18,7 @@ import com.google.common.cache.CacheLoader; import com.google.common.cache.CacheStats; import com.google.common.cache.LoadingCache; +import com.google.common.collect.ImmutableSet; import com.google.common.math.LongMath; import com.google.common.util.concurrent.UncheckedExecutionException; import com.google.inject.Inject; @@ -25,7 +26,7 @@ import io.trino.plugin.base.CatalogName; import io.trino.plugin.hive.metastore.HiveMetastore; import io.trino.plugin.hive.metastore.HiveMetastoreFactory; -import io.trino.plugin.hive.metastore.cache.CachingHiveMetastore.CachingHiveMetastoreBuilder; +import io.trino.plugin.hive.metastore.cache.CachingHiveMetastore.ObjectType; import io.trino.spi.NodeManager; import io.trino.spi.TrinoException; import io.trino.spi.security.ConnectorIdentity; @@ -36,6 +37,7 @@ import org.weakref.jmx.Nested; import java.util.Optional; +import java.util.Set; import java.util.concurrent.ExecutorService; import java.util.function.Function; @@ -51,11 +53,18 @@ public class SharedHiveMetastoreCache { private final boolean enabled; private final CatalogName catalogName; + + private final Duration metadataCacheTtl; + private final Duration statsCacheTtl; + + private final Optional metastoreRefreshInterval; + private final long metastoreCacheMaximumSize; private final int maxMetastoreRefreshThreads; private final Duration userMetastoreCacheTtl; private final long userMetastoreCacheMaximumSize; - private final CachingHiveMetastoreBuilder cachingMetastoreBuilder; + private final boolean metastorePartitionCacheEnabled; + private final Set cacheMissing; private ExecutorService executorService; @@ -70,34 +79,33 @@ public SharedHiveMetastoreCache( requireNonNull(catalogName, "catalogName is null"); this.catalogName = catalogName; + + metadataCacheTtl = config.getMetastoreCacheTtl(); + statsCacheTtl = config.getStatsCacheTtl(); maxMetastoreRefreshThreads = config.getMaxMetastoreRefreshThreads(); + metastoreRefreshInterval = config.getMetastoreRefreshInterval(); + metastoreCacheMaximumSize = config.getMetastoreCacheMaximumSize(); + metastorePartitionCacheEnabled = config.isPartitionCacheEnabled(); + ImmutableSet.Builder cacheMissing = ImmutableSet.builder(); + if (config.isCacheMissing()) { + cacheMissing.add(ObjectType.OTHER); + } + if (config.isCacheMissingPartitions()) { + cacheMissing.add(ObjectType.PARTITION); + } + if (config.isCacheMissingStats()) { + cacheMissing.add(ObjectType.STATS); + } + this.cacheMissing = cacheMissing.build(); + userMetastoreCacheTtl = impersonationCachingConfig.getUserMetastoreCacheTtl(); userMetastoreCacheMaximumSize = impersonationCachingConfig.getUserMetastoreCacheMaximumSize(); // Disable caching on workers, because there currently is no way to invalidate such a cache. // Note: while we could skip CachingHiveMetastoreModule altogether on workers, we retain it so that catalog // configuration can remain identical for all nodes, making cluster configuration easier. - Duration metastoreCacheTtl = config.getMetastoreCacheTtl(); - Duration statsCacheTtl = config.getStatsCacheTtl(); - if (metastoreCacheTtl.compareTo(statsCacheTtl) > 0) { - statsCacheTtl = metastoreCacheTtl; - } - - boolean metadataCacheEnabled = metastoreCacheTtl.toMillis() > 0; - boolean statsCacheEnabled = statsCacheTtl.toMillis() > 0; - enabled = (metadataCacheEnabled || statsCacheEnabled) && - nodeManager.getCurrentNode().isCoordinator() && - config.getMetastoreCacheMaximumSize() > 0; - - cachingMetastoreBuilder = CachingHiveMetastore.builder() - .metadataCacheEnabled(metadataCacheEnabled) - .statsCacheEnabled(statsCacheEnabled) - .cacheTtl(metastoreCacheTtl) - .statsCacheTtl(statsCacheTtl) - .refreshInterval(config.getMetastoreRefreshInterval()) - .maximumSize(config.getMetastoreCacheMaximumSize()) - .cacheMissing(config.isCacheMissing()) - .partitionCacheEnabled(config.isPartitionCacheEnabled()); + enabled = nodeManager.getCurrentNode().isCoordinator() && + (metadataCacheTtl.toMillis() > 0 || statsCacheTtl.toMillis() > 0); } @PostConstruct @@ -124,7 +132,7 @@ public boolean isEnabled() public HiveMetastoreFactory createCachingHiveMetastoreFactory(HiveMetastoreFactory metastoreFactory) { - if (!enabled) { + if (!enabled || metastoreFactory.hasBuiltInCaching()) { return metastoreFactory; } @@ -133,17 +141,27 @@ public HiveMetastoreFactory createCachingHiveMetastoreFactory(HiveMetastoreFacto if (userMetastoreCacheMaximumSize == 0 || userMetastoreCacheTtl.toMillis() == 0) { return metastoreFactory; } - return new ImpersonationCachingHiveMetastoreFactory(metastoreFactory); + return new ImpersonationCachingHiveMetastoreFactory( + user -> createCachingHiveMetastore(metastoreFactory, Optional.of(ConnectorIdentity.ofUser(user))), + userMetastoreCacheTtl, + userMetastoreCacheMaximumSize); } - CachingHiveMetastore cachingHiveMetastore = CachingHiveMetastore.builder(cachingMetastoreBuilder) - // Loading of cache entry in CachingHiveMetastore might trigger loading of another cache entry for different object type - // In case there are no empty executor slots, such operation would deadlock. Therefore, a reentrant executor needs to be - // used. - .delegate(metastoreFactory.createMetastore(Optional.empty())) - .executor(new ReentrantBoundedExecutor(executorService, maxMetastoreRefreshThreads)) - .build(); - return new CachingHiveMetastoreFactory(cachingHiveMetastore); + return new CachingHiveMetastoreFactory(createCachingHiveMetastore(metastoreFactory, Optional.empty())); + } + + private CachingHiveMetastore createCachingHiveMetastore(HiveMetastoreFactory metastoreFactory, Optional identity) + { + return CachingHiveMetastore.createCachingHiveMetastore( + metastoreFactory.createMetastore(identity), + metadataCacheTtl, + statsCacheTtl, + metastoreRefreshInterval, + new ReentrantBoundedExecutor(executorService, maxMetastoreRefreshThreads), + metastoreCacheMaximumSize, + CachingHiveMetastore.StatsRecording.ENABLED, + metastorePartitionCacheEnabled, + cacheMissing); } public static class CachingHiveMetastoreFactory @@ -176,20 +194,18 @@ public CachingHiveMetastore getMetastore() } } - public class ImpersonationCachingHiveMetastoreFactory + public static class ImpersonationCachingHiveMetastoreFactory implements HiveMetastoreFactory { - private final HiveMetastoreFactory metastoreFactory; private final LoadingCache cache; - public ImpersonationCachingHiveMetastoreFactory(HiveMetastoreFactory metastoreFactory) + public ImpersonationCachingHiveMetastoreFactory(Function cachingHiveMetastoreFactory, Duration userCacheTtl, long userCacheMaximumSize) { - this.metastoreFactory = metastoreFactory; cache = buildNonEvictableCache( CacheBuilder.newBuilder() - .expireAfterWrite(userMetastoreCacheTtl.toMillis(), MILLISECONDS) - .maximumSize(userMetastoreCacheMaximumSize), - CacheLoader.from(this::createUserCachingMetastore)); + .expireAfterWrite(userCacheTtl.toMillis(), MILLISECONDS) + .maximumSize(userCacheMaximumSize), + CacheLoader.from(cachingHiveMetastoreFactory::apply)); } @Override @@ -211,15 +227,6 @@ public HiveMetastore createMetastore(Optional identity) } } - private CachingHiveMetastore createUserCachingMetastore(String user) - { - ConnectorIdentity identity = ConnectorIdentity.ofUser(user); - return CachingHiveMetastore.builder(cachingMetastoreBuilder) - .delegate(metastoreFactory.createMetastore(Optional.of(identity))) - .executor(new ReentrantBoundedExecutor(executorService, maxMetastoreRefreshThreads)) - .build(); - } - @Managed public void flushCache() { @@ -249,30 +256,23 @@ public AggregateCacheStatsMBean getTableStats() @Managed @Nested - public AggregateCacheStatsMBean getTableNamesStats() - { - return new AggregateCacheStatsMBean(CachingHiveMetastore::getTableNamesCache); - } - - @Managed - @Nested - public AggregateCacheStatsMBean getAllTableNamesStats() + public AggregateCacheStatsMBean getTablesStats() { - return new AggregateCacheStatsMBean(CachingHiveMetastore::getAllTableNamesCache); + return new AggregateCacheStatsMBean(CachingHiveMetastore::getTablesCacheNew); } @Managed @Nested public AggregateCacheStatsMBean getTableWithParameterStats() { - return new AggregateCacheStatsMBean(CachingHiveMetastore::getTablesWithParameterCache); + return new AggregateCacheStatsMBean(CachingHiveMetastore::getTableNamesWithParametersCache); } @Managed @Nested - public AggregateCacheStatsMBean getTableStatisticsStats() + public AggregateCacheStatsMBean getTableColumnStatisticsCache() { - return new AggregateCacheStatsMBean(CachingHiveMetastore::getTableStatisticsCache); + return new AggregateCacheStatsMBean(CachingHiveMetastore::getTableColumnStatisticsCache); } @Managed @@ -282,20 +282,6 @@ public AggregateCacheStatsMBean getPartitionStatisticsStats() return new AggregateCacheStatsMBean(CachingHiveMetastore::getPartitionStatisticsCache); } - @Managed - @Nested - public AggregateCacheStatsMBean getViewNamesStats() - { - return new AggregateCacheStatsMBean(CachingHiveMetastore::getViewNamesCache); - } - - @Managed - @Nested - public AggregateCacheStatsMBean getAllViewNamesStats() - { - return new AggregateCacheStatsMBean(CachingHiveMetastore::getAllViewNamesCache); - } - @Managed @Nested public AggregateCacheStatsMBean getPartitionStats() @@ -331,13 +317,6 @@ public AggregateCacheStatsMBean getRoleGrantsStats() return new AggregateCacheStatsMBean(CachingHiveMetastore::getRoleGrantsCache); } - @Managed - @Nested - public AggregateCacheStatsMBean getGrantedPrincipalsStats() - { - return new AggregateCacheStatsMBean(CachingHiveMetastore::getGrantedPrincipalsCache); - } - @Managed @Nested public AggregateCacheStatsMBean getConfigValuesStats() @@ -410,16 +389,6 @@ public long getRequestCount() return requestCount; } - public long getHitCount() - { - return hitCount; - } - - public long getMissCount() - { - return missCount; - } - public double getHitRate() { return (requestCount == 0) ? 1.0 : (double) hitCount / requestCount; diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/file/ColumnStatistics.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/file/ColumnStatistics.java new file mode 100644 index 000000000000..c559d515d7a0 --- /dev/null +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/file/ColumnStatistics.java @@ -0,0 +1,147 @@ +/* + * 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.hive.metastore.file; + +import com.google.errorprone.annotations.Immutable; +import io.trino.plugin.hive.HiveBasicStatistics; +import io.trino.plugin.hive.metastore.HiveColumnStatistics; + +import java.math.BigDecimal; +import java.time.LocalDate; +import java.util.LinkedHashSet; +import java.util.Optional; +import java.util.OptionalDouble; +import java.util.OptionalLong; +import java.util.Set; + +import static com.google.common.base.Preconditions.checkArgument; +import static java.util.Objects.requireNonNull; + +@Immutable +public record ColumnStatistics( + Optional integerStatistics, + Optional doubleStatistics, + Optional decimalStatistics, + Optional dateStatistics, + Optional booleanStatistics, + OptionalLong maxValueSizeInBytes, + OptionalDouble averageColumnLength, + OptionalLong totalSizeInBytes, + OptionalLong nullsCount, + OptionalLong distinctValuesCount) +{ + public ColumnStatistics + { + requireNonNull(integerStatistics, "integerStatistics is null"); + requireNonNull(doubleStatistics, "doubleStatistics is null"); + requireNonNull(decimalStatistics, "decimalStatistics is null"); + requireNonNull(dateStatistics, "dateStatistics is null"); + requireNonNull(booleanStatistics, "booleanStatistics is null"); + requireNonNull(maxValueSizeInBytes, "maxValueSizeInBytes is null"); + requireNonNull(averageColumnLength, "averageColumnLength is null"); + requireNonNull(totalSizeInBytes, "totalSizeInBytes is null"); + checkArgument(averageColumnLength.isEmpty() || totalSizeInBytes.isEmpty(), "both averageColumnLength and totalSizeInBytes are present"); + requireNonNull(nullsCount, "nullsCount is null"); + requireNonNull(distinctValuesCount, "distinctValuesCount is null"); + + Set presentStatistics = new LinkedHashSet<>(); + integerStatistics.ifPresent(ignore -> presentStatistics.add("integerStatistics")); + doubleStatistics.ifPresent(ignore -> presentStatistics.add("doubleStatistics")); + decimalStatistics.ifPresent(ignore -> presentStatistics.add("decimalStatistics")); + dateStatistics.ifPresent(ignore -> presentStatistics.add("dateStatistics")); + booleanStatistics.ifPresent(ignore -> presentStatistics.add("booleanStatistics")); + checkArgument(presentStatistics.size() <= 1, "multiple type specific statistic objects are present: %s", presentStatistics); + } + + public static ColumnStatistics fromHiveColumnStatistics(HiveColumnStatistics hiveColumnStatistics) + { + return new ColumnStatistics( + hiveColumnStatistics.getIntegerStatistics().map(stat -> new IntegerStatistics(stat.getMin(), stat.getMax())), + hiveColumnStatistics.getDoubleStatistics().map(stat -> new DoubleStatistics(stat.getMin(), stat.getMax())), + hiveColumnStatistics.getDecimalStatistics().map(stat -> new DecimalStatistics(stat.getMin(), stat.getMax())), + hiveColumnStatistics.getDateStatistics().map(stat -> new DateStatistics(stat.getMin(), stat.getMax())), + hiveColumnStatistics.getBooleanStatistics().map(stat -> new BooleanStatistics(stat.getTrueCount(), stat.getFalseCount())), + hiveColumnStatistics.getMaxValueSizeInBytes(), + hiveColumnStatistics.getAverageColumnLength(), + OptionalLong.empty(), + hiveColumnStatistics.getNullsCount(), + hiveColumnStatistics.getDistinctValuesWithNullCount()); + } + + public HiveColumnStatistics toHiveColumnStatistics(HiveBasicStatistics basicStatistics) + { + OptionalDouble averageColumnLength = this.averageColumnLength; + if (totalSizeInBytes.isPresent() && basicStatistics.getRowCount().orElse(0) > 0 && nullsCount().isPresent()) { + long nonNullCount = basicStatistics.getRowCount().getAsLong() - nullsCount().orElseThrow(); + if (nonNullCount > 0) { + averageColumnLength = OptionalDouble.of(totalSizeInBytes.getAsLong() / (double) nonNullCount); + } + } + return new HiveColumnStatistics( + integerStatistics.map(stat -> new io.trino.plugin.hive.metastore.IntegerStatistics(stat.min(), stat.max())), + doubleStatistics.map(stat -> new io.trino.plugin.hive.metastore.DoubleStatistics(stat.min(), stat.max())), + decimalStatistics.map(stat -> new io.trino.plugin.hive.metastore.DecimalStatistics(stat.min(), stat.max())), + dateStatistics.map(stat -> new io.trino.plugin.hive.metastore.DateStatistics(stat.min(), stat.max())), + booleanStatistics.map(stat -> new io.trino.plugin.hive.metastore.BooleanStatistics(stat.trueCount(), stat.falseCount())), + maxValueSizeInBytes, + averageColumnLength, + nullsCount, + distinctValuesCount); + } + + public record IntegerStatistics(OptionalLong min, OptionalLong max) + { + public IntegerStatistics + { + requireNonNull(min, "min is null"); + requireNonNull(max, "max is null"); + } + } + + public record DoubleStatistics(OptionalDouble min, OptionalDouble max) + { + public DoubleStatistics + { + requireNonNull(min, "min is null"); + requireNonNull(max, "max is null"); + } + } + + public record DecimalStatistics(Optional min, Optional max) + { + public DecimalStatistics + { + requireNonNull(min, "min is null"); + requireNonNull(max, "max is null"); + } + } + + public record DateStatistics(Optional min, Optional max) + { + public DateStatistics + { + requireNonNull(min, "min is null"); + requireNonNull(max, "max is null"); + } + } + + public record BooleanStatistics(OptionalLong trueCount, OptionalLong falseCount) + { + public BooleanStatistics + { + requireNonNull(trueCount, "trueCount is null"); + requireNonNull(falseCount, "falseCount is null"); + } + } +} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/file/FileHiveMetastore.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/file/FileHiveMetastore.java index c1d8f652c862..659c1df53dad 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/file/FileHiveMetastore.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/file/FileHiveMetastore.java @@ -13,22 +13,20 @@ */ package io.trino.plugin.hive.metastore.file; -import com.google.common.cache.CacheLoader; -import com.google.common.cache.LoadingCache; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; import com.google.common.collect.ImmutableSet.Builder; +import com.google.common.collect.Lists; import com.google.common.collect.Sets; -import com.google.common.io.ByteStreams; import com.google.errorprone.annotations.ThreadSafe; -import com.google.errorprone.annotations.concurrent.GuardedBy; import io.airlift.json.JsonCodec; -import io.trino.cache.EvictableCacheBuilder; -import io.trino.hdfs.HdfsContext; -import io.trino.hdfs.HdfsEnvironment; +import io.trino.filesystem.FileIterator; +import io.trino.filesystem.Location; +import io.trino.filesystem.TrinoFileSystem; +import io.trino.filesystem.TrinoFileSystemFactory; +import io.trino.filesystem.TrinoOutputFile; import io.trino.plugin.hive.HiveBasicStatistics; -import io.trino.plugin.hive.HiveColumnStatisticType; import io.trino.plugin.hive.HiveType; import io.trino.plugin.hive.NodeVersion; import io.trino.plugin.hive.PartitionNotFoundException; @@ -36,7 +34,6 @@ import io.trino.plugin.hive.SchemaAlreadyExistsException; import io.trino.plugin.hive.TableAlreadyExistsException; import io.trino.plugin.hive.TableType; -import io.trino.plugin.hive.acid.AcidTransaction; import io.trino.plugin.hive.metastore.Database; import io.trino.plugin.hive.metastore.HiveColumnStatistics; import io.trino.plugin.hive.metastore.HiveMetastore; @@ -46,9 +43,10 @@ import io.trino.plugin.hive.metastore.Partition; import io.trino.plugin.hive.metastore.PartitionWithStatistics; import io.trino.plugin.hive.metastore.PrincipalPrivileges; +import io.trino.plugin.hive.metastore.StatisticsUpdateMode; import io.trino.plugin.hive.metastore.Table; +import io.trino.plugin.hive.metastore.TableInfo; import io.trino.plugin.hive.metastore.file.FileHiveMetastoreConfig.VersionCompatibility; -import io.trino.plugin.hive.metastore.thrift.ThriftMetastoreUtil; import io.trino.spi.TrinoException; import io.trino.spi.connector.ColumnNotFoundException; import io.trino.spi.connector.SchemaNotFoundException; @@ -57,17 +55,12 @@ import io.trino.spi.predicate.TupleDomain; import io.trino.spi.security.ConnectorIdentity; import io.trino.spi.security.RoleGrant; -import io.trino.spi.type.Type; -import org.apache.hadoop.fs.FSDataInputStream; -import org.apache.hadoop.fs.FileStatus; -import org.apache.hadoop.fs.FileSystem; -import org.apache.hadoop.fs.Path; +import java.io.FileNotFoundException; import java.io.IOException; +import java.io.InputStream; import java.io.OutputStream; -import java.util.ArrayDeque; import java.util.ArrayList; -import java.util.Arrays; import java.util.Collection; import java.util.EnumSet; import java.util.HashMap; @@ -80,41 +73,44 @@ import java.util.Map.Entry; import java.util.Objects; import java.util.Optional; +import java.util.OptionalLong; import java.util.Set; import java.util.function.Function; import java.util.function.Predicate; import java.util.stream.Collectors; -import static com.google.common.base.MoreObjects.toStringHelper; import static com.google.common.base.Preconditions.checkArgument; +import static com.google.common.base.Throwables.getCausalChain; +import static com.google.common.base.Verify.verify; import static com.google.common.collect.ImmutableList.toImmutableList; import static com.google.common.collect.ImmutableMap.toImmutableMap; import static com.google.common.collect.ImmutableSet.toImmutableSet; import static io.trino.plugin.hive.HiveErrorCode.HIVE_CONCURRENT_MODIFICATION_DETECTED; import static io.trino.plugin.hive.HiveErrorCode.HIVE_METASTORE_ERROR; -import static io.trino.plugin.hive.HiveMetadata.TABLE_COMMENT; +import static io.trino.plugin.hive.HiveMetadata.TRINO_QUERY_ID_NAME; import static io.trino.plugin.hive.HivePartitionManager.extractPartitionValues; import static io.trino.plugin.hive.TableType.EXTERNAL_TABLE; import static io.trino.plugin.hive.TableType.MANAGED_TABLE; import static io.trino.plugin.hive.TableType.MATERIALIZED_VIEW; -import static io.trino.plugin.hive.TableType.VIRTUAL_VIEW; +import static io.trino.plugin.hive.ViewReaderUtil.isSomeKindOfAView; import static io.trino.plugin.hive.metastore.HivePrivilegeInfo.HivePrivilege.OWNERSHIP; +import static io.trino.plugin.hive.metastore.MetastoreUtil.getHiveBasicStatistics; import static io.trino.plugin.hive.metastore.MetastoreUtil.makePartitionName; +import static io.trino.plugin.hive.metastore.MetastoreUtil.updateStatisticsParameters; import static io.trino.plugin.hive.metastore.MetastoreUtil.verifyCanDropColumn; +import static io.trino.plugin.hive.metastore.Table.TABLE_COMMENT; +import static io.trino.plugin.hive.metastore.file.ColumnStatistics.fromHiveColumnStatistics; import static io.trino.plugin.hive.metastore.file.FileHiveMetastore.SchemaType.DATABASE; import static io.trino.plugin.hive.metastore.file.FileHiveMetastore.SchemaType.PARTITION; import static io.trino.plugin.hive.metastore.file.FileHiveMetastore.SchemaType.TABLE; import static io.trino.plugin.hive.metastore.file.FileHiveMetastoreConfig.VERSION_COMPATIBILITY_CONFIG; import static io.trino.plugin.hive.metastore.file.FileHiveMetastoreConfig.VersionCompatibility.UNSAFE_ASSUME_COMPATIBILITY; -import static io.trino.plugin.hive.metastore.thrift.ThriftMetastoreUtil.getHiveBasicStatistics; -import static io.trino.plugin.hive.metastore.thrift.ThriftMetastoreUtil.updateStatisticsParameters; import static io.trino.plugin.hive.util.HiveUtil.DELTA_LAKE_PROVIDER; import static io.trino.plugin.hive.util.HiveUtil.SPARK_TABLE_PROVIDER_KEY; import static io.trino.plugin.hive.util.HiveUtil.escapeSchemaName; import static io.trino.plugin.hive.util.HiveUtil.escapeTableName; import static io.trino.plugin.hive.util.HiveUtil.isIcebergTable; import static io.trino.plugin.hive.util.HiveUtil.toPartitionValues; -import static io.trino.plugin.hive.util.HiveUtil.unescapePathName; import static io.trino.spi.StandardErrorCode.ALREADY_EXISTS; import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED; import static io.trino.spi.security.PrincipalType.ROLE; @@ -122,7 +118,6 @@ import static java.lang.String.format; import static java.util.Locale.ENGLISH; import static java.util.Objects.requireNonNull; -import static java.util.concurrent.TimeUnit.SECONDS; import static java.util.stream.Collectors.toList; import static java.util.stream.Collectors.toSet; @@ -134,8 +129,9 @@ public class FileHiveMetastore private static final String ADMIN_ROLE_NAME = "admin"; private static final String TRINO_SCHEMA_FILE_NAME_SUFFIX = ".trinoSchema"; private static final String TRINO_PERMISSIONS_DIRECTORY_NAME = ".trinoPermissions"; - public static final String ROLES_FILE_NAME = ".roles"; - public static final String ROLE_GRANTS_FILE_NAME = ".roleGrants"; + private static final String TRINO_FUNCTIONS_DIRECTORY_NAME = ".trinoFunction"; + private static final String ROLES_FILE_NAME = ".roles"; + private static final String ROLE_GRANTS_FILE_NAME = ".roleGrants"; // todo there should be a way to manage the admins list private static final Set ADMIN_USERS = ImmutableSet.of("admin", "hive", "hdfs"); @@ -144,12 +140,10 @@ public class FileHiveMetastore private final String currentVersion; private final VersionCompatibility versionCompatibility; - private final HdfsEnvironment hdfsEnvironment; - private final Path catalogDirectory; + private final TrinoFileSystem fileSystem; + private final Location catalogDirectory; private final boolean disableLocationChecks; - private final HdfsContext hdfsContext; private final boolean hideDeltaLakeTables; - private final FileSystem metadataFileSystem; private final JsonCodec databaseCodec = JsonCodec.jsonCodec(DatabaseMetadata.class); private final JsonCodec tableCodec = JsonCodec.jsonCodec(TableMetadata.class); @@ -158,28 +152,14 @@ public class FileHiveMetastore private final JsonCodec> rolesCodec = JsonCodec.listJsonCodec(String.class); private final JsonCodec> roleGrantsCodec = JsonCodec.listJsonCodec(RoleGrant.class); - // TODO Remove this speed-up workaround once that https://github.com/trinodb/trino/issues/13115 gets implemented - private final LoadingCache> listTablesCache; - - public FileHiveMetastore(NodeVersion nodeVersion, HdfsEnvironment hdfsEnvironment, boolean hideDeltaLakeTables, FileHiveMetastoreConfig config) + public FileHiveMetastore(NodeVersion nodeVersion, TrinoFileSystemFactory fileSystemFactory, boolean hideDeltaLakeTables, FileHiveMetastoreConfig config) { this.currentVersion = nodeVersion.toString(); this.versionCompatibility = requireNonNull(config.getVersionCompatibility(), "config.getVersionCompatibility() is null"); - this.hdfsEnvironment = requireNonNull(hdfsEnvironment, "hdfsEnvironment is null"); - this.catalogDirectory = new Path(requireNonNull(config.getCatalogDirectory(), "catalogDirectory is null")); + this.fileSystem = fileSystemFactory.create(ConnectorIdentity.ofUser(config.getMetastoreUser())); + this.catalogDirectory = Location.of(requireNonNull(config.getCatalogDirectory(), "catalogDirectory is null")); this.disableLocationChecks = config.isDisableLocationChecks(); - this.hdfsContext = new HdfsContext(ConnectorIdentity.ofUser(config.getMetastoreUser())); this.hideDeltaLakeTables = hideDeltaLakeTables; - try { - metadataFileSystem = hdfsEnvironment.getFileSystem(hdfsContext, this.catalogDirectory); - } - catch (IOException e) { - throw new TrinoException(HIVE_METASTORE_ERROR, e); - } - - listTablesCache = EvictableCacheBuilder.newBuilder() - .expireAfterWrite(10, SECONDS) - .build(CacheLoader.from(this::doListAllTables)); } @Override @@ -196,12 +176,23 @@ public synchronized void createDatabase(Database database) database.getParameters()); verifyDatabaseNameLength(database.getDatabaseName()); - verifyDatabaseNotExists(database.getDatabaseName()); - Path databaseMetadataDirectory = getDatabaseMetadataDirectory(database.getDatabaseName()); + Optional existingDatabase = getDatabase(database.getDatabaseName()); + if (existingDatabase.isPresent()) { + // Do not throw SchemaAlreadyExistsException if this query has already created the database. + // This may happen when an actually successful metastore create call is retried + // because of a timeout on our side. + String expectedQueryId = database.getParameters().get(TRINO_QUERY_ID_NAME); + if (expectedQueryId != null && expectedQueryId.equals(existingDatabase.get().getParameters().get(TRINO_QUERY_ID_NAME))) { + return; + } + throw new SchemaAlreadyExistsException(database.getDatabaseName()); + } + + Location databaseMetadataDirectory = getDatabaseMetadataDirectory(database.getDatabaseName()); writeSchemaFile(DATABASE, databaseMetadataDirectory, databaseCodec, new DatabaseMetadata(currentVersion, database), false); try { - metadataFileSystem.mkdirs(databaseMetadataDirectory); + fileSystem.createDirectory(databaseMetadataDirectory); } catch (IOException e) { throw new TrinoException(HIVE_METASTORE_ERROR, "Could not write database", e); @@ -217,7 +208,7 @@ public synchronized void dropDatabase(String databaseName, boolean deleteData) databaseName = databaseName.toLowerCase(ENGLISH); getRequiredDatabase(databaseName); - if (!getAllTables(databaseName).isEmpty()) { + if (!getTables(databaseName).isEmpty()) { throw new TrinoException(HIVE_METASTORE_ERROR, "Database " + databaseName + " is not empty"); } @@ -240,14 +231,11 @@ public synchronized void renameDatabase(String databaseName, String newDatabaseN getRequiredDatabase(databaseName); verifyDatabaseNotExists(newDatabaseName); - Path oldDatabaseMetadataDirectory = getDatabaseMetadataDirectory(databaseName); - Path newDatabaseMetadataDirectory = getDatabaseMetadataDirectory(newDatabaseName); + Location oldDatabaseMetadataDirectory = getDatabaseMetadataDirectory(databaseName); + Location newDatabaseMetadataDirectory = getDatabaseMetadataDirectory(newDatabaseName); try { renameSchemaFile(DATABASE, oldDatabaseMetadataDirectory, newDatabaseMetadataDirectory); - - if (!metadataFileSystem.rename(oldDatabaseMetadataDirectory, newDatabaseMetadataDirectory)) { - throw new TrinoException(HIVE_METASTORE_ERROR, "Could not rename database metadata directory"); - } + fileSystem.renameDirectory(oldDatabaseMetadataDirectory, newDatabaseMetadataDirectory); } catch (IOException e) { throw new TrinoException(HIVE_METASTORE_ERROR, e); @@ -258,7 +246,7 @@ public synchronized void renameDatabase(String databaseName, String newDatabaseN public synchronized void setDatabaseOwner(String databaseName, HivePrincipal principal) { Database database = getRequiredDatabase(databaseName); - Path databaseMetadataDirectory = getDatabaseMetadataDirectory(database.getDatabaseName()); + Location databaseMetadataDirectory = getDatabaseMetadataDirectory(database.getDatabaseName()); Database newDatabase = Database.builder(database) .setOwnerName(Optional.of(principal.getName())) .setOwnerType(Optional.of(principal.getType())) @@ -275,7 +263,7 @@ public synchronized Optional getDatabase(String databaseName) // Database names are stored lowercase. Accept non-lowercase name for compatibility with HMS (and Glue) String normalizedName = databaseName.toLowerCase(ENGLISH); - Path databaseMetadataDirectory = getDatabaseMetadataDirectory(normalizedName); + Location databaseMetadataDirectory = getDatabaseMetadataDirectory(normalizedName); return readSchemaFile(DATABASE, databaseMetadataDirectory, databaseCodec) .map(databaseMetadata -> { checkVersion(databaseMetadata.getWriterVersion()); @@ -289,14 +277,14 @@ private Database getRequiredDatabase(String databaseName) .orElseThrow(() -> new SchemaNotFoundException(databaseName)); } - private void verifyDatabaseNameLength(String databaseName) + private static void verifyDatabaseNameLength(String databaseName) { if (databaseName.length() > MAX_NAME_LENGTH) { throw new TrinoException(NOT_SUPPORTED, format("Schema name must be shorter than or equal to '%s' characters but got '%s'", MAX_NAME_LENGTH, databaseName.length())); } } - private void verifyTableNameLength(String tableName) + private static void verifyTableNameLength(String tableName) { if (tableName.length() > MAX_NAME_LENGTH) { throw new TrinoException(NOT_SUPPORTED, format("Table name must be shorter than or equal to '%s' characters but got '%s'", MAX_NAME_LENGTH, tableName.length())); @@ -313,9 +301,31 @@ private void verifyDatabaseNotExists(String databaseName) @Override public synchronized List getAllDatabases() { - return getChildSchemaDirectories(DATABASE, catalogDirectory).stream() - .map(Path::getName) - .collect(toImmutableList()); + try { + String prefix = catalogDirectory.toString(); + Set databases = new HashSet<>(); + + FileIterator iterator = fileSystem.listFiles(catalogDirectory); + while (iterator.hasNext()) { + Location location = iterator.next().location(); + + String child = location.toString().substring(prefix.length()); + if (child.startsWith("/")) { + child = child.substring(1); + } + + int length = child.length() - TRINO_SCHEMA_FILE_NAME_SUFFIX.length(); + if ((length > 1) && !child.contains("/") && child.startsWith(".") && + child.endsWith(TRINO_SCHEMA_FILE_NAME_SUFFIX)) { + databases.add(child.substring(1, length)); + } + } + + return ImmutableList.copyOf(databases); + } + catch (IOException e) { + throw new TrinoException(HIVE_METASTORE_ERROR, e); + } } @Override @@ -323,25 +333,35 @@ public synchronized void createTable(Table table, PrincipalPrivileges principalP { verifyTableNameLength(table.getTableName()); verifyDatabaseExists(table.getDatabaseName()); - verifyTableNotExists(table.getDatabaseName(), table.getTableName()); - Path tableMetadataDirectory = getTableMetadataDirectory(table); + Optional
existingTable = getTable(table.getDatabaseName(), table.getTableName()); + if (existingTable.isPresent()) { + // Do not throw TableAlreadyExistsException if this query has already created the table. + // This may happen when an actually successful metastore create call is retried + // because of a timeout on our side. + String expectedQueryId = table.getParameters().get(TRINO_QUERY_ID_NAME); + if (expectedQueryId != null && expectedQueryId.equals(existingTable.get().getParameters().get(TRINO_QUERY_ID_NAME))) { + return; + } + throw new TableAlreadyExistsException(new SchemaTableName(table.getDatabaseName(), table.getTableName())); + } + + Location tableMetadataDirectory = getTableMetadataDirectory(table); // validate table location - if (table.getTableType().equals(VIRTUAL_VIEW.name())) { + if (isSomeKindOfAView(table)) { checkArgument(table.getStorage().getLocation().isEmpty(), "Storage location for view must be empty"); } else if (table.getTableType().equals(MANAGED_TABLE.name())) { - if (!disableLocationChecks && !(new Path(table.getStorage().getLocation()).toString().contains(tableMetadataDirectory.toString()))) { + if (!disableLocationChecks && !table.getStorage().getLocation().contains(tableMetadataDirectory.toString())) { throw new TrinoException(HIVE_METASTORE_ERROR, "Table directory must be " + tableMetadataDirectory); } } else if (table.getTableType().equals(EXTERNAL_TABLE.name())) { if (!disableLocationChecks) { try { - Path externalLocation = new Path(table.getStorage().getLocation()); - FileSystem externalFileSystem = hdfsEnvironment.getFileSystem(hdfsContext, externalLocation); - if (!externalFileSystem.isDirectory(externalLocation)) { + Location externalLocation = Location.of(table.getStorage().getLocation()); + if (!fileSystem.directoryExists(externalLocation).orElse(true)) { throw new TrinoException(HIVE_METASTORE_ERROR, "External table location does not exist"); } } @@ -370,7 +390,7 @@ public synchronized Optional
getTable(String databaseName, String tableNa requireNonNull(databaseName, "databaseName is null"); requireNonNull(tableName, "tableName is null"); - Path tableMetadataDirectory = getTableMetadataDirectory(databaseName, tableName); + Location tableMetadataDirectory = getTableMetadataDirectory(databaseName, tableName); return readSchemaFile(TABLE, tableMetadataDirectory, tableCodec) .map(tableMetadata -> { checkVersion(tableMetadata.getWriterVersion()); @@ -387,7 +407,7 @@ public synchronized void setTableOwner(String databaseName, String tableName, Hi } Table table = getRequiredTable(databaseName, tableName); - Path tableMetadataDirectory = getTableMetadataDirectory(table); + Location tableMetadataDirectory = getTableMetadataDirectory(table); Table newTable = Table.builder(table) .setOwner(Optional.of(principal.getName())) .build(); @@ -396,42 +416,33 @@ public synchronized void setTableOwner(String databaseName, String tableName, Hi } @Override - public Set getSupportedColumnStatistics(Type type) - { - return ThriftMetastoreUtil.getSupportedColumnStatistics(type); - } - - @Override - public synchronized PartitionStatistics getTableStatistics(Table table) + public synchronized Map getTableColumnStatistics(String databaseName, String tableName, Set columnNames) { - return getTableStatistics(table.getDatabaseName(), table.getTableName()); - } - - private synchronized PartitionStatistics getTableStatistics(String databaseName, String tableName) - { - Path tableMetadataDirectory = getTableMetadataDirectory(databaseName, tableName); + checkArgument(!columnNames.isEmpty(), "columnNames is empty"); + Location tableMetadataDirectory = getTableMetadataDirectory(databaseName, tableName); TableMetadata tableMetadata = readSchemaFile(TABLE, tableMetadataDirectory, tableCodec) .orElseThrow(() -> new TableNotFoundException(new SchemaTableName(databaseName, tableName))); checkVersion(tableMetadata.getWriterVersion()); - HiveBasicStatistics basicStatistics = getHiveBasicStatistics(tableMetadata.getParameters()); - Map columnStatistics = tableMetadata.getColumnStatistics(); - return new PartitionStatistics(basicStatistics, columnStatistics); + return toHiveColumnStats(columnNames, tableMetadata.getParameters(), tableMetadata.getColumnStatistics()); } @Override - public synchronized Map getPartitionStatistics(Table table, List partitions) + public synchronized Map> getPartitionColumnStatistics(String databaseName, String tableName, Set partitionNames, Set columnNames) { - return partitions.stream() - .collect(toImmutableMap(partition -> makePartitionName(table, partition), partition -> getPartitionStatisticsInternal(table, partition.getValues()))); + checkArgument(!columnNames.isEmpty(), "columnNames is empty"); + ImmutableMap.Builder> result = ImmutableMap.builder(); + for (String partitionName : partitionNames) { + result.put(partitionName, getPartitionStatisticsInternal(databaseName, tableName, partitionName, columnNames)); + } + return result.buildOrThrow(); } - private synchronized PartitionStatistics getPartitionStatisticsInternal(Table table, List partitionValues) + private synchronized Map getPartitionStatisticsInternal(String databaseName, String tableName, String partitionName, Set columnNames) { - Path partitionDirectory = getPartitionMetadataDirectory(table, ImmutableList.copyOf(partitionValues)); + Location partitionDirectory = getPartitionMetadataDirectory(databaseName, tableName, partitionName); PartitionMetadata partitionMetadata = readSchemaFile(PARTITION, partitionDirectory, partitionCodec) - .orElseThrow(() -> new PartitionNotFoundException(table.getSchemaTableName(), partitionValues)); - HiveBasicStatistics basicStatistics = getHiveBasicStatistics(partitionMetadata.getParameters()); - return new PartitionStatistics(basicStatistics, partitionMetadata.getColumnStatistics()); + .orElseThrow(() -> new PartitionNotFoundException(new SchemaTableName(databaseName, tableName), extractPartitionValues(partitionName))); + return toHiveColumnStats(columnNames, partitionMetadata.getParameters(), partitionMetadata.getColumnStatistics()); } private Table getRequiredTable(String databaseName, String tableName) @@ -455,84 +466,58 @@ private void verifyTableNotExists(String newDatabaseName, String newTableName) } @Override - public synchronized void updateTableStatistics(String databaseName, String tableName, AcidTransaction transaction, Function update) + public synchronized void updateTableStatistics(String databaseName, String tableName, OptionalLong acidWriteId, StatisticsUpdateMode mode, PartitionStatistics statisticsUpdate) { - PartitionStatistics originalStatistics = getTableStatistics(databaseName, tableName); - PartitionStatistics updatedStatistics = update.apply(originalStatistics); - - Path tableMetadataDirectory = getTableMetadataDirectory(databaseName, tableName); + Location tableMetadataDirectory = getTableMetadataDirectory(databaseName, tableName); TableMetadata tableMetadata = readSchemaFile(TABLE, tableMetadataDirectory, tableCodec) .orElseThrow(() -> new TableNotFoundException(new SchemaTableName(databaseName, tableName))); checkVersion(tableMetadata.getWriterVersion()); + PartitionStatistics originalStatistics = toHivePartitionStatistics(tableMetadata.getParameters(), tableMetadata.getColumnStatistics()); + PartitionStatistics updatedStatistics = mode.updatePartitionStatistics(originalStatistics, statisticsUpdate); + TableMetadata updatedMetadata = tableMetadata .withParameters(currentVersion, updateStatisticsParameters(tableMetadata.getParameters(), updatedStatistics.getBasicStatistics())) - .withColumnStatistics(currentVersion, updatedStatistics.getColumnStatistics()); + .withColumnStatistics(currentVersion, fromHiveColumnStats(updatedStatistics.getColumnStatistics())); writeSchemaFile(TABLE, tableMetadataDirectory, tableCodec, updatedMetadata, true); } @Override - public synchronized void updatePartitionStatistics(Table table, Map> updates) + public synchronized void updatePartitionStatistics(Table table, StatisticsUpdateMode mode, Map partitionUpdates) { - updates.forEach((partitionName, update) -> { - PartitionStatistics originalStatistics = getPartitionStatisticsInternal(table, extractPartitionValues(partitionName)); - PartitionStatistics updatedStatistics = update.apply(originalStatistics); - - List partitionValues = extractPartitionValues(partitionName); - Path partitionDirectory = getPartitionMetadataDirectory(table, partitionValues); + partitionUpdates.forEach((partitionName, partitionUpdate) -> { + Location partitionDirectory = getPartitionMetadataDirectory(table, partitionName); PartitionMetadata partitionMetadata = readSchemaFile(PARTITION, partitionDirectory, partitionCodec) - .orElseThrow(() -> new PartitionNotFoundException(new SchemaTableName(table.getDatabaseName(), table.getTableName()), partitionValues)); + .orElseThrow(() -> new PartitionNotFoundException(table.getSchemaTableName(), extractPartitionValues(partitionName))); + PartitionStatistics originalStatistics = toHivePartitionStatistics(partitionMetadata.getParameters(), partitionMetadata.getColumnStatistics()); + + PartitionStatistics updatedStatistics = mode.updatePartitionStatistics(originalStatistics, partitionUpdate); PartitionMetadata updatedMetadata = partitionMetadata .withParameters(updateStatisticsParameters(partitionMetadata.getParameters(), updatedStatistics.getBasicStatistics())) - .withColumnStatistics(updatedStatistics.getColumnStatistics()); + .withColumnStatistics(fromHiveColumnStats(updatedStatistics.getColumnStatistics())); writeSchemaFile(PARTITION, partitionDirectory, partitionCodec, updatedMetadata, true); }); } @Override - public synchronized List getAllTables(String databaseName) - { - return listAllTables(databaseName).stream() - .filter(hideDeltaLakeTables - ? Predicate.not(ImmutableSet.copyOf(getTablesWithParameter(databaseName, SPARK_TABLE_PROVIDER_KEY, DELTA_LAKE_PROVIDER))::contains) - : tableName -> true) - .collect(toImmutableList()); - } - - @Override - public Optional> getAllTables() + public synchronized List getTables(String databaseName) { - return Optional.empty(); + return doListAllTables(databaseName, ignore -> true); } @Override - public synchronized List getTablesWithParameter(String databaseName, String parameterKey, String parameterValue) + public synchronized List getTableNamesWithParameters(String databaseName, String parameterKey, ImmutableSet parameterValues) { requireNonNull(parameterKey, "parameterKey is null"); - requireNonNull(parameterValue, "parameterValue is null"); - - List tables = listAllTables(databaseName); - - return tables.stream() - .map(tableName -> getTable(databaseName, tableName)) - .filter(Optional::isPresent) - .map(Optional::get) - .filter(table -> parameterValue.equals(table.getParameters().get(parameterKey))) - .map(Table::getTableName) + return doListAllTables(databaseName, table -> parameterValues.contains(table.getParameters().get(parameterKey))).stream() + .map(tableInfo -> tableInfo.tableName().getTableName()) .collect(toImmutableList()); } - @GuardedBy("this") - private List listAllTables(String databaseName) - { - return listTablesCache.getUnchecked(databaseName); - } - - @GuardedBy("this") - private List doListAllTables(String databaseName) + private synchronized List doListAllTables(String databaseName, Predicate tableMetadataPredicate) { requireNonNull(databaseName, "databaseName is null"); @@ -541,29 +526,37 @@ private List doListAllTables(String databaseName) return ImmutableList.of(); } - Path databaseMetadataDirectory = getDatabaseMetadataDirectory(databaseName); - List tables = getChildSchemaDirectories(TABLE, databaseMetadataDirectory).stream() - .map(Path::getName) - .collect(toImmutableList()); - return tables; - } + Location metadataDirectory = getDatabaseMetadataDirectory(databaseName); + try { + String prefix = metadataDirectory.toString(); + if (!prefix.endsWith("/")) { + prefix += "/"; + } + Set tables = new HashSet<>(); - @Override - public synchronized List getAllViews(String databaseName) - { - return getAllTables(databaseName).stream() - .map(tableName -> getTable(databaseName, tableName)) - .filter(Optional::isPresent) - .map(Optional::get) - .filter(table -> table.getTableType().equals(VIRTUAL_VIEW.name())) - .map(Table::getTableName) - .collect(toImmutableList()); - } + for (Location subdirectory : fileSystem.listDirectories(metadataDirectory)) { + String locationString = subdirectory.toString(); + verify(locationString.startsWith(prefix) && locationString.endsWith("/"), "Unexpected subdirectory %s when listing %s", subdirectory, metadataDirectory); - @Override - public Optional> getAllViews() - { - return Optional.empty(); + String tableName = locationString.substring(prefix.length(), locationString.length() - 1); + Location schemaFileLocation = subdirectory.appendPath(TRINO_SCHEMA_FILE_NAME_SUFFIX); + readFile("table schema", schemaFileLocation, tableCodec).ifPresent(tableMetadata -> { + checkVersion(tableMetadata.getWriterVersion()); + if ((hideDeltaLakeTables && DELTA_LAKE_PROVIDER.equals(tableMetadata.getParameters().get(SPARK_TABLE_PROVIDER_KEY))) + || !tableMetadataPredicate.test(tableMetadata)) { + return; + } + tables.add(new TableInfo( + new SchemaTableName(databaseName, tableName), + TableInfo.ExtendedRelationType.fromTableTypeAndComment(tableMetadata.getTableType(), tableMetadata.getParameters().get(TABLE_COMMENT)))); + }); + } + + return ImmutableList.copyOf(tables); + } + catch (IOException e) { + throw new TrinoException(HIVE_METASTORE_ERROR, e); + } } @Override @@ -574,7 +567,7 @@ public synchronized void dropTable(String databaseName, String tableName, boolea Table table = getRequiredTable(databaseName, tableName); - Path tableMetadataDirectory = getTableMetadataDirectory(databaseName, tableName); + Location tableMetadataDirectory = getTableMetadataDirectory(databaseName, tableName); if (deleteData) { deleteDirectoryAndSchema(TABLE, tableMetadataDirectory); @@ -586,7 +579,7 @@ public synchronized void dropTable(String databaseName, String tableName, boolea } @Override - public synchronized void replaceTable(String databaseName, String tableName, Table newTable, PrincipalPrivileges principalPrivileges) + public synchronized void replaceTable(String databaseName, String tableName, Table newTable, PrincipalPrivileges principalPrivileges, Map environmentContext) { Table table = getRequiredTable(databaseName, tableName); if (!table.getDatabaseName().equals(databaseName) || !table.getTableName().equals(tableName)) { @@ -597,7 +590,7 @@ public synchronized void replaceTable(String databaseName, String tableName, Tab throw new TrinoException(HIVE_CONCURRENT_MODIFICATION_DETECTED, "Cannot update Iceberg table: supplied previous location does not match current location"); } - Path tableMetadataDirectory = getTableMetadataDirectory(table); + Location tableMetadataDirectory = getTableMetadataDirectory(table); writeSchemaFile(TABLE, tableMetadataDirectory, tableCodec, new TableMetadata(currentVersion, newTable), true); // replace existing permissions @@ -626,32 +619,23 @@ public synchronized void renameTable(String databaseName, String tableName, Stri verifyTableNameLength(newTableName); verifyTableNotExists(newDatabaseName, newTableName); - Path oldPath = getTableMetadataDirectory(databaseName, tableName); - Path newPath = getTableMetadataDirectory(newDatabaseName, newTableName); + Location oldPath = getTableMetadataDirectory(databaseName, tableName); + Location newPath = getTableMetadataDirectory(newDatabaseName, newTableName); try { if (isIcebergTable(table)) { - if (!metadataFileSystem.mkdirs(newPath)) { - throw new TrinoException(HIVE_METASTORE_ERROR, "Could not create new table directory"); - } - // Iceberg metadata references files in old path, so these cannot be moved. Moving table description (metadata from metastore perspective) only. - if (!metadataFileSystem.rename(getSchemaPath(TABLE, oldPath), getSchemaPath(TABLE, newPath))) { - throw new TrinoException(HIVE_METASTORE_ERROR, "Could not rename table schema file"); - } + fileSystem.createDirectory(newPath); + // Iceberg metadata references files in the old path, so these cannot be moved. Moving table description (metadata from metastore perspective) only. + fileSystem.renameFile(getSchemaFile(TABLE, oldPath), getSchemaFile(TABLE, newPath)); // TODO drop data files when table is being dropped } else { - if (!metadataFileSystem.rename(oldPath, newPath)) { - throw new TrinoException(HIVE_METASTORE_ERROR, "Could not rename table directory"); - } + fileSystem.renameDirectory(oldPath, newPath); } } catch (IOException e) { throw new TrinoException(HIVE_METASTORE_ERROR, e); } - finally { - listTablesCache.invalidateAll(); - } } @Override @@ -660,7 +644,7 @@ public synchronized void commentTable(String databaseName, String tableName, Opt alterTable(databaseName, tableName, oldTable -> { Map parameters = oldTable.getParameters().entrySet().stream() .filter(entry -> !entry.getKey().equals(TABLE_COMMENT)) - .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)); + .collect(Collectors.toMap(Entry::getKey, Entry::getValue)); comment.ifPresent(value -> parameters.put(TABLE_COMMENT, value)); return oldTable.withParameters(currentVersion, parameters); @@ -670,24 +654,16 @@ public synchronized void commentTable(String databaseName, String tableName, Opt @Override public synchronized void commentColumn(String databaseName, String tableName, String columnName, Optional comment) { - alterTable(databaseName, tableName, oldTable -> { - if (oldTable.getColumn(columnName).isEmpty()) { - SchemaTableName name = new SchemaTableName(databaseName, tableName); - throw new ColumnNotFoundException(name, columnName); - } - - ImmutableList.Builder newDataColumns = ImmutableList.builder(); - for (Column fieldSchema : oldTable.getDataColumns()) { - if (fieldSchema.getName().equals(columnName)) { - newDataColumns.add(new Column(columnName, fieldSchema.getType(), comment, fieldSchema.getProperties())); - } - else { - newDataColumns.add(fieldSchema); - } - } + alterTable(databaseName, tableName, table -> table + .withDataColumns(currentVersion, updateColumnComment(table.getDataColumns(), columnName, comment)) + .withPartitionColumns(currentVersion, updateColumnComment(table.getPartitionColumns(), columnName, comment))); + } - return oldTable.withDataColumns(currentVersion, newDataColumns.build()); - }); + private static List updateColumnComment(List originalColumns, String columnName, Optional comment) + { + return Lists.transform(originalColumns, column -> column.getName().equals(columnName) + ? new Column(column.getName(), column.getType(), comment, column.getProperties()) + : column); } @Override @@ -764,7 +740,7 @@ private void alterTable(String databaseName, String tableName, Function new TableNotFoundException(new SchemaTableName(databaseName, tableName))); @@ -791,23 +767,24 @@ public synchronized void addPartitions(String databaseName, String tableName, Li checkArgument(EnumSet.of(MANAGED_TABLE, EXTERNAL_TABLE).contains(tableType), "Invalid table type: %s", tableType); try { - Map schemaFiles = new LinkedHashMap<>(); + Map schemaFiles = new LinkedHashMap<>(); for (PartitionWithStatistics partitionWithStatistics : partitions) { Partition partition = partitionWithStatistics.getPartition(); verifiedPartition(table, partition); - Path partitionMetadataDirectory = getPartitionMetadataDirectory(table, partition.getValues()); - Path schemaPath = getSchemaPath(PARTITION, partitionMetadataDirectory); - if (metadataFileSystem.exists(schemaPath)) { + Location partitionMetadataDirectory = getPartitionMetadataDirectory(table, partition.getValues()); + Location schemaPath = getSchemaFile(PARTITION, partitionMetadataDirectory); + + if (fileSystem.directoryExists(schemaPath).orElse(false)) { throw new TrinoException(HIVE_METASTORE_ERROR, "Partition already exists"); } byte[] schemaJson = partitionCodec.toJsonBytes(new PartitionMetadata(table, partitionWithStatistics)); schemaFiles.put(schemaPath, schemaJson); } - Set createdFiles = new LinkedHashSet<>(); + Set createdFiles = new LinkedHashSet<>(); try { - for (Entry entry : schemaFiles.entrySet()) { - try (OutputStream outputStream = metadataFileSystem.create(entry.getKey())) { + for (Entry entry : schemaFiles.entrySet()) { + try (OutputStream outputStream = fileSystem.newOutputFile(entry.getKey()).create()) { createdFiles.add(entry.getKey()); outputStream.write(entry.getValue()); } @@ -817,11 +794,12 @@ public synchronized void addPartitions(String databaseName, String tableName, Li } } catch (Throwable e) { - for (Path createdFile : createdFiles) { - try { - metadataFileSystem.delete(createdFile, false); - } - catch (IOException ignored) { + try { + fileSystem.deleteFiles(createdFiles); + } + catch (IOException ex) { + if (!e.equals(ex)) { + e.addSuppressed(ex); } } throw e; @@ -834,21 +812,20 @@ public synchronized void addPartitions(String databaseName, String tableName, Li private void verifiedPartition(Table table, Partition partition) { - Path partitionMetadataDirectory = getPartitionMetadataDirectory(table, partition.getValues()); + Location partitionMetadataDirectory = getPartitionMetadataDirectory(table, partition.getValues()); if (table.getTableType().equals(MANAGED_TABLE.name())) { - if (!partitionMetadataDirectory.equals(new Path(partition.getStorage().getLocation()))) { + if (!partitionMetadataDirectory.equals(Location.of(partition.getStorage().getLocation()))) { throw new TrinoException(HIVE_METASTORE_ERROR, "Partition directory must be " + partitionMetadataDirectory); } } else if (table.getTableType().equals(EXTERNAL_TABLE.name())) { try { - Path externalLocation = new Path(partition.getStorage().getLocation()); - FileSystem externalFileSystem = hdfsEnvironment.getFileSystem(hdfsContext, externalLocation); - if (!externalFileSystem.isDirectory(externalLocation)) { + Location externalLocation = Location.of(partition.getStorage().getLocation()); + if (!fileSystem.directoryExists(externalLocation).orElse(true)) { throw new TrinoException(HIVE_METASTORE_ERROR, "External partition location does not exist"); } - if (isChildDirectory(catalogDirectory, externalLocation)) { + if (externalLocation.toString().startsWith(catalogDirectory.toString())) { throw new TrinoException(HIVE_METASTORE_ERROR, "External partition location cannot be inside the system metadata directory"); } } @@ -874,7 +851,7 @@ public synchronized void dropPartition(String databaseName, String tableName, Li } Table table = tableReference.get(); - Path partitionMetadataDirectory = getPartitionMetadataDirectory(table, partitionValues); + Location partitionMetadataDirectory = getPartitionMetadataDirectory(table, partitionValues); if (deleteData) { deleteDirectoryAndSchema(PARTITION, partitionMetadataDirectory); } @@ -891,7 +868,7 @@ public synchronized void alterPartition(String databaseName, String tableName, P Partition partition = partitionWithStatistics.getPartition(); verifiedPartition(table, partition); - Path partitionMetadataDirectory = getPartitionMetadataDirectory(table, partition.getValues()); + Location partitionMetadataDirectory = getPartitionMetadataDirectory(table, partition.getValues()); writeSchemaFile(PARTITION, partitionMetadataDirectory, partitionCodec, new PartitionMetadata(table, partitionWithStatistics), true); } @@ -917,7 +894,7 @@ public synchronized void dropRole(String role) public synchronized Set listRoles() { Set roles = new HashSet<>(); - // Hive SQL standard assumes admin role already exists, so until that is fixed always add it here + // Hive SQL standard assumes the admin role already exists, so until that is fixed always add it here roles.add("admin"); readFile("roles", getRolesFile(), rolesCodec).ifPresent(roles::addAll); return ImmutableSet.copyOf(roles); @@ -983,14 +960,6 @@ public synchronized void revokeRoles(Set roles, Set grant } } - @Override - public synchronized Set listGrantedPrincipals(String role) - { - return listRoleGrantsSanitized().stream() - .filter(grant -> grant.getRoleName().equals(role)) - .collect(toImmutableSet()); - } - @Override public synchronized Set listRoleGrants(HivePrincipal principal) { @@ -1014,11 +983,11 @@ private synchronized Set listRoleGrantsSanitized() return removeDuplicatedEntries(removeNonExistingRoles(grants, existingRoles)); } - private Set removeDuplicatedEntries(Set grants) + private static Set removeDuplicatedEntries(Set grants) { - Map map = new HashMap<>(); + Map map = new HashMap<>(); for (RoleGrant grant : grants) { - RoleGranteeTuple tuple = new RoleGranteeTuple(grant.getRoleName(), HivePrincipal.from(grant.getGrantee())); + RoleGrantee tuple = new RoleGrantee(grant.getRoleName(), HivePrincipal.from(grant.getGrantee())); map.merge(tuple, grant, (first, second) -> first.isGrantable() ? first : second); } return ImmutableSet.copyOf(map.values()); @@ -1061,9 +1030,9 @@ private synchronized Optional> getAllPartitionNames(String database } Table table = tableReference.get(); - Path tableMetadataDirectory = getTableMetadataDirectory(table); + Location tableMetadataDirectory = getTableMetadataDirectory(table); - List> partitions = listPartitions(tableMetadataDirectory, table.getPartitionColumns()); + List> partitions = listPartitions(tableMetadataDirectory, table.getPartitionColumns()); List partitionNames = partitions.stream() .map(partitionValues -> makePartitionName(table.getPartitionColumns(), ImmutableList.copyOf(partitionValues))) @@ -1075,44 +1044,40 @@ private synchronized Optional> getAllPartitionNames(String database private boolean isValidPartition(Table table, String partitionName) { + Location location = getSchemaFile(PARTITION, getPartitionMetadataDirectory(table, partitionName)); try { - return metadataFileSystem.exists(getSchemaPath(PARTITION, getPartitionMetadataDirectory(table, partitionName))); + return fileSystem.newInputFile(location).exists(); } catch (IOException e) { return false; } } - private List> listPartitions(Path director, List partitionColumns) + private List> listPartitions(Location directory, List partitionColumns) { if (partitionColumns.isEmpty()) { return ImmutableList.of(); } try { - String directoryPrefix = partitionColumns.get(0).getName() + '='; - - List> partitionValues = new ArrayList<>(); - for (FileStatus fileStatus : metadataFileSystem.listStatus(director)) { - if (!fileStatus.isDirectory()) { - continue; - } - if (!fileStatus.getPath().getName().startsWith(directoryPrefix)) { - continue; + List> partitionValues = new ArrayList<>(); + FileIterator iterator = fileSystem.listFiles(directory); + while (iterator.hasNext()) { + Location location = iterator.next().location(); + String path = location.toString().substring(directory.toString().length()); + + if (path.startsWith("/")) { + path = path.substring(1); } - List> childPartitionValues; - if (partitionColumns.size() == 1) { - childPartitionValues = ImmutableList.of(new ArrayDeque<>()); - } - else { - childPartitionValues = listPartitions(fileStatus.getPath(), partitionColumns.subList(1, partitionColumns.size())); + if (!path.endsWith("/" + TRINO_SCHEMA_FILE_NAME_SUFFIX)) { + continue; } + path = path.substring(0, path.length() - TRINO_SCHEMA_FILE_NAME_SUFFIX.length() - 1); - String value = unescapePathName(fileStatus.getPath().getName().substring(directoryPrefix.length())); - for (ArrayDeque childPartition : childPartitionValues) { - childPartition.addFirst(value); - partitionValues.add(childPartition); + List values = toPartitionValues(path); + if (values.size() == partitionColumns.size()) { + partitionValues.add(values); } } return partitionValues; @@ -1128,7 +1093,7 @@ public synchronized Optional getPartition(Table table, List p requireNonNull(table, "table is null"); requireNonNull(partitionValues, "partitionValues is null"); - Path partitionDirectory = getPartitionMetadataDirectory(table, partitionValues); + Location partitionDirectory = getPartitionMetadataDirectory(table, partitionValues); return readSchemaFile(PARTITION, partitionDirectory, partitionCodec) .map(partitionMetadata -> partitionMetadata.toPartition(table.getDatabaseName(), table.getTableName(), partitionValues, partitionDirectory.toString())); } @@ -1158,7 +1123,7 @@ public synchronized Map> getPartitionsByNames(Table public synchronized Set listTablePrivileges(String databaseName, String tableName, Optional tableOwner, Optional principal) { Table table = getRequiredTable(databaseName, tableName); - Path permissionsDirectory = getPermissionsDirectory(table); + Location permissionsDirectory = getPermissionsDirectory(table); if (principal.isEmpty()) { Builder privileges = ImmutableSet.builder() .addAll(readAllPermissions(permissionsDirectory)); @@ -1210,14 +1175,11 @@ private synchronized void setTablePrivileges( try { Table table = getRequiredTable(databaseName, tableName); - Path permissionsDirectory = getPermissionsDirectory(table); + Location permissionsDirectory = getPermissionsDirectory(table); - boolean created = metadataFileSystem.mkdirs(permissionsDirectory); - if (!created && !metadataFileSystem.isDirectory(permissionsDirectory)) { - throw new TrinoException(HIVE_METASTORE_ERROR, "Could not create permissions directory"); - } + fileSystem.createDirectory(permissionsDirectory); - Path permissionFilePath = getPermissionsPath(permissionsDirectory, grantee); + Location permissionFilePath = getPermissionsPath(permissionsDirectory, grantee); List permissions = privileges.stream() .map(hivePrivilegeInfo -> new PermissionMetadata(hivePrivilegeInfo.getHivePrivilege(), hivePrivilegeInfo.isGrantOption(), grantee)) .collect(toList()); @@ -1231,67 +1193,44 @@ private synchronized void setTablePrivileges( private synchronized void deleteTablePrivileges(Table table) { try { - Path permissionsDirectory = getPermissionsDirectory(table); - metadataFileSystem.delete(permissionsDirectory, true); + Location permissionsDirectory = getPermissionsDirectory(table); + fileSystem.deleteDirectory(permissionsDirectory); } catch (IOException e) { throw new TrinoException(HIVE_METASTORE_ERROR, "Could not delete table permissions", e); } } - private List getChildSchemaDirectories(SchemaType type, Path metadataDirectory) - { - try { - if (!metadataFileSystem.isDirectory(metadataDirectory)) { - return ImmutableList.of(); - } - - ImmutableList.Builder childSchemaDirectories = ImmutableList.builder(); - for (FileStatus child : metadataFileSystem.listStatus(metadataDirectory)) { - if (!child.isDirectory()) { - continue; - } - Path childPath = child.getPath(); - if (childPath.getName().startsWith(".")) { - continue; - } - if (metadataFileSystem.isFile(getSchemaPath(type, childPath))) { - childSchemaDirectories.add(childPath); - } - } - return childSchemaDirectories.build(); - } - catch (IOException e) { - throw new TrinoException(HIVE_METASTORE_ERROR, e); - } - } - - private Set readPermissionsFile(Path permissionFilePath) + private Set readPermissionsFile(Location permissionFilePath) { return readFile("permissions", permissionFilePath, permissionsCodec).orElse(ImmutableList.of()).stream() .map(PermissionMetadata::toHivePrivilegeInfo) .collect(toImmutableSet()); } - private Set readAllPermissions(Path permissionsDirectory) + private Set readAllPermissions(Location permissionsDirectory) { try { - return Arrays.stream(metadataFileSystem.listStatus(permissionsDirectory)) - .filter(FileStatus::isFile) - .filter(file -> !file.getPath().getName().startsWith(".")) - .flatMap(file -> readPermissionsFile(file.getPath()).stream()) - .collect(toImmutableSet()); + ImmutableSet.Builder permissions = ImmutableSet.builder(); + FileIterator iterator = fileSystem.listFiles(permissionsDirectory); + while (iterator.hasNext()) { + Location location = iterator.next().location(); + if (!location.fileName().startsWith(".")) { + permissions.addAll(readPermissionsFile(location)); + } + } + return permissions.build(); } catch (IOException e) { throw new TrinoException(HIVE_METASTORE_ERROR, e); } } - private void deleteDirectoryAndSchema(SchemaType type, Path metadataDirectory) + private void deleteDirectoryAndSchema(SchemaType type, Location metadataDirectory) { try { - Path schemaPath = getSchemaPath(type, metadataDirectory); - if (!metadataFileSystem.isFile(schemaPath)) { + Location schemaPath = getSchemaFile(type, metadataDirectory); + if (!fileSystem.newInputFile(schemaPath).exists()) { // if there is no schema file, assume this is not a database, partition or table return; } @@ -1300,9 +1239,7 @@ private void deleteDirectoryAndSchema(SchemaType type, Path metadataDirectory) // (For cases when the schema file isn't in the metadata directory.) deleteSchemaFile(type, metadataDirectory); - if (!metadataFileSystem.delete(metadataDirectory, true)) { - throw new TrinoException(HIVE_METASTORE_ERROR, "Could not delete metadata directory"); - } + fileSystem.deleteDirectory(metadataDirectory); } catch (IOException e) { throw new TrinoException(HIVE_METASTORE_ERROR, e); @@ -1329,205 +1266,186 @@ private void checkVersion(Optional writerVersion) UNSAFE_ASSUME_COMPATIBILITY)); } - private Optional readSchemaFile(SchemaType type, Path metadataDirectory, JsonCodec codec) + private Optional readSchemaFile(SchemaType type, Location metadataDirectory, JsonCodec codec) { - return readFile(type + " schema", getSchemaPath(type, metadataDirectory), codec); + return readFile(type + " schema", getSchemaFile(type, metadataDirectory), codec); } - private Optional readFile(String type, Path path, JsonCodec codec) + private Optional readFile(String type, Location file, JsonCodec codec) { try { - if (!metadataFileSystem.isFile(path)) { - return Optional.empty(); - } - - try (FSDataInputStream inputStream = metadataFileSystem.open(path)) { - byte[] json = ByteStreams.toByteArray(inputStream); - return Optional.of(codec.fromJson(json)); + try (InputStream inputStream = fileSystem.newInputFile(file).newStream()) { + return Optional.of(codec.fromJson(inputStream.readAllBytes())); } } catch (Exception e) { + if (getCausalChain(e).stream().anyMatch(FileNotFoundException.class::isInstance)) { + return Optional.empty(); + } throw new TrinoException(HIVE_METASTORE_ERROR, "Could not read " + type, e); } } - private void writeSchemaFile(SchemaType type, Path directory, JsonCodec codec, T value, boolean overwrite) + private void writeSchemaFile(SchemaType type, Location directory, JsonCodec codec, T value, boolean overwrite) { - writeFile(type + " schema", getSchemaPath(type, directory), codec, value, overwrite); + writeFile(type + " schema", getSchemaFile(type, directory), codec, value, overwrite); } - private void writeFile(String type, Path path, JsonCodec codec, T value, boolean overwrite) + private void writeFile(String type, Location location, JsonCodec codec, T value, boolean overwrite) { try { byte[] json = codec.toJsonBytes(value); - if (!overwrite) { - if (metadataFileSystem.exists(path)) { + TrinoOutputFile output = fileSystem.newOutputFile(location); + if (overwrite) { + output.createOrOverwrite(json); + } + else { + // best-effort exclusive and atomic creation + if (fileSystem.newInputFile(location).exists()) { throw new TrinoException(HIVE_METASTORE_ERROR, type + " file already exists"); } - } - - metadataFileSystem.mkdirs(path.getParent()); - - // todo implement safer overwrite code - try (OutputStream outputStream = metadataFileSystem.create(path, overwrite)) { - outputStream.write(json); + try { + output.createExclusive(json); + } + catch (UnsupportedOperationException ignore) { + // fall back to non-exclusive creation, relying on synchronization and above exists check + try (OutputStream out = output.create()) { + out.write(json); + } + } } } catch (Exception e) { throw new TrinoException(HIVE_METASTORE_ERROR, "Could not write " + type, e); } - finally { - listTablesCache.invalidateAll(); - } } - private void renameSchemaFile(SchemaType type, Path oldMetadataDirectory, Path newMetadataDirectory) + private void renameSchemaFile(SchemaType type, Location oldMetadataDirectory, Location newMetadataDirectory) { try { - if (!metadataFileSystem.rename(getSchemaPath(type, oldMetadataDirectory), getSchemaPath(type, newMetadataDirectory))) { - throw new TrinoException(HIVE_METASTORE_ERROR, "Could not rename " + type + " schema"); - } + fileSystem.renameFile(getSchemaFile(type, oldMetadataDirectory), getSchemaFile(type, newMetadataDirectory)); } catch (IOException e) { throw new TrinoException(HIVE_METASTORE_ERROR, "Could not rename " + type + " schema", e); } - finally { - listTablesCache.invalidateAll(); - } } - private void deleteSchemaFile(SchemaType type, Path metadataDirectory) + private void deleteSchemaFile(SchemaType type, Location metadataDirectory) { try { - if (!metadataFileSystem.delete(getSchemaPath(type, metadataDirectory), false)) { - throw new TrinoException(HIVE_METASTORE_ERROR, "Could not delete " + type + " schema"); - } + fileSystem.deleteFile(getSchemaFile(type, metadataDirectory)); } catch (IOException e) { throw new TrinoException(HIVE_METASTORE_ERROR, "Could not delete " + type + " schema", e); } - finally { - listTablesCache.invalidateAll(); - } } - private Path getDatabaseMetadataDirectory(String databaseName) + private Location getDatabaseMetadataDirectory(String databaseName) + { + return catalogDirectory.appendPath(escapeSchemaName(databaseName)); + } + + private Location getFunctionsDirectory(String databaseName) { - return new Path(catalogDirectory, escapeSchemaName(databaseName)); + return getDatabaseMetadataDirectory(databaseName).appendPath(TRINO_FUNCTIONS_DIRECTORY_NAME); } - private Path getTableMetadataDirectory(Table table) + private Location getTableMetadataDirectory(Table table) { return getTableMetadataDirectory(table.getDatabaseName(), table.getTableName()); } - private Path getTableMetadataDirectory(String databaseName, String tableName) + private Location getTableMetadataDirectory(String databaseName, String tableName) { - return new Path(getDatabaseMetadataDirectory(databaseName), escapeTableName(tableName)); + return getDatabaseMetadataDirectory(databaseName).appendPath(escapeTableName(tableName)); } - private Path getPartitionMetadataDirectory(Table table, List values) + private Location getPartitionMetadataDirectory(Table table, List values) { String partitionName = makePartitionName(table.getPartitionColumns(), values); return getPartitionMetadataDirectory(table, partitionName); } - private Path getPartitionMetadataDirectory(Table table, String partitionName) + private Location getPartitionMetadataDirectory(Table table, String partitionName) { - Path tableMetadataDirectory = getTableMetadataDirectory(table); - return new Path(tableMetadataDirectory, partitionName); + return getPartitionMetadataDirectory(table.getDatabaseName(), table.getTableName(), partitionName); } - private Path getPermissionsDirectory(Table table) + private Location getPartitionMetadataDirectory(String databaseName, String tableName, String partitionName) { - return new Path(getTableMetadataDirectory(table), TRINO_PERMISSIONS_DIRECTORY_NAME); + return getTableMetadataDirectory(databaseName, tableName).appendPath(partitionName); } - private static Path getPermissionsPath(Path permissionsDirectory, HivePrincipal grantee) + private Location getPermissionsDirectory(Table table) { - return new Path(permissionsDirectory, grantee.getType().toString().toLowerCase(Locale.US) + "_" + grantee.getName()); + return getTableMetadataDirectory(table).appendPath(TRINO_PERMISSIONS_DIRECTORY_NAME); } - private Path getRolesFile() + private static Location getPermissionsPath(Location permissionsDirectory, HivePrincipal grantee) { - return new Path(catalogDirectory, ROLES_FILE_NAME); + String granteeType = grantee.getType().toString().toLowerCase(Locale.US); + return permissionsDirectory.appendPath(granteeType + "_" + grantee.getName()); } - private Path getRoleGrantsFile() + private Location getRolesFile() { - return new Path(catalogDirectory, ROLE_GRANTS_FILE_NAME); + return catalogDirectory.appendPath(ROLES_FILE_NAME); } - private static Path getSchemaPath(SchemaType type, Path metadataDirectory) + private Location getRoleGrantsFile() { - if (type == DATABASE) { - return new Path( - requireNonNull(metadataDirectory.getParent(), "Can't use root directory as database path"), - format(".%s%s", metadataDirectory.getName(), TRINO_SCHEMA_FILE_NAME_SUFFIX)); - } - return new Path(metadataDirectory, TRINO_SCHEMA_FILE_NAME_SUFFIX); + return catalogDirectory.appendPath(ROLE_GRANTS_FILE_NAME); } - private static boolean isChildDirectory(Path parentDirectory, Path childDirectory) + private static Location getSchemaFile(SchemaType type, Location metadataDirectory) { - if (parentDirectory.equals(childDirectory)) { - return true; - } - if (childDirectory.isRoot()) { - return false; + if (type == DATABASE) { + String path = metadataDirectory.toString(); + if (path.endsWith("/")) { + path = path.substring(0, path.length() - 1); + } + checkArgument(!path.isEmpty(), "Can't use root directory as database path: %s", metadataDirectory); + int index = path.lastIndexOf('/'); + if (index >= 0) { + path = path.substring(0, index + 1) + "." + path.substring(index + 1); + } + else { + path = "." + path; + } + return Location.of(path).appendSuffix(TRINO_SCHEMA_FILE_NAME_SUFFIX); } - return isChildDirectory(parentDirectory, childDirectory.getParent()); + return metadataDirectory.appendPath(TRINO_SCHEMA_FILE_NAME_SUFFIX); } - private static class RoleGranteeTuple + private static PartitionStatistics toHivePartitionStatistics(Map parameters, Map columnStatistics) { - private final String role; - private final HivePrincipal grantee; - - private RoleGranteeTuple(String role, HivePrincipal grantee) - { - this.role = requireNonNull(role, "role is null"); - this.grantee = requireNonNull(grantee, "grantee is null"); - } - - public String getRole() - { - return role; - } - - public HivePrincipal getGrantee() - { - return grantee; - } + HiveBasicStatistics basicStatistics = getHiveBasicStatistics(parameters); + Map hiveColumnStatistics = columnStatistics.entrySet().stream() + .collect(toImmutableMap(Entry::getKey, column -> column.getValue().toHiveColumnStatistics(basicStatistics))); + return new PartitionStatistics(basicStatistics, hiveColumnStatistics); + } - @Override - public boolean equals(Object o) - { - if (this == o) { - return true; - } - if (o == null || getClass() != o.getClass()) { - return false; - } - RoleGranteeTuple that = (RoleGranteeTuple) o; - return Objects.equals(role, that.role) && - Objects.equals(grantee, that.grantee); - } + private static Map fromHiveColumnStats(Map columnStatistics) + { + return columnStatistics.entrySet().stream() + .collect(toImmutableMap(Entry::getKey, entry -> fromHiveColumnStatistics(entry.getValue()))); + } - @Override - public int hashCode() - { - return Objects.hash(role, grantee); - } + private static Map toHiveColumnStats(Set columnNames, Map partitionMetadata, Map columnStatistics) + { + HiveBasicStatistics basicStatistics = getHiveBasicStatistics(partitionMetadata); + return columnStatistics.entrySet().stream() + .filter(entry -> columnNames.contains(entry.getKey())) + .collect(toImmutableMap(Entry::getKey, entry -> entry.getValue().toHiveColumnStatistics(basicStatistics))); + } - @Override - public String toString() + private record RoleGrantee(String role, HivePrincipal grantee) + { + private RoleGrantee { - return toStringHelper(this) - .add("role", role) - .add("grantee", grantee) - .toString(); + requireNonNull(role, "role is null"); + requireNonNull(grantee, "grantee is null"); } } diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/file/FileHiveMetastoreFactory.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/file/FileHiveMetastoreFactory.java index cc7e310a81c5..017860d73d28 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/file/FileHiveMetastoreFactory.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/file/FileHiveMetastoreFactory.java @@ -14,7 +14,7 @@ package io.trino.plugin.hive.metastore.file; import com.google.inject.Inject; -import io.trino.hdfs.HdfsEnvironment; +import io.trino.filesystem.TrinoFileSystemFactory; import io.trino.plugin.hive.HideDeltaLakeTables; import io.trino.plugin.hive.NodeVersion; import io.trino.plugin.hive.metastore.HiveMetastore; @@ -26,13 +26,17 @@ public class FileHiveMetastoreFactory implements HiveMetastoreFactory { - private final FileHiveMetastore metastore; + private final HiveMetastore metastore; @Inject - public FileHiveMetastoreFactory(NodeVersion nodeVersion, HdfsEnvironment hdfsEnvironment, @HideDeltaLakeTables boolean hideDeltaLakeTables, FileHiveMetastoreConfig config) + public FileHiveMetastoreFactory( + NodeVersion nodeVersion, + TrinoFileSystemFactory fileSystemFactory, + @HideDeltaLakeTables boolean hideDeltaLakeTables, + FileHiveMetastoreConfig config) { - // file metastore does not support impersonation, so just create a single shared instance - metastore = new FileHiveMetastore(nodeVersion, hdfsEnvironment, hideDeltaLakeTables, config); + // file metastore does not support impersonation, so create a single shared instance + metastore = new FileHiveMetastore(nodeVersion, fileSystemFactory, hideDeltaLakeTables, config); } @Override diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/file/PartitionMetadata.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/file/PartitionMetadata.java index 7e2ec117e292..c638d893dbe4 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/file/PartitionMetadata.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/file/PartitionMetadata.java @@ -21,7 +21,6 @@ import io.trino.plugin.hive.HiveStorageFormat; import io.trino.plugin.hive.PartitionStatistics; import io.trino.plugin.hive.metastore.Column; -import io.trino.plugin.hive.metastore.HiveColumnStatistics; import io.trino.plugin.hive.metastore.Partition; import io.trino.plugin.hive.metastore.PartitionWithStatistics; import io.trino.plugin.hive.metastore.Storage; @@ -33,8 +32,10 @@ import java.util.Map; import java.util.Optional; +import static com.google.common.collect.ImmutableMap.toImmutableMap; import static io.trino.plugin.hive.TableType.EXTERNAL_TABLE; import static io.trino.plugin.hive.metastore.StorageFormat.VIEW_STORAGE_FORMAT; +import static io.trino.plugin.hive.metastore.file.ColumnStatistics.fromHiveColumnStatistics; import static io.trino.plugin.hive.metastore.thrift.ThriftMetastoreUtil.updateStatisticsParameters; import static java.util.Objects.requireNonNull; @@ -49,7 +50,7 @@ public class PartitionMetadata private final Optional externalLocation; - private final Map columnStatistics; + private final Map columnStatistics; @JsonCreator public PartitionMetadata( @@ -59,7 +60,7 @@ public PartitionMetadata( @JsonProperty("bucketProperty") Optional bucketProperty, @JsonProperty("serdeParameters") Map serdeParameters, @JsonProperty("externalLocation") Optional externalLocation, - @JsonProperty("columnStatistics") Map columnStatistics) + @JsonProperty("columnStatistics") Map columnStatistics) { this.columns = ImmutableList.copyOf(requireNonNull(columns, "columns is null")); this.parameters = ImmutableMap.copyOf(requireNonNull(parameters, "parameters is null")); @@ -94,7 +95,8 @@ public PartitionMetadata(Table table, PartitionWithStatistics partitionWithStati bucketProperty = partition.getStorage().getBucketProperty(); serdeParameters = partition.getStorage().getSerdeParameters(); - columnStatistics = ImmutableMap.copyOf(statistics.getColumnStatistics()); + columnStatistics = statistics.getColumnStatistics().entrySet().stream() + .collect(toImmutableMap(Map.Entry::getKey, entry -> fromHiveColumnStatistics(entry.getValue()))); } @JsonProperty @@ -134,7 +136,7 @@ public Optional getExternalLocation() } @JsonProperty - public Map getColumnStatistics() + public Map getColumnStatistics() { return columnStatistics; } @@ -144,7 +146,7 @@ public PartitionMetadata withParameters(Map parameters) return new PartitionMetadata(columns, parameters, storageFormat, bucketProperty, serdeParameters, externalLocation, columnStatistics); } - public PartitionMetadata withColumnStatistics(Map columnStatistics) + public PartitionMetadata withColumnStatistics(Map columnStatistics) { return new PartitionMetadata(columns, parameters, storageFormat, bucketProperty, serdeParameters, externalLocation, columnStatistics); } diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/file/TableMetadata.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/file/TableMetadata.java index 770b6fd6dea3..ea94e9697e23 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/file/TableMetadata.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/file/TableMetadata.java @@ -19,7 +19,6 @@ import com.google.common.collect.ImmutableMap; import io.trino.plugin.hive.HiveBucketProperty; import io.trino.plugin.hive.HiveStorageFormat; -import io.trino.plugin.hive.metastore.HiveColumnStatistics; import io.trino.plugin.hive.metastore.Storage; import io.trino.plugin.hive.metastore.StorageFormat; import io.trino.plugin.hive.metastore.Table; @@ -55,7 +54,7 @@ public class TableMetadata private final Optional viewOriginalText; private final Optional viewExpandedText; - private final Map columnStatistics; + private final Map columnStatistics; @JsonCreator public TableMetadata( @@ -72,7 +71,7 @@ public TableMetadata( @JsonProperty("externalLocation") Optional externalLocation, @JsonProperty("viewOriginalText") Optional viewOriginalText, @JsonProperty("viewExpandedText") Optional viewExpandedText, - @JsonProperty("columnStatistics") Map columnStatistics) + @JsonProperty("columnStatistics") Map columnStatistics) { this.writerVersion = requireNonNull(writerVersion, "writerVersion is null"); this.owner = requireNonNull(owner, "owner is null"); @@ -227,7 +226,7 @@ public Optional getViewExpandedText() } @JsonProperty - public Map getColumnStatistics() + public Map getColumnStatistics() { return columnStatistics; } @@ -251,6 +250,25 @@ public TableMetadata withDataColumns(String currentVersion, List dataCol columnStatistics); } + public TableMetadata withPartitionColumns(String currentVersion, List partitionColumns) + { + return new TableMetadata( + Optional.of(requireNonNull(currentVersion, "currentVersion is null")), + owner, + tableType, + dataColumns, + partitionColumns, + parameters, + storageFormat, + originalStorageFormat, + bucketProperty, + serdeParameters, + externalLocation, + viewOriginalText, + viewExpandedText, + columnStatistics); + } + public TableMetadata withParameters(String currentVersion, Map parameters) { return new TableMetadata( @@ -270,7 +288,7 @@ public TableMetadata withParameters(String currentVersion, Map p columnStatistics); } - public TableMetadata withColumnStatistics(String currentVersion, Map columnStatistics) + public TableMetadata withColumnStatistics(String currentVersion, Map columnStatistics) { return new TableMetadata( Optional.of(requireNonNull(currentVersion, "currentVersion is null")), diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/aws/AwsApiCallStats.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/AwsApiCallStats.java similarity index 97% rename from plugin/trino-hive/src/main/java/io/trino/plugin/hive/aws/AwsApiCallStats.java rename to plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/AwsApiCallStats.java index ff650da566cb..8b45d5bc4af0 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/aws/AwsApiCallStats.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/AwsApiCallStats.java @@ -11,7 +11,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package io.trino.plugin.hive.aws; +package io.trino.plugin.hive.metastore.glue; import com.google.errorprone.annotations.ThreadSafe; import io.airlift.stats.CounterStat; diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/AwsSdkUtil.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/AwsSdkUtil.java index 646252513c11..0e15916e6e45 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/AwsSdkUtil.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/AwsSdkUtil.java @@ -14,7 +14,6 @@ package io.trino.plugin.hive.metastore.glue; import com.google.common.collect.AbstractIterator; -import io.trino.plugin.hive.aws.AwsApiCallStats; import java.util.Iterator; import java.util.function.BiConsumer; diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/DefaultGlueColumnStatisticsProvider.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/DefaultGlueColumnStatisticsProvider.java deleted file mode 100644 index 92e85f14e748..000000000000 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/DefaultGlueColumnStatisticsProvider.java +++ /dev/null @@ -1,303 +0,0 @@ -/* - * 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.hive.metastore.glue; - -import com.amazonaws.services.glue.AWSGlueAsync; -import com.amazonaws.services.glue.model.ColumnStatistics; -import com.amazonaws.services.glue.model.ColumnStatisticsData; -import com.amazonaws.services.glue.model.ColumnStatisticsType; -import com.amazonaws.services.glue.model.DateColumnStatisticsData; -import com.amazonaws.services.glue.model.DecimalColumnStatisticsData; -import com.amazonaws.services.glue.model.DeleteColumnStatisticsForPartitionRequest; -import com.amazonaws.services.glue.model.DeleteColumnStatisticsForTableRequest; -import com.amazonaws.services.glue.model.DoubleColumnStatisticsData; -import com.amazonaws.services.glue.model.EntityNotFoundException; -import com.amazonaws.services.glue.model.GetColumnStatisticsForPartitionRequest; -import com.amazonaws.services.glue.model.GetColumnStatisticsForPartitionResult; -import com.amazonaws.services.glue.model.GetColumnStatisticsForTableRequest; -import com.amazonaws.services.glue.model.GetColumnStatisticsForTableResult; -import com.amazonaws.services.glue.model.LongColumnStatisticsData; -import com.amazonaws.services.glue.model.UpdateColumnStatisticsForPartitionRequest; -import com.amazonaws.services.glue.model.UpdateColumnStatisticsForTableRequest; -import com.google.common.collect.ImmutableList; -import com.google.common.collect.ImmutableMap; -import com.google.common.collect.Lists; -import io.trino.plugin.hive.HiveBasicStatistics; -import io.trino.plugin.hive.HiveColumnStatisticType; -import io.trino.plugin.hive.metastore.Column; -import io.trino.plugin.hive.metastore.HiveColumnStatistics; -import io.trino.plugin.hive.metastore.Partition; -import io.trino.plugin.hive.metastore.Table; -import io.trino.plugin.hive.metastore.thrift.ThriftMetastoreUtil; -import io.trino.spi.TrinoException; -import io.trino.spi.type.Type; - -import java.util.ArrayList; -import java.util.Collection; -import java.util.HashMap; -import java.util.List; -import java.util.Map; -import java.util.Set; -import java.util.concurrent.CompletableFuture; -import java.util.concurrent.Executor; - -import static com.google.common.collect.ImmutableList.toImmutableList; -import static com.google.common.collect.Sets.difference; -import static io.airlift.concurrent.MoreFutures.getFutureValue; -import static io.trino.plugin.hive.HiveErrorCode.HIVE_METASTORE_ERROR; -import static io.trino.plugin.hive.HiveErrorCode.HIVE_PARTITION_NOT_FOUND; -import static io.trino.plugin.hive.metastore.glue.converter.GlueStatConverter.fromGlueColumnStatistics; -import static io.trino.plugin.hive.metastore.glue.converter.GlueStatConverter.toGlueColumnStatistics; -import static io.trino.plugin.hive.metastore.thrift.ThriftMetastoreUtil.getHiveBasicStatistics; -import static java.util.concurrent.CompletableFuture.allOf; -import static java.util.concurrent.CompletableFuture.runAsync; -import static java.util.concurrent.CompletableFuture.supplyAsync; -import static java.util.stream.Collectors.toUnmodifiableList; - -public class DefaultGlueColumnStatisticsProvider - implements GlueColumnStatisticsProvider -{ - // Read limit for AWS Glue API GetColumnStatisticsForPartition - // https://docs.aws.amazon.com/glue/latest/dg/aws-glue-api-catalog-partitions.html#aws-glue-api-catalog-partitions-GetColumnStatisticsForPartition - private static final int GLUE_COLUMN_READ_STAT_PAGE_SIZE = 100; - - // Write limit for AWS Glue API UpdateColumnStatisticsForPartition - // https://docs.aws.amazon.com/glue/latest/dg/aws-glue-api-catalog-partitions.html#aws-glue-api-catalog-partitions-UpdateColumnStatisticsForPartition - private static final int GLUE_COLUMN_WRITE_STAT_PAGE_SIZE = 25; - - private final GlueMetastoreStats stats; - private final AWSGlueAsync glueClient; - private final Executor readExecutor; - private final Executor writeExecutor; - - public DefaultGlueColumnStatisticsProvider(AWSGlueAsync glueClient, Executor readExecutor, Executor writeExecutor, GlueMetastoreStats stats) - { - this.glueClient = glueClient; - this.readExecutor = readExecutor; - this.writeExecutor = writeExecutor; - this.stats = stats; - } - - @Override - public Set getSupportedColumnStatistics(Type type) - { - return ThriftMetastoreUtil.getSupportedColumnStatistics(type); - } - - @Override - public Map getTableColumnStatistics(Table table) - { - try { - List columnNames = getAllColumns(table); - List> columnChunks = Lists.partition(columnNames, GLUE_COLUMN_READ_STAT_PAGE_SIZE); - List> getStatsFutures = columnChunks.stream() - .map(partialColumns -> supplyAsync(() -> { - GetColumnStatisticsForTableRequest request = new GetColumnStatisticsForTableRequest() - .withDatabaseName(table.getDatabaseName()) - .withTableName(table.getTableName()) - .withColumnNames(partialColumns); - return stats.getGetColumnStatisticsForTable().call(() -> glueClient.getColumnStatisticsForTable(request)); - }, readExecutor)).collect(toImmutableList()); - - HiveBasicStatistics tableStatistics = getHiveBasicStatistics(table.getParameters()); - ImmutableMap.Builder columnStatsMapBuilder = ImmutableMap.builder(); - for (CompletableFuture future : getStatsFutures) { - GetColumnStatisticsForTableResult tableColumnsStats = getFutureValue(future, TrinoException.class); - for (ColumnStatistics columnStatistics : tableColumnsStats.getColumnStatisticsList()) { - columnStatsMapBuilder.put( - columnStatistics.getColumnName(), - fromGlueColumnStatistics(columnStatistics.getStatisticsData(), tableStatistics.getRowCount())); - } - } - return columnStatsMapBuilder.buildOrThrow(); - } - catch (RuntimeException ex) { - throw new TrinoException(HIVE_METASTORE_ERROR, ex); - } - } - - @Override - public Map> getPartitionColumnStatistics(Collection partitions) - { - Map>> resultsForPartition = new HashMap<>(); - for (Partition partition : partitions) { - ImmutableList.Builder> futures = ImmutableList.builder(); - List> columnChunks = Lists.partition(partition.getColumns(), GLUE_COLUMN_READ_STAT_PAGE_SIZE); - for (List partialPartitionColumns : columnChunks) { - List columnsNames = partialPartitionColumns.stream() - .map(Column::getName) - .collect(toImmutableList()); - GetColumnStatisticsForPartitionRequest request = new GetColumnStatisticsForPartitionRequest() - .withDatabaseName(partition.getDatabaseName()) - .withTableName(partition.getTableName()) - .withColumnNames(columnsNames) - .withPartitionValues(partition.getValues()); - futures.add(supplyAsync(() -> stats.getGetColumnStatisticsForPartition().call(() -> glueClient.getColumnStatisticsForPartition(request)), readExecutor)); - } - resultsForPartition.put(partition, futures.build()); - } - - try { - ImmutableMap.Builder> partitionStatistics = ImmutableMap.builder(); - resultsForPartition.forEach((partition, futures) -> { - HiveBasicStatistics tableStatistics = getHiveBasicStatistics(partition.getParameters()); - ImmutableMap.Builder columnStatsMapBuilder = ImmutableMap.builder(); - - for (CompletableFuture getColumnStatisticsResultFuture : futures) { - GetColumnStatisticsForPartitionResult getColumnStatisticsResult = getFutureValue(getColumnStatisticsResultFuture); - getColumnStatisticsResult.getColumnStatisticsList().forEach(columnStatistics -> - columnStatsMapBuilder.put( - columnStatistics.getColumnName(), - fromGlueColumnStatistics(columnStatistics.getStatisticsData(), tableStatistics.getRowCount()))); - } - - partitionStatistics.put(partition, columnStatsMapBuilder.buildOrThrow()); - }); - - return partitionStatistics.buildOrThrow(); - } - catch (RuntimeException ex) { - if (ex.getCause() != null && ex.getCause() instanceof EntityNotFoundException) { - throw new TrinoException(HIVE_PARTITION_NOT_FOUND, ex.getCause()); - } - throw new TrinoException(HIVE_METASTORE_ERROR, ex); - } - } - - // Glue will accept null as min/max values but return 0 when reading - // to avoid incorrect stats we skip writes for column statistics that have min/max null - // this can be removed once glue fix this behaviour - private boolean isGlueWritable(ColumnStatistics stats) - { - ColumnStatisticsData statisticsData = stats.getStatisticsData(); - String columnType = stats.getStatisticsData().getType(); - if (columnType.equals(ColumnStatisticsType.DATE.toString())) { - DateColumnStatisticsData data = statisticsData.getDateColumnStatisticsData(); - return data.getMaximumValue() != null && data.getMinimumValue() != null; - } - if (columnType.equals(ColumnStatisticsType.DECIMAL.toString())) { - DecimalColumnStatisticsData data = statisticsData.getDecimalColumnStatisticsData(); - return data.getMaximumValue() != null && data.getMinimumValue() != null; - } - if (columnType.equals(ColumnStatisticsType.DOUBLE.toString())) { - DoubleColumnStatisticsData data = statisticsData.getDoubleColumnStatisticsData(); - return data.getMaximumValue() != null && data.getMinimumValue() != null; - } - if (columnType.equals(ColumnStatisticsType.LONG.toString())) { - LongColumnStatisticsData data = statisticsData.getLongColumnStatisticsData(); - return data.getMaximumValue() != null && data.getMinimumValue() != null; - } - return true; - } - - @Override - public void updateTableColumnStatistics(Table table, Map updatedTableColumnStatistics) - { - try { - HiveBasicStatistics tableStats = getHiveBasicStatistics(table.getParameters()); - List columnStats = toGlueColumnStatistics(table, updatedTableColumnStatistics, tableStats.getRowCount()).stream() - .filter(this::isGlueWritable) - .collect(toUnmodifiableList()); - - List> columnChunks = Lists.partition(columnStats, GLUE_COLUMN_WRITE_STAT_PAGE_SIZE); - - List> updateFutures = columnChunks.stream().map(columnChunk -> runAsync( - () -> stats.getUpdateColumnStatisticsForTable().call(() -> glueClient.updateColumnStatisticsForTable( - new UpdateColumnStatisticsForTableRequest() - .withDatabaseName(table.getDatabaseName()) - .withTableName(table.getTableName()) - .withColumnStatisticsList(columnChunk))), this.writeExecutor)) - .collect(toUnmodifiableList()); - - Map currentTableColumnStatistics = this.getTableColumnStatistics(table); - Set removedStatistics = difference(currentTableColumnStatistics.keySet(), updatedTableColumnStatistics.keySet()); - List> deleteFutures = removedStatistics.stream() - .map(column -> runAsync(() -> stats.getDeleteColumnStatisticsForTable().call(() -> - glueClient.deleteColumnStatisticsForTable( - new DeleteColumnStatisticsForTableRequest() - .withDatabaseName(table.getDatabaseName()) - .withTableName(table.getTableName()) - .withColumnName(column))), this.writeExecutor)) - .collect(toUnmodifiableList()); - - ImmutableList> updateOperationsFutures = ImmutableList.>builder() - .addAll(updateFutures) - .addAll(deleteFutures) - .build(); - - getFutureValue(allOf(updateOperationsFutures.toArray(CompletableFuture[]::new))); - } - catch (RuntimeException ex) { - throw new TrinoException(HIVE_METASTORE_ERROR, ex); - } - } - - @Override - public void updatePartitionStatistics(Set partitionStatisticsUpdates) - { - Map> currentStatistics = getPartitionColumnStatistics( - partitionStatisticsUpdates.stream() - .map(PartitionStatisticsUpdate::getPartition).collect(toImmutableList())); - - List> updateFutures = new ArrayList<>(); - for (PartitionStatisticsUpdate update : partitionStatisticsUpdates) { - Partition partition = update.getPartition(); - Map updatedColumnStatistics = update.getColumnStatistics(); - - HiveBasicStatistics partitionStats = getHiveBasicStatistics(partition.getParameters()); - List columnStats = toGlueColumnStatistics(partition, updatedColumnStatistics, partitionStats.getRowCount()).stream() - .filter(this::isGlueWritable) - .collect(toUnmodifiableList()); - - List> columnChunks = Lists.partition(columnStats, GLUE_COLUMN_WRITE_STAT_PAGE_SIZE); - columnChunks.forEach(columnChunk -> - updateFutures.add(runAsync(() -> stats.getUpdateColumnStatisticsForPartition().call(() -> - glueClient.updateColumnStatisticsForPartition( - new UpdateColumnStatisticsForPartitionRequest() - .withDatabaseName(partition.getDatabaseName()) - .withTableName(partition.getTableName()) - .withPartitionValues(partition.getValues()) - .withColumnStatisticsList(columnChunk))), - writeExecutor))); - - Set removedStatistics = difference(currentStatistics.get(partition).keySet(), updatedColumnStatistics.keySet()); - removedStatistics.forEach(column -> - updateFutures.add(runAsync(() -> stats.getDeleteColumnStatisticsForPartition().call(() -> - glueClient.deleteColumnStatisticsForPartition( - new DeleteColumnStatisticsForPartitionRequest() - .withDatabaseName(partition.getDatabaseName()) - .withTableName(partition.getTableName()) - .withPartitionValues(partition.getValues()) - .withColumnName(column))), - writeExecutor))); - } - try { - getFutureValue(allOf(updateFutures.toArray(CompletableFuture[]::new))); - } - catch (RuntimeException ex) { - if (ex.getCause() != null && ex.getCause() instanceof EntityNotFoundException) { - throw new TrinoException(HIVE_PARTITION_NOT_FOUND, ex.getCause()); - } - throw new TrinoException(HIVE_METASTORE_ERROR, ex); - } - } - - private List getAllColumns(Table table) - { - ImmutableList.Builder allColumns = ImmutableList.builderWithExpectedSize(table.getDataColumns().size() + table.getPartitionColumns().size()); - table.getDataColumns().stream().map(Column::getName).forEach(allColumns::add); - table.getPartitionColumns().stream().map(Column::getName).forEach(allColumns::add); - return allColumns.build(); - } -} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/DefaultGlueColumnStatisticsProviderFactory.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/DefaultGlueColumnStatisticsProviderFactory.java deleted file mode 100644 index c08a3f545aa0..000000000000 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/DefaultGlueColumnStatisticsProviderFactory.java +++ /dev/null @@ -1,47 +0,0 @@ -/* - * 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.hive.metastore.glue; - -import com.amazonaws.services.glue.AWSGlueAsync; -import com.google.inject.Inject; - -import java.util.concurrent.Executor; - -import static java.util.Objects.requireNonNull; - -public class DefaultGlueColumnStatisticsProviderFactory - implements GlueColumnStatisticsProviderFactory -{ - private final Executor statisticsReadExecutor; - private final Executor statisticsWriteExecutor; - - @Inject - public DefaultGlueColumnStatisticsProviderFactory( - @ForGlueColumnStatisticsRead Executor statisticsReadExecutor, - @ForGlueColumnStatisticsWrite Executor statisticsWriteExecutor) - { - this.statisticsReadExecutor = requireNonNull(statisticsReadExecutor, "statisticsReadExecutor is null"); - this.statisticsWriteExecutor = requireNonNull(statisticsWriteExecutor, "statisticsWriteExecutor is null"); - } - - @Override - public GlueColumnStatisticsProvider createGlueColumnStatisticsProvider(AWSGlueAsync glueClient, GlueMetastoreStats stats) - { - return new DefaultGlueColumnStatisticsProvider( - glueClient, - statisticsReadExecutor, - statisticsWriteExecutor, - stats); - } -} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/DefaultGlueMetastoreTableFilterProvider.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/DefaultGlueMetastoreTableFilterProvider.java deleted file mode 100644 index 8e463704a527..000000000000 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/DefaultGlueMetastoreTableFilterProvider.java +++ /dev/null @@ -1,46 +0,0 @@ -/* - * 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.hive.metastore.glue; - -import com.amazonaws.services.glue.model.Table; -import com.google.inject.Inject; -import com.google.inject.Provider; -import io.trino.plugin.hive.HideDeltaLakeTables; - -import java.util.function.Predicate; - -import static io.trino.plugin.hive.metastore.glue.converter.GlueToTrinoConverter.getTableParameters; -import static io.trino.plugin.hive.util.HiveUtil.isDeltaLakeTable; -import static java.util.function.Predicate.not; - -public class DefaultGlueMetastoreTableFilterProvider - implements Provider> -{ - private final boolean hideDeltaLakeTables; - - @Inject - public DefaultGlueMetastoreTableFilterProvider(@HideDeltaLakeTables boolean hideDeltaLakeTables) - { - this.hideDeltaLakeTables = hideDeltaLakeTables; - } - - @Override - public Predicate
get() - { - if (hideDeltaLakeTables) { - return not(table -> isDeltaLakeTable(getTableParameters(table))); - } - return table -> true; - } -} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/DisabledGlueColumnStatisticsProvider.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/DisabledGlueColumnStatisticsProvider.java deleted file mode 100644 index a0b4024fcb3a..000000000000 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/DisabledGlueColumnStatisticsProvider.java +++ /dev/null @@ -1,69 +0,0 @@ -/* - * 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.hive.metastore.glue; - -import com.google.common.collect.ImmutableMap; -import com.google.common.collect.ImmutableSet; -import io.trino.plugin.hive.HiveColumnStatisticType; -import io.trino.plugin.hive.metastore.HiveColumnStatistics; -import io.trino.plugin.hive.metastore.Partition; -import io.trino.plugin.hive.metastore.Table; -import io.trino.spi.TrinoException; -import io.trino.spi.type.Type; - -import java.util.Collection; -import java.util.Map; -import java.util.Set; - -import static com.google.common.collect.ImmutableMap.toImmutableMap; -import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED; -import static java.util.function.UnaryOperator.identity; - -public class DisabledGlueColumnStatisticsProvider - implements GlueColumnStatisticsProvider -{ - @Override - public Set getSupportedColumnStatistics(Type type) - { - return ImmutableSet.of(); - } - - @Override - public Map getTableColumnStatistics(Table table) - { - return ImmutableMap.of(); - } - - @Override - public Map> getPartitionColumnStatistics(Collection partitions) - { - return partitions.stream().collect(toImmutableMap(identity(), partition -> ImmutableMap.of())); - } - - @Override - public void updateTableColumnStatistics(Table table, Map columnStatistics) - { - if (!columnStatistics.isEmpty()) { - throw new TrinoException(NOT_SUPPORTED, "Glue metastore column level statistics are disabled"); - } - } - - @Override - public void updatePartitionStatistics(Set partitionStatisticsUpdates) - { - if (partitionStatisticsUpdates.stream().anyMatch(update -> !update.getColumnStatistics().isEmpty())) { - throw new TrinoException(NOT_SUPPORTED, "Glue metastore column level statistics are disabled"); - } - } -} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/GlueCache.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/GlueCache.java new file mode 100644 index 000000000000..9c8f01b9905c --- /dev/null +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/GlueCache.java @@ -0,0 +1,79 @@ +/* + * 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.hive.metastore.glue; + +import io.trino.plugin.hive.metastore.Database; +import io.trino.plugin.hive.metastore.HiveColumnStatistics; +import io.trino.plugin.hive.metastore.Partition; +import io.trino.plugin.hive.metastore.Table; +import io.trino.plugin.hive.metastore.TableInfo; + +import java.util.Collection; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.function.BiFunction; +import java.util.function.Consumer; +import java.util.function.Function; +import java.util.function.Supplier; + +public interface GlueCache +{ + GlueCache NOOP = new NoopGlueCache(); + + List getDatabaseNames(Function, List> loader); + + /** + * Invalidate the database cache and cascade to all nested elements in the database (table, partition, function, etc.). + */ + void invalidateDatabase(String databaseName); + + void invalidateDatabaseNames(); + + Optional getDatabase(String databaseName, Supplier> loader); + + List getTables(String databaseName, Function, List> loader); + + void invalidateTables(String databaseName); + + Optional
getTable(String databaseName, String tableName, Supplier> loader); + + void invalidateTable(String databaseName, String tableName, boolean cascade); + + Map getTableColumnStatistics(String databaseName, String tableName, Set columnNames, Function, Map> loader); + + void invalidateTableColumnStatistics(String databaseName, String tableName); + + Set getPartitionNames(String databaseName, String tableName, String glueExpression, Function, Set> loader); + + Optional getPartition(String databaseName, String tableName, PartitionName partitionName, Supplier> loader); + + Collection batchGetPartitions( + String databaseName, + String tableName, + Collection partitionNames, + BiFunction, Collection, Collection> loader); + + void invalidatePartition(String databaseName, String tableName, PartitionName partitionName); + + Map getPartitionColumnStatistics( + String databaseName, + String tableName, + PartitionName partitionName, + Set columnNames, + Function, Map> loader); + + void flushCache(); +} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/GlueClientUtil.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/GlueClientUtil.java deleted file mode 100644 index dbfdcf2067ca..000000000000 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/GlueClientUtil.java +++ /dev/null @@ -1,65 +0,0 @@ -/* - * 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.hive.metastore.glue; - -import com.amazonaws.ClientConfiguration; -import com.amazonaws.auth.AWSCredentialsProvider; -import com.amazonaws.client.builder.AwsClientBuilder; -import com.amazonaws.handlers.RequestHandler2; -import com.amazonaws.metrics.RequestMetricCollector; -import com.amazonaws.services.glue.AWSGlueAsync; -import com.amazonaws.services.glue.AWSGlueAsyncClientBuilder; - -import java.util.Set; - -import static com.google.common.base.Preconditions.checkArgument; -import static io.trino.hdfs.s3.AwsCurrentRegionHolder.getCurrentRegionFromEC2Metadata; - -public final class GlueClientUtil -{ - private GlueClientUtil() {} - - public static AWSGlueAsync createAsyncGlueClient( - GlueHiveMetastoreConfig config, - AWSCredentialsProvider credentialsProvider, - Set requestHandlers, - RequestMetricCollector metricsCollector) - { - ClientConfiguration clientConfig = new ClientConfiguration() - .withMaxConnections(config.getMaxGlueConnections()) - .withMaxErrorRetry(config.getMaxGlueErrorRetries()); - AWSGlueAsyncClientBuilder asyncGlueClientBuilder = AWSGlueAsyncClientBuilder.standard() - .withMetricsCollector(metricsCollector) - .withClientConfiguration(clientConfig); - - asyncGlueClientBuilder.setRequestHandlers(requestHandlers.toArray(RequestHandler2[]::new)); - - if (config.getGlueEndpointUrl().isPresent()) { - checkArgument(config.getGlueRegion().isPresent(), "Glue region must be set when Glue endpoint URL is set"); - asyncGlueClientBuilder.setEndpointConfiguration(new AwsClientBuilder.EndpointConfiguration( - config.getGlueEndpointUrl().get(), - config.getGlueRegion().get())); - } - else if (config.getGlueRegion().isPresent()) { - asyncGlueClientBuilder.setRegion(config.getGlueRegion().get()); - } - else if (config.getPinGlueClientToCurrentRegion()) { - asyncGlueClientBuilder.setRegion(getCurrentRegionFromEC2Metadata().getName()); - } - - asyncGlueClientBuilder.setCredentials(credentialsProvider); - - return asyncGlueClientBuilder.build(); - } -} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/GlueColumnStatisticsProvider.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/GlueColumnStatisticsProvider.java index 3ab0bee62322..deddbf14b4aa 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/GlueColumnStatisticsProvider.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/GlueColumnStatisticsProvider.java @@ -33,8 +33,16 @@ public interface GlueColumnStatisticsProvider Map getTableColumnStatistics(Table table); + Map getTableColumnStatistics(String databaseName, String tableName, Set columnNames); + Map> getPartitionColumnStatistics(Collection partitions); + Map> getPartitionColumnStatistics( + String databaseName, + String tableName, + Set partitionNames, + Set columns); + default Map getPartitionColumnStatistics(Partition partition) { return getPartitionColumnStatistics(ImmutableSet.of(partition)).get(partition); diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/GlueContext.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/GlueContext.java new file mode 100644 index 000000000000..35103b9ba10a --- /dev/null +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/GlueContext.java @@ -0,0 +1,164 @@ +/* + * 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.hive.metastore.glue; + +import com.google.errorprone.annotations.CanIgnoreReturnValue; +import com.google.inject.Inject; +import software.amazon.awssdk.awscore.AwsRequest; +import software.amazon.awssdk.services.glue.model.BatchCreatePartitionRequest; +import software.amazon.awssdk.services.glue.model.BatchGetPartitionRequest; +import software.amazon.awssdk.services.glue.model.BatchUpdatePartitionRequest; +import software.amazon.awssdk.services.glue.model.CreateDatabaseRequest; +import software.amazon.awssdk.services.glue.model.CreateTableRequest; +import software.amazon.awssdk.services.glue.model.CreateUserDefinedFunctionRequest; +import software.amazon.awssdk.services.glue.model.DeleteColumnStatisticsForPartitionRequest; +import software.amazon.awssdk.services.glue.model.DeleteColumnStatisticsForTableRequest; +import software.amazon.awssdk.services.glue.model.DeleteDatabaseRequest; +import software.amazon.awssdk.services.glue.model.DeletePartitionRequest; +import software.amazon.awssdk.services.glue.model.DeleteTableRequest; +import software.amazon.awssdk.services.glue.model.DeleteUserDefinedFunctionRequest; +import software.amazon.awssdk.services.glue.model.GetColumnStatisticsForPartitionRequest; +import software.amazon.awssdk.services.glue.model.GetColumnStatisticsForTableRequest; +import software.amazon.awssdk.services.glue.model.GetDatabaseRequest; +import software.amazon.awssdk.services.glue.model.GetDatabasesRequest; +import software.amazon.awssdk.services.glue.model.GetPartitionRequest; +import software.amazon.awssdk.services.glue.model.GetPartitionsRequest; +import software.amazon.awssdk.services.glue.model.GetTableRequest; +import software.amazon.awssdk.services.glue.model.GetTablesRequest; +import software.amazon.awssdk.services.glue.model.GetUserDefinedFunctionRequest; +import software.amazon.awssdk.services.glue.model.GetUserDefinedFunctionsRequest; +import software.amazon.awssdk.services.glue.model.UpdateColumnStatisticsForPartitionRequest; +import software.amazon.awssdk.services.glue.model.UpdateColumnStatisticsForTableRequest; +import software.amazon.awssdk.services.glue.model.UpdateDatabaseRequest; +import software.amazon.awssdk.services.glue.model.UpdatePartitionRequest; +import software.amazon.awssdk.services.glue.model.UpdateTableRequest; +import software.amazon.awssdk.services.glue.model.UpdateUserDefinedFunctionRequest; + +import java.util.Optional; + +import static java.util.Objects.requireNonNull; + +public class GlueContext +{ + private final Optional catalogId; + + @Inject + public GlueContext(GlueHiveMetastoreConfig config) + { + this(config.getCatalogId()); + } + + public GlueContext(Optional catalogId) + { + this.catalogId = requireNonNull(catalogId, "catalogId is null"); + } + + @CanIgnoreReturnValue + public B configureClient(B baseRequestBuilder) + { + catalogId.ifPresent(id -> setCatalogId(baseRequestBuilder, id)); + return baseRequestBuilder; + } + + private static void setCatalogId(AwsRequest.Builder request, String catalogId) + { + if (request instanceof GetDatabasesRequest.Builder builder) { + builder.catalogId(catalogId); + } + else if (request instanceof GetDatabaseRequest.Builder builder) { + builder.catalogId(catalogId); + } + else if (request instanceof CreateDatabaseRequest.Builder builder) { + builder.catalogId(catalogId); + } + else if (request instanceof UpdateDatabaseRequest.Builder builder) { + builder.catalogId(catalogId); + } + else if (request instanceof DeleteDatabaseRequest.Builder builder) { + builder.catalogId(catalogId); + } + else if (request instanceof GetTablesRequest.Builder builder) { + builder.catalogId(catalogId); + } + else if (request instanceof GetTableRequest.Builder builder) { + builder.catalogId(catalogId); + } + else if (request instanceof CreateTableRequest.Builder builder) { + builder.catalogId(catalogId); + } + else if (request instanceof UpdateTableRequest.Builder builder) { + builder.catalogId(catalogId); + } + else if (request instanceof DeleteTableRequest.Builder builder) { + builder.catalogId(catalogId); + } + else if (request instanceof GetPartitionsRequest.Builder builder) { + builder.catalogId(catalogId); + } + else if (request instanceof GetPartitionRequest.Builder builder) { + builder.catalogId(catalogId); + } + else if (request instanceof UpdatePartitionRequest.Builder builder) { + builder.catalogId(catalogId); + } + else if (request instanceof DeletePartitionRequest.Builder builder) { + builder.catalogId(catalogId); + } + else if (request instanceof BatchGetPartitionRequest.Builder builder) { + builder.catalogId(catalogId); + } + else if (request instanceof BatchCreatePartitionRequest.Builder builder) { + builder.catalogId(catalogId); + } + else if (request instanceof BatchUpdatePartitionRequest.Builder builder) { + builder.catalogId(catalogId); + } + else if (request instanceof GetColumnStatisticsForTableRequest.Builder builder) { + builder.catalogId(catalogId); + } + else if (request instanceof UpdateColumnStatisticsForTableRequest.Builder builder) { + builder.catalogId(catalogId); + } + else if (request instanceof DeleteColumnStatisticsForTableRequest.Builder builder) { + builder.catalogId(catalogId); + } + else if (request instanceof GetColumnStatisticsForPartitionRequest.Builder builder) { + builder.catalogId(catalogId); + } + else if (request instanceof UpdateColumnStatisticsForPartitionRequest.Builder builder) { + builder.catalogId(catalogId); + } + else if (request instanceof DeleteColumnStatisticsForPartitionRequest.Builder builder) { + builder.catalogId(catalogId); + } + else if (request instanceof GetUserDefinedFunctionsRequest.Builder builder) { + builder.catalogId(catalogId); + } + else if (request instanceof GetUserDefinedFunctionRequest.Builder builder) { + builder.catalogId(catalogId); + } + else if (request instanceof CreateUserDefinedFunctionRequest.Builder builder) { + builder.catalogId(catalogId); + } + else if (request instanceof UpdateUserDefinedFunctionRequest.Builder builder) { + builder.catalogId(catalogId); + } + else if (request instanceof DeleteUserDefinedFunctionRequest.Builder builder) { + builder.catalogId(catalogId); + } + else { + throw new IllegalArgumentException("Unsupported request: " + request); + } + } +} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/GlueConverter.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/GlueConverter.java new file mode 100644 index 000000000000..1de4b4d3664e --- /dev/null +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/GlueConverter.java @@ -0,0 +1,636 @@ +/* + * 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.hive.metastore.glue; + +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import io.trino.plugin.hive.HiveBucketProperty; +import io.trino.plugin.hive.HiveStorageFormat; +import io.trino.plugin.hive.HiveType; +import io.trino.plugin.hive.metastore.BooleanStatistics; +import io.trino.plugin.hive.metastore.Column; +import io.trino.plugin.hive.metastore.Database; +import io.trino.plugin.hive.metastore.DateStatistics; +import io.trino.plugin.hive.metastore.DecimalStatistics; +import io.trino.plugin.hive.metastore.DoubleStatistics; +import io.trino.plugin.hive.metastore.HiveColumnStatistics; +import io.trino.plugin.hive.metastore.IntegerStatistics; +import io.trino.plugin.hive.metastore.Partition; +import io.trino.plugin.hive.metastore.SortingColumn; +import io.trino.plugin.hive.metastore.Storage; +import io.trino.plugin.hive.metastore.StorageFormat; +import io.trino.plugin.hive.metastore.Table; +import io.trino.plugin.hive.type.PrimitiveCategory; +import io.trino.plugin.hive.type.PrimitiveTypeInfo; +import io.trino.plugin.hive.type.TypeInfo; +import io.trino.spi.TrinoException; +import io.trino.spi.security.PrincipalType; +import jakarta.annotation.Nullable; +import org.gaul.modernizer_maven_annotations.SuppressModernizer; +import software.amazon.awssdk.core.SdkBytes; +import software.amazon.awssdk.services.glue.model.BinaryColumnStatisticsData; +import software.amazon.awssdk.services.glue.model.BooleanColumnStatisticsData; +import software.amazon.awssdk.services.glue.model.ColumnStatisticsData; +import software.amazon.awssdk.services.glue.model.ColumnStatisticsType; +import software.amazon.awssdk.services.glue.model.DatabaseInput; +import software.amazon.awssdk.services.glue.model.DateColumnStatisticsData; +import software.amazon.awssdk.services.glue.model.DecimalColumnStatisticsData; +import software.amazon.awssdk.services.glue.model.DecimalNumber; +import software.amazon.awssdk.services.glue.model.DoubleColumnStatisticsData; +import software.amazon.awssdk.services.glue.model.LongColumnStatisticsData; +import software.amazon.awssdk.services.glue.model.Order; +import software.amazon.awssdk.services.glue.model.PartitionInput; +import software.amazon.awssdk.services.glue.model.SerDeInfo; +import software.amazon.awssdk.services.glue.model.StorageDescriptor; +import software.amazon.awssdk.services.glue.model.StringColumnStatisticsData; +import software.amazon.awssdk.services.glue.model.TableInput; + +import java.math.BigDecimal; +import java.math.BigInteger; +import java.time.Instant; +import java.time.LocalDate; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Locale; +import java.util.Map; +import java.util.Optional; +import java.util.OptionalDouble; +import java.util.OptionalLong; +import java.util.concurrent.TimeUnit; + +import static com.google.common.base.MoreObjects.firstNonNull; +import static com.google.common.base.Strings.emptyToNull; +import static com.google.common.base.Strings.lenientFormat; +import static com.google.common.collect.ImmutableList.toImmutableList; +import static com.google.common.collect.ImmutableMap.toImmutableMap; +import static io.trino.plugin.hive.HiveErrorCode.HIVE_INVALID_METADATA; +import static io.trino.plugin.hive.HiveErrorCode.HIVE_UNSUPPORTED_FORMAT; +import static io.trino.plugin.hive.TableType.EXTERNAL_TABLE; +import static io.trino.plugin.hive.ViewReaderUtil.isTrinoMaterializedView; +import static io.trino.plugin.hive.ViewReaderUtil.isTrinoView; +import static io.trino.plugin.hive.metastore.HiveColumnStatistics.createBinaryColumnStatistics; +import static io.trino.plugin.hive.metastore.HiveColumnStatistics.createBooleanColumnStatistics; +import static io.trino.plugin.hive.metastore.HiveColumnStatistics.createDateColumnStatistics; +import static io.trino.plugin.hive.metastore.HiveColumnStatistics.createDecimalColumnStatistics; +import static io.trino.plugin.hive.metastore.HiveColumnStatistics.createDoubleColumnStatistics; +import static io.trino.plugin.hive.metastore.HiveColumnStatistics.createIntegerColumnStatistics; +import static io.trino.plugin.hive.metastore.HiveColumnStatistics.createStringColumnStatistics; +import static io.trino.plugin.hive.metastore.SortingColumn.Order.ASCENDING; +import static io.trino.plugin.hive.metastore.SortingColumn.Order.DESCENDING; +import static io.trino.plugin.hive.metastore.Table.TABLE_COMMENT; +import static io.trino.plugin.hive.metastore.thrift.ThriftMetastoreUtil.fromMetastoreNullsCount; +import static io.trino.plugin.hive.util.HiveUtil.isDeltaLakeTable; +import static io.trino.plugin.hive.util.HiveUtil.isIcebergTable; +import static java.util.Objects.requireNonNull; + +public final class GlueConverter +{ + static final String PUBLIC_OWNER = "PUBLIC"; + private static final Storage FAKE_PARQUET_STORAGE = new Storage( + StorageFormat.create( + "org.apache.hadoop.hive.ql.io.parquet.serde.ParquetHiveSerDe", + "org.apache.hadoop.hive.ql.io.parquet.MapredParquetInputFormat", + "org.apache.hadoop.hive.ql.io.parquet.MapredParquetOutputFormat"), + Optional.empty(), + Optional.empty(), + false, + ImmutableMap.of()); + private static final Column FAKE_COLUMN = new Column("ignored", HiveType.HIVE_INT, Optional.empty(), ImmutableMap.of()); + private static final long SECONDS_PER_DAY = TimeUnit.DAYS.toSeconds(1); + + private GlueConverter() {} + + public static String getTableType(software.amazon.awssdk.services.glue.model.Table glueTable) + { + // Athena treats a missing table type as EXTERNAL_TABLE. + return firstNonNull(getTableTypeNullable(glueTable), EXTERNAL_TABLE.name()); + } + + @Nullable + @SuppressModernizer // Usage of `Table.tableType` is not allowed. Only this method can call that. + public static String getTableTypeNullable(software.amazon.awssdk.services.glue.model.Table glueTable) + { + return glueTable.tableType(); + } + + public static Database fromGlueDatabase(software.amazon.awssdk.services.glue.model.Database glueDb) + { + return new Database( + glueDb.name(), + Optional.ofNullable(emptyToNull(glueDb.locationUri())), + Optional.of(PUBLIC_OWNER), + Optional.of(PrincipalType.ROLE), + Optional.ofNullable(glueDb.description()), + glueDb.parameters()); + } + + public static DatabaseInput toGlueDatabaseInput(Database database) + { + return DatabaseInput.builder() + .name(database.getDatabaseName()) + .parameters(database.getParameters()) + .description(database.getComment().orElse(null)) + .locationUri(database.getLocation().orElse(null)) + .build(); + } + + public static Table fromGlueTable(software.amazon.awssdk.services.glue.model.Table glueTable, String databaseName) + { + String tableType = getTableType(glueTable); + + Map tableParameters = glueTable.parameters(); + if (glueTable.description() != null) { + // Glue description overrides the comment field in the parameters + tableParameters = new LinkedHashMap<>(tableParameters); + tableParameters.put(TABLE_COMMENT, glueTable.description()); + } + + Storage storage; + List partitionColumns; + List dataColumns; + StorageDescriptor sd = glueTable.storageDescriptor(); + if (sd == null) { + if (!isIcebergTable(tableParameters) && !isDeltaLakeTable(tableParameters) && !isTrinoMaterializedView(tableType, tableParameters)) { + throw new TrinoException(HIVE_UNSUPPORTED_FORMAT, "Table StorageDescriptor is null for table '%s' %s".formatted(databaseName, glueTable.name())); + } + + dataColumns = ImmutableList.of(FAKE_COLUMN); + partitionColumns = ImmutableList.of(); + storage = FAKE_PARQUET_STORAGE; + } + else if (isIcebergTable(tableParameters)) { + // todo: any reason to not do this for trino mv? + if (sd.columns() == null) { + dataColumns = ImmutableList.of(FAKE_COLUMN); + } + else { + dataColumns = fromGlueColumns(sd.columns(), ColumnType.DATA, false); + } + partitionColumns = ImmutableList.of(); + storage = FAKE_PARQUET_STORAGE; + } + else if (isDeltaLakeTable(tableParameters)) { + dataColumns = ImmutableList.of(FAKE_COLUMN); + partitionColumns = ImmutableList.of(); + storage = fromGlueStorage(sd, databaseName + "." + glueTable.name()); + } + else { + boolean isCsv = sd.serdeInfo() != null && HiveStorageFormat.CSV.getSerde().equals(sd.serdeInfo().serializationLibrary()); + dataColumns = fromGlueColumns(sd.columns(), ColumnType.DATA, isCsv); + if (glueTable.partitionKeys() != null) { + partitionColumns = fromGlueColumns(glueTable.partitionKeys(), ColumnType.PARTITION, isCsv); + } + else { + partitionColumns = ImmutableList.of(); + } + storage = fromGlueStorage(sd, databaseName + "." + glueTable.name()); + } + + return new Table( + databaseName, + glueTable.name(), + Optional.ofNullable(glueTable.owner()), + tableType, + storage, + dataColumns, + partitionColumns, + tableParameters, + Optional.ofNullable(glueTable.viewOriginalText()), + Optional.ofNullable(glueTable.viewExpandedText()), + OptionalLong.empty()); + } + + public static TableInput toGlueTableInput(Table table) + { + Map tableParameters = table.getParameters(); + Optional comment = Optional.empty(); + if (!isTrinoView(table.getTableType(), table.getParameters()) && !isTrinoMaterializedView(table.getTableType(), table.getParameters())) { + comment = Optional.ofNullable(tableParameters.get(TABLE_COMMENT)); + tableParameters = tableParameters.entrySet().stream() + .filter(entry -> !entry.getKey().equals(TABLE_COMMENT)) + .collect(toImmutableMap(Map.Entry::getKey, Map.Entry::getValue)); + } + + return TableInput.builder() + .name(table.getTableName()) + .owner(table.getOwner().orElse(null)) + .tableType(table.getTableType()) + .storageDescriptor(toGlueStorage(table.getStorage(), table.getDataColumns())) + .partitionKeys(table.getPartitionColumns().stream().map(GlueConverter::toGlueColumn).collect(toImmutableList())) + .parameters(tableParameters) + .viewOriginalText(table.getViewOriginalText().orElse(null)) + .viewExpandedText(table.getViewExpandedText().orElse(null)) + .description(comment.orElse(null)) + .build(); + } + + public static Partition fromGluePartition(String databaseName, String tableName, software.amazon.awssdk.services.glue.model.Partition gluePartition) + { + requireNonNull(gluePartition.storageDescriptor(), "Partition StorageDescriptor is null"); + + if (!databaseName.equals(gluePartition.databaseName())) { + throw new IllegalArgumentException("Unexpected databaseName, expected: %s, but found: %s".formatted(databaseName, gluePartition.databaseName())); + } + if (!tableName.equals(gluePartition.tableName())) { + throw new IllegalArgumentException("Unexpected tableName, expected: %s, but found: %s".formatted(tableName, gluePartition.tableName())); + } + + StorageDescriptor sd = gluePartition.storageDescriptor(); + boolean isCsv = sd.serdeInfo() != null && HiveStorageFormat.CSV.getSerde().equals(sd.serdeInfo().serializationLibrary()); + List partitionName = gluePartition.values(); + return new Partition( + databaseName, + tableName, + partitionName, + fromGlueStorage(sd, databaseName + "." + tableName + "/" + partitionName), + fromGlueColumns(sd.columns(), ColumnType.DATA, isCsv), + gluePartition.parameters()); + } + + public static PartitionInput toGluePartitionInput(Partition partition) + { + return PartitionInput.builder() + .values(partition.getValues()) + .storageDescriptor(toGlueStorage(partition.getStorage(), partition.getColumns())) + .parameters(partition.getParameters()) + .build(); + } + + private static List fromGlueColumns(List glueColumns, ColumnType columnType, boolean isCsv) + { + return glueColumns.stream() + .map(glueColumn -> fromGlueColumn(glueColumn, columnType, isCsv)) + .collect(toImmutableList()); + } + + private static Column fromGlueColumn(software.amazon.awssdk.services.glue.model.Column glueColumn, ColumnType columnType, boolean isCsv) + { + // OpenCSVSerde deserializes columns from csv file into strings, so we set the column type from the metastore + // to string to avoid cast exceptions. + if (columnType == ColumnType.DATA && isCsv) { + //TODO(https://github.com/trinodb/trino/issues/7240) Add tests + return new Column(glueColumn.name(), HiveType.HIVE_STRING, Optional.ofNullable(glueColumn.comment()), glueColumn.parameters()); + } + return new Column(glueColumn.name(), HiveType.valueOf(glueColumn.type().toLowerCase(Locale.ROOT)), Optional.ofNullable(glueColumn.comment()), glueColumn.parameters()); + } + + private static software.amazon.awssdk.services.glue.model.Column toGlueColumn(Column trinoColumn) + { + return software.amazon.awssdk.services.glue.model.Column.builder() + .name(trinoColumn.getName()) + .type(trinoColumn.getType().toString()) + .comment(trinoColumn.getComment().orElse(null)) + .parameters(trinoColumn.getProperties()) + .build(); + } + + private static Storage fromGlueStorage(StorageDescriptor sd, String tablePartitionName) + { + Optional bucketProperty = Optional.empty(); + if (sd.numberOfBuckets() > 0) { + if (sd.bucketColumns().isEmpty()) { + throw new TrinoException(HIVE_INVALID_METADATA, "Table/partition metadata has 'numBuckets' set, but 'bucketCols' is not set: " + tablePartitionName); + } + + List sortBy = ImmutableList.of(); + if (!sd.sortColumns().isEmpty()) { + sortBy = sd.sortColumns().stream() + .map(order -> new SortingColumn(order.column(), fromGlueSortOrder(order.sortOrder(), tablePartitionName))) + .collect(toImmutableList()); + } + + bucketProperty = Optional.of(new HiveBucketProperty(sd.bucketColumns(), sd.numberOfBuckets(), sortBy)); + } + + SerDeInfo serdeInfo = requireNonNull(sd.serdeInfo(), () -> "StorageDescriptor SerDeInfo is null: " + tablePartitionName); + return new Storage( + StorageFormat.createNullable(serdeInfo.serializationLibrary(), sd.inputFormat(), sd.outputFormat()), + Optional.ofNullable(sd.location()), + bucketProperty, + sd.skewedInfo() != null && !sd.skewedInfo().skewedColumnNames().isEmpty(), + serdeInfo.parameters()); + } + + private static StorageDescriptor toGlueStorage(Storage storage, List columns) + { + if (storage.isSkewed()) { + throw new IllegalArgumentException("Writing to skewed table/partition is not supported"); + } + SerDeInfo serdeInfo = SerDeInfo.builder() + .serializationLibrary(storage.getStorageFormat().getSerDeNullable()) + .parameters(storage.getSerdeParameters()) + .build(); + + StorageDescriptor.Builder builder = StorageDescriptor.builder() + .location(storage.getLocation()) + .columns(columns.stream().map(GlueConverter::toGlueColumn).collect(toImmutableList())) + .serdeInfo(serdeInfo) + .inputFormat(storage.getStorageFormat().getInputFormatNullable()) + .outputFormat(storage.getStorageFormat().getOutputFormatNullable()) + .parameters(ImmutableMap.of()); + + Optional bucketProperty = storage.getBucketProperty(); + if (bucketProperty.isPresent()) { + builder.numberOfBuckets(bucketProperty.get().getBucketCount()); + builder.bucketColumns(bucketProperty.get().getBucketedBy()); + if (!bucketProperty.get().getSortedBy().isEmpty()) { + builder.sortColumns(bucketProperty.get().getSortedBy().stream() + .map(GlueConverter::toGlueSortOrder) + .collect(toImmutableList())); + } + } + + return builder.build(); + } + + private static SortingColumn.Order fromGlueSortOrder(Integer value, String tablePartitionName) + { + if (value == 0) { + return DESCENDING; + } + if (value == 1) { + return ASCENDING; + } + throw new TrinoException(HIVE_INVALID_METADATA, "Table/partition metadata has invalid sorting order: " + tablePartitionName); + } + + private static Order toGlueSortOrder(SortingColumn column) + { + int order = switch (column.getOrder()) { + case ASCENDING -> 1; + case DESCENDING -> 0; + }; + return Order.builder() + .column(column.getColumnName()) + .sortOrder(order) + .build(); + } + + public static Map fromGlueStatistics(List> glueColumnStatistics) + { + ImmutableMap.Builder columnStatistics = ImmutableMap.builder(); + for (var columns : glueColumnStatistics) { + for (var column : columns) { + fromGlueColumnStatistics(column.statisticsData()) + .ifPresent(stats -> columnStatistics.put(column.columnName(), stats)); + } + } + return columnStatistics.buildOrThrow(); + } + + private static Optional fromGlueColumnStatistics(ColumnStatisticsData catalogColumnStatisticsData) + { + return switch (catalogColumnStatisticsData.type()) { + case BINARY -> { + BinaryColumnStatisticsData data = catalogColumnStatisticsData.binaryColumnStatisticsData(); + yield Optional.of(createBinaryColumnStatistics( + OptionalLong.of(data.maximumLength()), + OptionalDouble.of(data.averageLength()), + fromMetastoreNullsCount(data.numberOfNulls()))); + } + case BOOLEAN -> { + BooleanColumnStatisticsData catalogBooleanData = catalogColumnStatisticsData.booleanColumnStatisticsData(); + yield Optional.of(createBooleanColumnStatistics( + OptionalLong.of(catalogBooleanData.numberOfTrues()), + OptionalLong.of(catalogBooleanData.numberOfFalses()), + fromMetastoreNullsCount(catalogBooleanData.numberOfNulls()))); + } + case DATE -> { + DateColumnStatisticsData data = catalogColumnStatisticsData.dateColumnStatisticsData(); + yield Optional.of(createDateColumnStatistics( + dateToLocalDate(data.minimumValue()), + dateToLocalDate(data.maximumValue()), + fromMetastoreNullsCount(data.numberOfNulls()), + OptionalLong.of(data.numberOfDistinctValues()))); + } + case DECIMAL -> { + DecimalColumnStatisticsData data = catalogColumnStatisticsData.decimalColumnStatisticsData(); + yield Optional.of(createDecimalColumnStatistics( + fromGlueDecimal(data.minimumValue()), + fromGlueDecimal(data.maximumValue()), + fromMetastoreNullsCount(data.numberOfNulls()), + OptionalLong.of(data.numberOfDistinctValues()))); + } + case DOUBLE -> { + DoubleColumnStatisticsData data = catalogColumnStatisticsData.doubleColumnStatisticsData(); + yield Optional.of(createDoubleColumnStatistics( + OptionalDouble.of(data.minimumValue()), + OptionalDouble.of(data.maximumValue()), + fromMetastoreNullsCount(data.numberOfNulls()), + OptionalLong.of(data.numberOfDistinctValues()))); + } + case LONG -> { + LongColumnStatisticsData data = catalogColumnStatisticsData.longColumnStatisticsData(); + yield Optional.of(createIntegerColumnStatistics( + OptionalLong.of(data.minimumValue()), + OptionalLong.of(data.maximumValue()), + fromMetastoreNullsCount(data.numberOfNulls()), + OptionalLong.of(data.numberOfDistinctValues()))); + } + case STRING -> { + StringColumnStatisticsData data = catalogColumnStatisticsData.stringColumnStatisticsData(); + yield Optional.of(createStringColumnStatistics( + OptionalLong.of(data.maximumLength()), + OptionalDouble.of(data.averageLength()), + fromMetastoreNullsCount(data.numberOfNulls()), + OptionalLong.of(data.numberOfDistinctValues()))); + } + case UNKNOWN_TO_SDK_VERSION -> Optional.empty(); + }; + } + + public static List toGlueColumnStatistics(Map columnStatistics) + { + return columnStatistics.entrySet().stream() + .map(e -> toGlueColumnStatistics(e.getKey(), e.getValue())) + .flatMap(Optional::stream) + .collect(toImmutableList()); + } + + private static Optional toGlueColumnStatistics(Column column, HiveColumnStatistics statistics) + { + return toGlueColumnStatisticsData(statistics, column.getType()) + .map(columnStatisticsData -> software.amazon.awssdk.services.glue.model.ColumnStatistics.builder() + .columnName(column.getName()) + .columnType(column.getType().toString()) + .statisticsData(columnStatisticsData) + .analyzedTime(Instant.now()) + .build()); + } + + private static Optional toGlueColumnStatisticsData(HiveColumnStatistics statistics, HiveType columnType) + { + if (!isGlueWritable(statistics)) { + return Optional.empty(); + } + + if (statistics.getBooleanStatistics().isPresent()) { + BooleanStatistics booleanStatistics = statistics.getBooleanStatistics().get(); + return Optional.ofNullable(ColumnStatisticsData.builder() + .type(ColumnStatisticsType.BOOLEAN) + .booleanColumnStatisticsData(builder -> builder + .numberOfTrues(boxedValue(booleanStatistics.getTrueCount())) + .numberOfFalses(boxedValue(booleanStatistics.getFalseCount())) + .numberOfNulls(boxedValue(statistics.getNullsCount()))) + .build()); + } + if (statistics.getDateStatistics().isPresent()) { + DateStatistics dateStatistics = statistics.getDateStatistics().get(); + return Optional.ofNullable(ColumnStatisticsData.builder() + .type(ColumnStatisticsType.DATE) + .dateColumnStatisticsData(builder -> builder + .minimumValue(dateStatistics.getMin().map(GlueConverter::localDateToDate).orElse(null)) + .maximumValue(dateStatistics.getMax().map(GlueConverter::localDateToDate).orElse(null)) + .numberOfNulls(boxedValue(statistics.getNullsCount())) + .numberOfDistinctValues(boxedValue(statistics.getDistinctValuesWithNullCount()))) + .build()); + } + if (statistics.getDecimalStatistics().isPresent()) { + DecimalStatistics decimalStatistics = statistics.getDecimalStatistics().get(); + return Optional.ofNullable(ColumnStatisticsData.builder() + .type(ColumnStatisticsType.DECIMAL) + .decimalColumnStatisticsData(builder -> builder + .minimumValue(toGlueDecimal(decimalStatistics.getMin())) + .maximumValue(toGlueDecimal(decimalStatistics.getMax())) + .numberOfNulls(boxedValue(statistics.getNullsCount())) + .numberOfDistinctValues(boxedValue(statistics.getDistinctValuesWithNullCount()))) + .build()); + } + if (statistics.getDoubleStatistics().isPresent()) { + DoubleStatistics doubleStatistics = statistics.getDoubleStatistics().get(); + return Optional.ofNullable(ColumnStatisticsData.builder() + .type(ColumnStatisticsType.DOUBLE) + .doubleColumnStatisticsData(builder -> builder + .minimumValue(boxedValue(doubleStatistics.getMin())) + .maximumValue(boxedValue(doubleStatistics.getMax())) + .numberOfNulls(boxedValue(statistics.getNullsCount())) + .numberOfDistinctValues(boxedValue(statistics.getDistinctValuesWithNullCount()))) + .build()); + } + if (statistics.getIntegerStatistics().isPresent()) { + IntegerStatistics integerStatistics = statistics.getIntegerStatistics().get(); + return Optional.ofNullable(ColumnStatisticsData.builder() + .type(ColumnStatisticsType.LONG) + .longColumnStatisticsData(builder -> builder + .minimumValue(boxedValue(integerStatistics.getMin())) + .maximumValue(boxedValue(integerStatistics.getMax())) + .numberOfNulls(boxedValue(statistics.getNullsCount())) + .numberOfDistinctValues(boxedValue(statistics.getDistinctValuesWithNullCount()))) + .build()); + } + + TypeInfo typeInfo = columnType.getTypeInfo(); + if (!(typeInfo instanceof PrimitiveTypeInfo primitiveTypeInfo)) { + throw new IllegalArgumentException(lenientFormat("Unsupported statistics type: %s", columnType)); + } + PrimitiveCategory primitiveCategory = primitiveTypeInfo.getPrimitiveCategory(); + + if (PrimitiveCategory.BINARY == primitiveCategory) { + return Optional.ofNullable(ColumnStatisticsData.builder() + .type(ColumnStatisticsType.BINARY) + .binaryColumnStatisticsData(builder -> builder + .maximumLength(statistics.getMaxValueSizeInBytes().orElse(0)) + .numberOfNulls(boxedValue(statistics.getNullsCount())) + .averageLength(boxedValue(statistics.getAverageColumnLength()))) + .build()); + } + if (PrimitiveCategory.VARCHAR == primitiveCategory || (PrimitiveCategory.CHAR == primitiveCategory) || (PrimitiveCategory.STRING == primitiveCategory)) { + OptionalLong distinctValuesCount = statistics.getDistinctValuesWithNullCount(); + return Optional.ofNullable(ColumnStatisticsData.builder() + .type(ColumnStatisticsType.STRING) + .stringColumnStatisticsData(builder -> builder + .numberOfNulls(boxedValue(statistics.getNullsCount())) + .numberOfDistinctValues(boxedValue(distinctValuesCount)) + .maximumLength(statistics.getMaxValueSizeInBytes().orElse(0)) + .averageLength(boxedValue(statistics.getAverageColumnLength()))) + .build()); + } + return Optional.empty(); + } + + // Glue will accept null as min/max values, but return 0 when reading + // to avoid incorrect statistics we skip writes for columns that have min/max null + // this can be removed once Glue fixes this behavior + private static boolean isGlueWritable(HiveColumnStatistics columnStatistics) + { + if (columnStatistics.getDateStatistics().isPresent()) { + DateStatistics dateStatistics = columnStatistics.getDateStatistics().get(); + return dateStatistics.getMin().isPresent() && dateStatistics.getMax().isPresent(); + } + if (columnStatistics.getDecimalStatistics().isPresent()) { + DecimalStatistics decimalStatistics = columnStatistics.getDecimalStatistics().get(); + return decimalStatistics.getMin().isPresent() && decimalStatistics.getMax().isPresent(); + } + if (columnStatistics.getDoubleStatistics().isPresent()) { + DoubleStatistics doubleStatistics = columnStatistics.getDoubleStatistics().get(); + return doubleStatistics.getMin().isPresent() && doubleStatistics.getMax().isPresent(); + } + if (columnStatistics.getIntegerStatistics().isPresent()) { + IntegerStatistics integerStatistics = columnStatistics.getIntegerStatistics().get(); + return integerStatistics.getMin().isPresent() && integerStatistics.getMax().isPresent(); + } + return true; + } + + private static Long boxedValue(OptionalLong optionalLong) + { + return optionalLong.isPresent() ? optionalLong.getAsLong() : null; + } + + private static Double boxedValue(OptionalDouble optionalDouble) + { + return optionalDouble.isPresent() ? optionalDouble.getAsDouble() : null; + } + + private static Optional fromGlueDecimal(DecimalNumber number) + { + if (number == null) { + return Optional.empty(); + } + return Optional.of(new BigDecimal(new BigInteger(number.unscaledValue().asByteArray()), number.scale())); + } + + private static DecimalNumber toGlueDecimal(Optional optionalDecimal) + { + if (optionalDecimal.isEmpty()) { + return null; + } + BigDecimal decimal = optionalDecimal.get(); + return DecimalNumber.builder() + .unscaledValue(SdkBytes.fromByteArray(decimal.unscaledValue().toByteArray())) + .scale(decimal.scale()) + .build(); + } + + private static Optional dateToLocalDate(Instant date) + { + if (date == null) { + return Optional.empty(); + } + long daysSinceEpoch = date.getEpochSecond() / SECONDS_PER_DAY; + return Optional.of(LocalDate.ofEpochDay(daysSinceEpoch)); + } + + private static Instant localDateToDate(LocalDate date) + { + long secondsSinceEpoch = date.toEpochDay() * SECONDS_PER_DAY; + return Instant.ofEpochSecond(secondsSinceEpoch); + } + + private enum ColumnType + { + DATA, + PARTITION, + } +} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/GlueCredentialsProvider.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/GlueCredentialsProvider.java deleted file mode 100644 index a6f13e0b6d2a..000000000000 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/GlueCredentialsProvider.java +++ /dev/null @@ -1,93 +0,0 @@ -/* - * 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.hive.metastore.glue; - -import com.amazonaws.auth.AWSCredentialsProvider; -import com.amazonaws.auth.AWSStaticCredentialsProvider; -import com.amazonaws.auth.BasicAWSCredentials; -import com.amazonaws.auth.DefaultAWSCredentialsProviderChain; -import com.amazonaws.auth.STSAssumeRoleSessionCredentialsProvider; -import com.amazonaws.client.builder.AwsClientBuilder; -import com.amazonaws.services.securitytoken.AWSSecurityTokenServiceClientBuilder; -import com.google.inject.Inject; -import com.google.inject.Provider; - -import static io.trino.hdfs.s3.AwsCurrentRegionHolder.getCurrentRegionFromEC2Metadata; -import static java.lang.String.format; - -public class GlueCredentialsProvider - implements Provider -{ - private final AWSCredentialsProvider credentialsProvider; - - @Inject - public GlueCredentialsProvider(GlueHiveMetastoreConfig config) - { - if (config.getAwsCredentialsProvider().isPresent()) { - this.credentialsProvider = getCustomAWSCredentialsProvider(config.getAwsCredentialsProvider().get()); - } - else { - AWSCredentialsProvider provider; - if (config.getAwsAccessKey().isPresent() && config.getAwsSecretKey().isPresent()) { - provider = new AWSStaticCredentialsProvider( - new BasicAWSCredentials(config.getAwsAccessKey().get(), config.getAwsSecretKey().get())); - } - else { - provider = DefaultAWSCredentialsProviderChain.getInstance(); - } - if (config.getIamRole().isPresent()) { - AWSSecurityTokenServiceClientBuilder stsClientBuilder = AWSSecurityTokenServiceClientBuilder - .standard() - .withCredentials(provider); - - if (config.getGlueStsEndpointUrl().isPresent() && config.getGlueStsRegion().isPresent()) { - stsClientBuilder.setEndpointConfiguration(new AwsClientBuilder.EndpointConfiguration(config.getGlueStsEndpointUrl().get(), config.getGlueStsRegion().get())); - } - else if (config.getGlueStsRegion().isPresent()) { - stsClientBuilder.setRegion(config.getGlueStsRegion().get()); - } - else if (config.getPinGlueClientToCurrentRegion()) { - stsClientBuilder.setRegion(getCurrentRegionFromEC2Metadata().getName()); - } - - provider = new STSAssumeRoleSessionCredentialsProvider - .Builder(config.getIamRole().get(), "trino-session") - .withExternalId(config.getExternalId().orElse(null)) - .withStsClient(stsClientBuilder.build()) - .build(); - } - this.credentialsProvider = provider; - } - } - - @Override - public AWSCredentialsProvider get() - { - return credentialsProvider; - } - - private static AWSCredentialsProvider getCustomAWSCredentialsProvider(String providerClass) - { - try { - Object instance = Class.forName(providerClass).getConstructor().newInstance(); - if (!(instance instanceof AWSCredentialsProvider)) { - throw new RuntimeException("Invalid credentials provider class: " + instance.getClass().getName()); - } - return (AWSCredentialsProvider) instance; - } - catch (ReflectiveOperationException e) { - throw new RuntimeException(format("Error creating an instance of %s", providerClass), e); - } - } -} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/GlueHiveExecutionInterceptor.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/GlueHiveExecutionInterceptor.java new file mode 100644 index 000000000000..1f62c0e78dac --- /dev/null +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/GlueHiveExecutionInterceptor.java @@ -0,0 +1,42 @@ +/* + * 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.hive.metastore.glue; + +import com.google.inject.Inject; +import software.amazon.awssdk.core.SdkRequest; +import software.amazon.awssdk.core.interceptor.Context; +import software.amazon.awssdk.core.interceptor.ExecutionAttributes; +import software.amazon.awssdk.core.interceptor.ExecutionInterceptor; +import software.amazon.awssdk.services.glue.model.UpdateTableRequest; + +public class GlueHiveExecutionInterceptor + implements ExecutionInterceptor +{ + private final boolean skipArchive; + + @Inject + GlueHiveExecutionInterceptor(GlueHiveMetastoreConfig config) + { + this.skipArchive = config.isSkipArchive(); + } + + @Override + public SdkRequest modifyRequest(Context.ModifyRequest context, ExecutionAttributes executionAttributes) + { + if (context.request() instanceof UpdateTableRequest updateTableRequest) { + return updateTableRequest.toBuilder().skipArchive(skipArchive).build(); + } + return context.request(); + } +} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/GlueHiveMetastore.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/GlueHiveMetastore.java index 8936b05c9a42..11ea2d7eb2e2 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/GlueHiveMetastore.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/GlueHiveMetastore.java @@ -13,77 +13,23 @@ */ package io.trino.plugin.hive.metastore.glue; -import com.amazonaws.AmazonServiceException; -import com.amazonaws.AmazonWebServiceRequest; -import com.amazonaws.auth.DefaultAWSCredentialsProviderChain; -import com.amazonaws.handlers.AsyncHandler; -import com.amazonaws.services.glue.AWSGlueAsync; -import com.amazonaws.services.glue.model.AccessDeniedException; -import com.amazonaws.services.glue.model.AlreadyExistsException; -import com.amazonaws.services.glue.model.BatchCreatePartitionRequest; -import com.amazonaws.services.glue.model.BatchCreatePartitionResult; -import com.amazonaws.services.glue.model.BatchGetPartitionRequest; -import com.amazonaws.services.glue.model.BatchGetPartitionResult; -import com.amazonaws.services.glue.model.BatchUpdatePartitionRequest; -import com.amazonaws.services.glue.model.BatchUpdatePartitionRequestEntry; -import com.amazonaws.services.glue.model.BatchUpdatePartitionResult; -import com.amazonaws.services.glue.model.ConcurrentModificationException; -import com.amazonaws.services.glue.model.CreateDatabaseRequest; -import com.amazonaws.services.glue.model.CreateTableRequest; -import com.amazonaws.services.glue.model.DatabaseInput; -import com.amazonaws.services.glue.model.DeleteDatabaseRequest; -import com.amazonaws.services.glue.model.DeletePartitionRequest; -import com.amazonaws.services.glue.model.DeleteTableRequest; -import com.amazonaws.services.glue.model.EntityNotFoundException; -import com.amazonaws.services.glue.model.ErrorDetail; -import com.amazonaws.services.glue.model.GetDatabaseRequest; -import com.amazonaws.services.glue.model.GetDatabaseResult; -import com.amazonaws.services.glue.model.GetDatabasesRequest; -import com.amazonaws.services.glue.model.GetDatabasesResult; -import com.amazonaws.services.glue.model.GetPartitionRequest; -import com.amazonaws.services.glue.model.GetPartitionResult; -import com.amazonaws.services.glue.model.GetPartitionsRequest; -import com.amazonaws.services.glue.model.GetPartitionsResult; -import com.amazonaws.services.glue.model.GetTableRequest; -import com.amazonaws.services.glue.model.GetTableResult; -import com.amazonaws.services.glue.model.GetTablesRequest; -import com.amazonaws.services.glue.model.GetTablesResult; -import com.amazonaws.services.glue.model.PartitionError; -import com.amazonaws.services.glue.model.PartitionInput; -import com.amazonaws.services.glue.model.PartitionValueList; -import com.amazonaws.services.glue.model.Segment; -import com.amazonaws.services.glue.model.TableInput; -import com.amazonaws.services.glue.model.UpdateDatabaseRequest; -import com.amazonaws.services.glue.model.UpdatePartitionRequest; -import com.amazonaws.services.glue.model.UpdateTableRequest; -import com.google.common.annotations.VisibleForTesting; -import com.google.common.base.Stopwatch; -import com.google.common.base.Throwables; +import com.google.common.base.Strings; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; -import com.google.common.collect.Iterables; import com.google.common.collect.Lists; import com.google.inject.Inject; -import dev.failsafe.Failsafe; -import dev.failsafe.RetryPolicy; -import io.airlift.concurrent.MoreFutures; import io.airlift.log.Logger; -import io.trino.hdfs.DynamicHdfsConfiguration; -import io.trino.hdfs.HdfsConfig; -import io.trino.hdfs.HdfsConfiguration; -import io.trino.hdfs.HdfsConfigurationInitializer; -import io.trino.hdfs.HdfsContext; -import io.trino.hdfs.HdfsEnvironment; -import io.trino.hdfs.authentication.NoHdfsAuthentication; -import io.trino.plugin.hive.HiveColumnStatisticType; +import io.trino.filesystem.Location; +import io.trino.filesystem.TrinoFileSystem; +import io.trino.filesystem.TrinoFileSystemFactory; +import io.trino.plugin.hive.HiveBasicStatistics; +import io.trino.plugin.hive.HivePartitionManager; import io.trino.plugin.hive.HiveType; import io.trino.plugin.hive.PartitionNotFoundException; import io.trino.plugin.hive.PartitionStatistics; import io.trino.plugin.hive.SchemaAlreadyExistsException; import io.trino.plugin.hive.TableAlreadyExistsException; -import io.trino.plugin.hive.acid.AcidTransaction; -import io.trino.plugin.hive.aws.AwsApiCallStats; import io.trino.plugin.hive.metastore.Column; import io.trino.plugin.hive.metastore.Database; import io.trino.plugin.hive.metastore.HiveColumnStatistics; @@ -94,72 +40,102 @@ import io.trino.plugin.hive.metastore.Partition; import io.trino.plugin.hive.metastore.PartitionWithStatistics; import io.trino.plugin.hive.metastore.PrincipalPrivileges; +import io.trino.plugin.hive.metastore.StatisticsUpdateMode; import io.trino.plugin.hive.metastore.Table; -import io.trino.plugin.hive.metastore.glue.converter.GlueInputConverter; -import io.trino.plugin.hive.metastore.glue.converter.GlueToTrinoConverter; -import io.trino.plugin.hive.metastore.glue.converter.GlueToTrinoConverter.GluePartitionConverter; -import io.trino.plugin.hive.util.HiveUtil; -import io.trino.plugin.hive.util.HiveWriteUtils; +import io.trino.plugin.hive.metastore.TableInfo; +import io.trino.spi.ErrorCode; import io.trino.spi.TrinoException; -import io.trino.spi.connector.ColumnNotFoundException; import io.trino.spi.connector.SchemaNotFoundException; import io.trino.spi.connector.SchemaTableName; import io.trino.spi.connector.TableNotFoundException; import io.trino.spi.predicate.TupleDomain; import io.trino.spi.security.ConnectorIdentity; import io.trino.spi.security.RoleGrant; -import io.trino.spi.type.Type; import jakarta.annotation.Nullable; -import org.apache.hadoop.fs.Path; +import jakarta.annotation.PreDestroy; import org.weakref.jmx.Flatten; import org.weakref.jmx.Managed; - -import java.time.Duration; -import java.util.AbstractMap.SimpleEntry; +import software.amazon.awssdk.core.exception.SdkException; +import software.amazon.awssdk.services.glue.GlueClient; +import software.amazon.awssdk.services.glue.model.AlreadyExistsException; +import software.amazon.awssdk.services.glue.model.BatchGetPartitionResponse; +import software.amazon.awssdk.services.glue.model.BatchUpdatePartitionRequestEntry; +import software.amazon.awssdk.services.glue.model.ColumnStatistics; +import software.amazon.awssdk.services.glue.model.CreateTableRequest; +import software.amazon.awssdk.services.glue.model.DatabaseInput; +import software.amazon.awssdk.services.glue.model.DeleteTableRequest; +import software.amazon.awssdk.services.glue.model.EntityNotFoundException; +import software.amazon.awssdk.services.glue.model.GetDatabaseResponse; +import software.amazon.awssdk.services.glue.model.GetDatabasesResponse; +import software.amazon.awssdk.services.glue.model.GetPartitionResponse; +import software.amazon.awssdk.services.glue.model.GetPartitionsResponse; +import software.amazon.awssdk.services.glue.model.GetTableResponse; +import software.amazon.awssdk.services.glue.model.GetTablesResponse; +import software.amazon.awssdk.services.glue.model.PartitionInput; +import software.amazon.awssdk.services.glue.model.PartitionValueList; +import software.amazon.awssdk.services.glue.model.Segment; +import software.amazon.awssdk.services.glue.model.TableInput; + +import java.io.IOException; import java.util.ArrayList; -import java.util.Comparator; +import java.util.Collection; +import java.util.EnumSet; import java.util.List; import java.util.Map; import java.util.Map.Entry; +import java.util.Objects; import java.util.Optional; import java.util.OptionalLong; import java.util.Set; -import java.util.concurrent.CompletionService; +import java.util.concurrent.Callable; import java.util.concurrent.ExecutionException; -import java.util.concurrent.Executor; -import java.util.concurrent.ExecutorCompletionService; -import java.util.concurrent.Future; -import java.util.function.Function; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.Consumer; import java.util.function.Predicate; +import java.util.function.UnaryOperator; +import java.util.stream.Collectors; +import java.util.stream.IntStream; +import java.util.stream.Stream; -import static com.google.common.base.Strings.isNullOrEmpty; +import static com.google.common.base.Preconditions.checkArgument; +import static com.google.common.base.Throwables.throwIfUnchecked; import static com.google.common.base.Verify.verify; -import static com.google.common.collect.Comparators.lexicographical; import static com.google.common.collect.ImmutableList.toImmutableList; import static com.google.common.collect.ImmutableMap.toImmutableMap; import static com.google.common.collect.ImmutableSet.toImmutableSet; -import static com.google.common.util.concurrent.MoreExecutors.directExecutor; +import static io.airlift.concurrent.Threads.daemonThreadsNamed; +import static io.opentelemetry.context.Context.taskWrapping; +import static io.trino.plugin.base.util.ExecutorUtil.processWithAdditionalThreads; +import static io.trino.plugin.hive.HiveErrorCode.HIVE_FILESYSTEM_ERROR; import static io.trino.plugin.hive.HiveErrorCode.HIVE_METASTORE_ERROR; +import static io.trino.plugin.hive.HiveMetadata.TRINO_QUERY_ID_NAME; import static io.trino.plugin.hive.TableType.MANAGED_TABLE; -import static io.trino.plugin.hive.TableType.VIRTUAL_VIEW; +import static io.trino.plugin.hive.metastore.MetastoreUtil.getHiveBasicStatistics; import static io.trino.plugin.hive.metastore.MetastoreUtil.makePartitionName; import static io.trino.plugin.hive.metastore.MetastoreUtil.toPartitionName; -import static io.trino.plugin.hive.metastore.MetastoreUtil.verifyCanDropColumn; -import static io.trino.plugin.hive.metastore.glue.AwsSdkUtil.getPaginatedResults; -import static io.trino.plugin.hive.metastore.glue.GlueClientUtil.createAsyncGlueClient; -import static io.trino.plugin.hive.metastore.glue.converter.GlueInputConverter.convertPartition; -import static io.trino.plugin.hive.metastore.glue.converter.GlueToTrinoConverter.getTableParameters; -import static io.trino.plugin.hive.metastore.glue.converter.GlueToTrinoConverter.getTableTypeNullable; -import static io.trino.plugin.hive.metastore.glue.converter.GlueToTrinoConverter.mappedCopy; -import static io.trino.plugin.hive.metastore.thrift.ThriftMetastoreUtil.getHiveBasicStatistics; -import static io.trino.plugin.hive.metastore.thrift.ThriftMetastoreUtil.updateStatisticsParameters; +import static io.trino.plugin.hive.metastore.MetastoreUtil.updateStatisticsParameters; +import static io.trino.plugin.hive.metastore.Table.TABLE_COMMENT; +import static io.trino.plugin.hive.metastore.glue.GlueConverter.fromGlueStatistics; +import static io.trino.plugin.hive.metastore.glue.GlueConverter.toGlueColumnStatistics; +import static io.trino.plugin.hive.metastore.glue.GlueConverter.toGlueDatabaseInput; +import static io.trino.plugin.hive.metastore.glue.GlueConverter.toGluePartitionInput; +import static io.trino.plugin.hive.metastore.glue.GlueConverter.toGlueTableInput; import static io.trino.plugin.hive.util.HiveUtil.escapeSchemaName; -import static io.trino.plugin.hive.util.HiveUtil.toPartitionValues; +import static io.trino.plugin.hive.util.HiveUtil.isDeltaLakeTable; +import static io.trino.plugin.hive.util.HiveUtil.isHudiTable; +import static io.trino.plugin.hive.util.HiveUtil.isIcebergTable; +import static io.trino.spi.ErrorType.EXTERNAL; import static io.trino.spi.StandardErrorCode.ALREADY_EXISTS; +import static io.trino.spi.StandardErrorCode.NOT_FOUND; import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED; +import static io.trino.spi.StandardErrorCode.TABLE_NOT_FOUND; +import static io.trino.spi.StandardErrorCode.UNSUPPORTED_TABLE_TYPE; +import static io.trino.spi.connector.SchemaTableName.schemaTableName; import static io.trino.spi.security.PrincipalType.USER; +import static java.util.Map.entry; import static java.util.Objects.requireNonNull; -import static java.util.concurrent.TimeUnit.NANOSECONDS; +import static java.util.concurrent.Executors.newFixedThreadPool; import static java.util.function.Predicate.not; import static java.util.function.UnaryOperator.identity; import static java.util.stream.Collectors.toCollection; @@ -172,68 +148,78 @@ public class GlueHiveMetastore private static final String PUBLIC_ROLE_NAME = "public"; private static final String DEFAULT_METASTORE_USER = "presto"; + + // Read limit for AWS Glue API GetColumnStatisticsForPartition + // https://docs.aws.amazon.com/glue/latest/dg/aws-glue-api-catalog-partitions.html#aws-glue-api-catalog-partitions-GetColumnStatisticsForPartition + private static final int GLUE_COLUMN_READ_STAT_PAGE_SIZE = 100; + + // Write limit for AWS Glue API UpdateColumnStatisticsForPartition + // https://docs.aws.amazon.com/glue/latest/dg/aws-glue-api-catalog-partitions.html#aws-glue-api-catalog-partitions-UpdateColumnStatisticsForPartition + private static final int GLUE_COLUMN_WRITE_STAT_PAGE_SIZE = 25; + private static final int BATCH_GET_PARTITION_MAX_PAGE_SIZE = 1000; - private static final int BATCH_CREATE_PARTITION_MAX_PAGE_SIZE = 100; private static final int BATCH_UPDATE_PARTITION_MAX_PAGE_SIZE = 100; private static final int AWS_GLUE_GET_PARTITIONS_MAX_RESULTS = 1000; - private static final Comparator> PARTITION_VALUE_COMPARATOR = lexicographical(String.CASE_INSENSITIVE_ORDER); - private static final Predicate VIEWS_FILTER = table -> VIRTUAL_VIEW.name().equals(getTableTypeNullable(table)); - private static final RetryPolicy CONCURRENT_MODIFICATION_EXCEPTION_RETRY_POLICY = RetryPolicy.builder() - .handleIf(throwable -> Throwables.getRootCause(throwable) instanceof ConcurrentModificationException) - .withDelay(Duration.ofMillis(100)) - .withMaxRetries(3) - .build(); - - private final HdfsEnvironment hdfsEnvironment; - private final HdfsContext hdfsContext; - private final AWSGlueAsync glueClient; + + private static final AtomicInteger poolCounter = new AtomicInteger(); + + private final GlueClient glueClient; + private final GlueContext glueContext; + private final GlueCache glueCache; + private final TrinoFileSystem fileSystem; private final Optional defaultDir; private final int partitionSegments; - private final Executor partitionsReadExecutor; - private final GlueMetastoreStats stats; - private final GlueColumnStatisticsProvider columnStatisticsProvider; private final boolean assumeCanonicalPartitionKeys; - private final Predicate tableFilter; + private final GlueMetastoreStats stats = new GlueMetastoreStats(); + private final Predicate tableVisibilityFilter; + private final ExecutorService executor; @Inject public GlueHiveMetastore( - HdfsEnvironment hdfsEnvironment, - GlueHiveMetastoreConfig glueConfig, - @ForGlueHiveMetastore Executor partitionsReadExecutor, - GlueColumnStatisticsProviderFactory columnStatisticsProviderFactory, - AWSGlueAsync glueClient, - @ForGlueHiveMetastore GlueMetastoreStats stats, - @ForGlueHiveMetastore Predicate tableFilter) - { - this.hdfsEnvironment = requireNonNull(hdfsEnvironment, "hdfsEnvironment is null"); - this.hdfsContext = new HdfsContext(ConnectorIdentity.ofUser(DEFAULT_METASTORE_USER)); + GlueClient glueClient, + GlueContext glueContext, + GlueCache glueCache, + TrinoFileSystemFactory fileSystemFactory, + GlueHiveMetastoreConfig config, + Set visibleTableKinds) + { + this( + glueClient, + glueContext, + glueCache, + fileSystemFactory.create(ConnectorIdentity.ofUser(DEFAULT_METASTORE_USER)), + config.getDefaultWarehouseDir(), + config.getPartitionSegments(), + config.isAssumeCanonicalPartitionKeys(), + visibleTableKinds, + newFixedThreadPool(config.getThreads(), daemonThreadsNamed("glue-%s-%%s".formatted(poolCounter.getAndIncrement())))); + } + + private GlueHiveMetastore( + GlueClient glueClient, + GlueContext glueContext, + GlueCache glueCache, TrinoFileSystem fileSystem, + Optional defaultDir, + int partitionSegments, + boolean assumeCanonicalPartitionKeys, + Set visibleTableKinds, + ExecutorService executor) + { this.glueClient = requireNonNull(glueClient, "glueClient is null"); - this.defaultDir = glueConfig.getDefaultWarehouseDir(); - this.partitionSegments = glueConfig.getPartitionSegments(); - this.partitionsReadExecutor = requireNonNull(partitionsReadExecutor, "partitionsReadExecutor is null"); - this.assumeCanonicalPartitionKeys = glueConfig.isAssumeCanonicalPartitionKeys(); - this.tableFilter = requireNonNull(tableFilter, "tableFilter is null"); - this.stats = requireNonNull(stats, "stats is null"); - this.columnStatisticsProvider = columnStatisticsProviderFactory.createGlueColumnStatisticsProvider(glueClient, stats); - } - - @VisibleForTesting - public static GlueHiveMetastore createTestingGlueHiveMetastore(java.nio.file.Path defaultWarehouseDir) - { - HdfsConfig hdfsConfig = new HdfsConfig(); - HdfsConfiguration hdfsConfiguration = new DynamicHdfsConfiguration(new HdfsConfigurationInitializer(hdfsConfig), ImmutableSet.of()); - HdfsEnvironment hdfsEnvironment = new HdfsEnvironment(hdfsConfiguration, hdfsConfig, new NoHdfsAuthentication()); - GlueMetastoreStats stats = new GlueMetastoreStats(); - GlueHiveMetastoreConfig glueConfig = new GlueHiveMetastoreConfig() - .setDefaultWarehouseDir(defaultWarehouseDir.toUri().toString()); - return new GlueHiveMetastore( - hdfsEnvironment, - glueConfig, - directExecutor(), - new DefaultGlueColumnStatisticsProviderFactory(directExecutor(), directExecutor()), - createAsyncGlueClient(glueConfig, DefaultAWSCredentialsProviderChain.getInstance(), ImmutableSet.of(), stats.newRequestMetricsCollector()), - stats, - table -> true); + this.glueContext = requireNonNull(glueContext, "glueContext is null"); + this.glueCache = glueCache; + this.fileSystem = requireNonNull(fileSystem, "fileSystem is null"); + this.defaultDir = requireNonNull(defaultDir, "defaultDir is null"); + this.partitionSegments = partitionSegments; + this.assumeCanonicalPartitionKeys = assumeCanonicalPartitionKeys; + this.tableVisibilityFilter = createTablePredicate(visibleTableKinds); + this.executor = taskWrapping(requireNonNull(executor, "executor is null")); + } + + @PreDestroy + public void shutdown() + { + executor.shutdownNow(); } @Managed @@ -243,374 +229,362 @@ public GlueMetastoreStats getStats() return stats; } - @Override - public Optional getDatabase(String databaseName) + @Managed + @Flatten + public GlueCache getGlueCache() { - try { - GetDatabaseResult result = stats.getGetDatabase().call(() -> - glueClient.getDatabase(new GetDatabaseRequest().withName(databaseName))); - return Optional.of(GlueToTrinoConverter.convertDatabase(result.getDatabase())); - } - catch (EntityNotFoundException e) { - return Optional.empty(); - } - catch (AmazonServiceException e) { - throw new TrinoException(HIVE_METASTORE_ERROR, e); - } + return glueCache; } @Override public List getAllDatabases() + { + return glueCache.getDatabaseNames(this::getDatabasesInternal); + } + + private List getDatabasesInternal(Consumer cacheDatabase) { try { - List databaseNames = getPaginatedResults( - glueClient::getDatabases, - new GetDatabasesRequest(), - GetDatabasesRequest::setNextToken, - GetDatabasesResult::getNextToken, - stats.getGetDatabases()) - .map(GetDatabasesResult::getDatabaseList) + return stats.getGetDatabases().call(() -> glueClient.getDatabasesPaginator(glueContext::configureClient).stream() + .map(GetDatabasesResponse::databaseList) .flatMap(List::stream) - .map(com.amazonaws.services.glue.model.Database::getName) - .collect(toImmutableList()); - return databaseNames; + .map(GlueConverter::fromGlueDatabase) + .peek(cacheDatabase) + .map(Database::getDatabaseName) + .toList()); } - catch (AmazonServiceException e) { + catch (SdkException e) { throw new TrinoException(HIVE_METASTORE_ERROR, e); } } @Override - public Optional
getTable(String databaseName, String tableName) + public Optional getDatabase(String databaseName) + { + return glueCache.getDatabase(databaseName, () -> getDatabaseInternal(databaseName)); + } + + private Optional getDatabaseInternal(String databaseName) { try { - GetTableResult result = stats.getGetTable().call(() -> - glueClient.getTable(new GetTableRequest() - .withDatabaseName(databaseName) - .withName(tableName))); - return Optional.of(GlueToTrinoConverter.convertTable(result.getTable(), databaseName)); + GetDatabaseResponse response = stats.getGetDatabase().call(() -> glueClient.getDatabase(builder -> builder + .applyMutation(glueContext::configureClient) + .name(databaseName))); + return Optional.of(GlueConverter.fromGlueDatabase(response.database())); } catch (EntityNotFoundException e) { return Optional.empty(); } - catch (AmazonServiceException e) { + catch (SdkException e) { throw new TrinoException(HIVE_METASTORE_ERROR, e); } } @Override - public Set getSupportedColumnStatistics(Type type) - { - return columnStatisticsProvider.getSupportedColumnStatistics(type); - } - - private Table getExistingTable(String databaseName, String tableName) - { - return getTable(databaseName, tableName) - .orElseThrow(() -> new TableNotFoundException(new SchemaTableName(databaseName, tableName))); - } - - @Override - public PartitionStatistics getTableStatistics(Table table) - { - return new PartitionStatistics(getHiveBasicStatistics(table.getParameters()), columnStatisticsProvider.getTableColumnStatistics(table)); - } - - @Override - public Map getPartitionStatistics(Table table, List partitions) + public void createDatabase(Database database) { - Map partitionBasicStatistics = columnStatisticsProvider.getPartitionColumnStatistics(partitions).entrySet().stream() - .collect(toImmutableMap( - entry -> makePartitionName(table, entry.getKey()), - entry -> new PartitionStatistics(getHiveBasicStatistics(entry.getKey().getParameters()), entry.getValue()))); + if (database.getLocation().isEmpty() && defaultDir.isPresent()) { + Location location = Location.of(defaultDir.get()) + .appendPath(escapeSchemaName(database.getDatabaseName())); + database = Database.builder(database) + .setLocation(Optional.of(location.toString())) + .build(); + } - long tableRowCount = partitionBasicStatistics.values().stream() - .mapToLong(partitionStatistics -> partitionStatistics.getBasicStatistics().getRowCount().orElse(0)) - .sum(); - if (!partitionBasicStatistics.isEmpty() && tableRowCount == 0) { - // When the table has partitions, but row count statistics are set to zero, we treat this case as empty - // statistics to avoid underestimation in the CBO. This scenario may be caused when other engines are - // used to ingest data into partitioned hive tables. - partitionBasicStatistics = partitionBasicStatistics.keySet().stream() - .map(key -> new SimpleEntry<>(key, PartitionStatistics.empty())) - .collect(toImmutableMap(SimpleEntry::getKey, SimpleEntry::getValue)); + try { + DatabaseInput databaseInput = toGlueDatabaseInput(database); + stats.getCreateDatabase().call(() -> glueClient.createDatabase(builder -> builder + .applyMutation(glueContext::configureClient) + .databaseInput(databaseInput))); + } + catch (AlreadyExistsException e) { + // Do not throw SchemaAlreadyExistsException if this query has already created the database. + // This may happen when an actually successful metastore create call is retried + // because of a timeout on our side. + String expectedQueryId = database.getParameters().get(TRINO_QUERY_ID_NAME); + if (expectedQueryId != null) { + String existingQueryId = getDatabase(database.getDatabaseName()) + .map(Database::getParameters) + .map(parameters -> parameters.get(TRINO_QUERY_ID_NAME)) + .orElse(null); + if (expectedQueryId.equals(existingQueryId)) { + return; + } + } + throw new SchemaAlreadyExistsException(database.getDatabaseName(), e); + } + catch (SdkException e) { + throw new TrinoException(HIVE_METASTORE_ERROR, e); + } + finally { + glueCache.invalidateDatabase(database.getDatabaseName()); + glueCache.invalidateDatabaseNames(); } - return partitionBasicStatistics; + if (database.getLocation().isPresent()) { + Location location = Location.of(database.getLocation().get()); + try { + fileSystem.createDirectory(location); + } + catch (IOException e) { + throw new TrinoException(HIVE_FILESYSTEM_ERROR, "Failed to create directory: " + location, e); + } + } } @Override - public void updateTableStatistics(String databaseName, String tableName, AcidTransaction transaction, Function update) + public void dropDatabase(String databaseName, boolean deleteData) { - Table table = getExistingTable(databaseName, tableName); - if (transaction.isAcidTransactionRunning()) { - table = Table.builder(table).setWriteId(OptionalLong.of(transaction.getWriteId())).build(); + Optional location = Optional.empty(); + if (deleteData) { + location = getDatabase(databaseName) + .orElseThrow(() -> new SchemaNotFoundException(databaseName)) + .getLocation() + .map(Location::of); } - PartitionStatistics currentStatistics = getTableStatistics(table); - PartitionStatistics updatedStatistics = update.apply(currentStatistics); try { - TableInput tableInput = GlueInputConverter.convertTable(table); - final Map statisticsParameters = updateStatisticsParameters(table.getParameters(), updatedStatistics.getBasicStatistics()); - tableInput.setParameters(statisticsParameters); - table = Table.builder(table).setParameters(statisticsParameters).build(); - stats.getUpdateTable().call(() -> glueClient.updateTable(new UpdateTableRequest() - .withDatabaseName(databaseName) - .withTableInput(tableInput))); - columnStatisticsProvider.updateTableColumnStatistics(table, updatedStatistics.getColumnStatistics()); + stats.getDeleteDatabase().call(() -> glueClient.deleteDatabase(builder -> builder + .applyMutation(glueContext::configureClient) + .name(databaseName))); } catch (EntityNotFoundException e) { - throw new TableNotFoundException(new SchemaTableName(databaseName, tableName)); + throw new SchemaNotFoundException(databaseName, e); } - catch (AmazonServiceException e) { + catch (SdkException e) { throw new TrinoException(HIVE_METASTORE_ERROR, e); } - } + finally { + glueCache.invalidateDatabase(databaseName); + glueCache.invalidateDatabaseNames(); + } - @Override - public void updatePartitionStatistics(Table table, Map> updates) - { - Iterables.partition(updates.entrySet(), BATCH_CREATE_PARTITION_MAX_PAGE_SIZE).forEach(partitionUpdates -> - updatePartitionStatisticsBatch(table, partitionUpdates.stream().collect(toImmutableMap(Entry::getKey, Entry::getValue)))); + location.ifPresent(this::deleteDir); } - private void updatePartitionStatisticsBatch(Table table, Map> updates) + @Override + public void renameDatabase(String databaseName, String newDatabaseName) { - ImmutableList.Builder partitionUpdateRequests = ImmutableList.builder(); - ImmutableSet.Builder columnStatisticsUpdates = ImmutableSet.builder(); - - Map, String> partitionValuesToName = updates.keySet().stream() - .collect(toImmutableMap(HiveUtil::toPartitionValues, identity())); - - List partitions = batchGetPartition(table, ImmutableList.copyOf(updates.keySet())); - Map partitionsStatistics = getPartitionStatistics(table, partitions); - - partitions.forEach(partition -> { - Function update = updates.get(partitionValuesToName.get(partition.getValues())); - - PartitionStatistics currentStatistics = partitionsStatistics.get(makePartitionName(table, partition)); - PartitionStatistics updatedStatistics = update.apply(currentStatistics); - - Map updatedStatisticsParameters = updateStatisticsParameters(partition.getParameters(), updatedStatistics.getBasicStatistics()); - - partition = Partition.builder(partition).setParameters(updatedStatisticsParameters).build(); - Map updatedColumnStatistics = updatedStatistics.getColumnStatistics(); - - PartitionInput partitionInput = GlueInputConverter.convertPartition(partition); - partitionInput.setParameters(partition.getParameters()); - - partitionUpdateRequests.add(new BatchUpdatePartitionRequestEntry() - .withPartitionValueList(partition.getValues()) - .withPartitionInput(partitionInput)); - columnStatisticsUpdates.add(new GlueColumnStatisticsProvider.PartitionStatisticsUpdate(partition, updatedColumnStatistics)); - }); - - List> partitionUpdateRequestsPartitioned = Lists.partition(partitionUpdateRequests.build(), BATCH_UPDATE_PARTITION_MAX_PAGE_SIZE); - List> partitionUpdateRequestsFutures = new ArrayList<>(); - partitionUpdateRequestsPartitioned.forEach(partitionUpdateRequestsPartition -> { - // Update basic statistics - partitionUpdateRequestsFutures.add(glueClient.batchUpdatePartitionAsync(new BatchUpdatePartitionRequest() - .withDatabaseName(table.getDatabaseName()) - .withTableName(table.getTableName()) - .withEntries(partitionUpdateRequestsPartition), - new StatsRecordingAsyncHandler<>(stats.getBatchUpdatePartition()))); - }); - try { - // Update column statistics - columnStatisticsProvider.updatePartitionStatistics(columnStatisticsUpdates.build()); - // Don't block on the batch update call until the column statistics have finished updating - partitionUpdateRequestsFutures.forEach(MoreFutures::getFutureValue); + var database = stats.getGetDatabase().call(() -> glueClient.getDatabase(builder -> builder + .applyMutation(glueContext::configureClient) + .name(databaseName)).database()); + DatabaseInput renamedDatabase = DatabaseInput.builder() + .name(newDatabaseName) + .parameters(database.parameters()) + .description(database.description()) + .locationUri(database.locationUri()) + .build(); + stats.getUpdateDatabase().call(() -> glueClient.updateDatabase(builder -> builder + .applyMutation(glueContext::configureClient) + .name(databaseName).databaseInput(renamedDatabase))); + } + catch (EntityNotFoundException e) { + throw new SchemaNotFoundException(databaseName, e); + } + catch (AlreadyExistsException e) { + throw new SchemaAlreadyExistsException(newDatabaseName, e); } - catch (AmazonServiceException e) { + catch (SdkException e) { throw new TrinoException(HIVE_METASTORE_ERROR, e); } + finally { + glueCache.invalidateDatabase(databaseName); + glueCache.invalidateDatabase(newDatabaseName); + glueCache.invalidateDatabaseNames(); + } } @Override - public List getAllTables(String databaseName) - { - return getTableNames(databaseName, tableFilter); - } - - @Override - public Optional> getAllTables() - { - return Optional.empty(); - } - - @Override - public List getTablesWithParameter(String databaseName, String parameterKey, String parameterValue) + public void setDatabaseOwner(String databaseName, HivePrincipal principal) { - return getTableNames(databaseName, table -> parameterValue.equals(getTableParameters(table).get(parameterKey))); + throw new TrinoException(NOT_SUPPORTED, "setting the database owner is not supported by Glue"); } @Override - public List getAllViews(String databaseName) + public List getTables(String databaseName) { - return getTableNames(databaseName, VIEWS_FILTER); + return glueCache.getTables(databaseName, cacheTable -> getTablesInternal(cacheTable, databaseName, ignore -> true)); } @Override - public Optional> getAllViews() + public List getTableNamesWithParameters(String databaseName, String parameterKey, ImmutableSet parameterValues) { - return Optional.empty(); + return getTablesInternal( + ignore -> {}, + databaseName, + table -> table.parameters() != null && parameterValues.contains(table.parameters().get(parameterKey))).stream() + .map(tableInfo -> tableInfo.tableName().getTableName()) + .collect(toImmutableList()); } - private List getTableNames(String databaseName, Predicate filter) + private List getTablesInternal(Consumer
cacheTable, String databaseName, Predicate filter) { try { - List tableNames = getPaginatedResults( - glueClient::getTables, - new GetTablesRequest() - .withDatabaseName(databaseName), - GetTablesRequest::setNextToken, - GetTablesResult::getNextToken, - stats.getGetTables()) - .map(GetTablesResult::getTableList) - .flatMap(List::stream) + ImmutableList glueTables = stats.getGetTables() + .call(() -> glueClient.getTablesPaginator(builder -> builder + .applyMutation(glueContext::configureClient) + .databaseName(databaseName)).stream() + .map(GetTablesResponse::tableList) + .flatMap(List::stream)) + .filter(tableVisibilityFilter) .filter(filter) - .map(com.amazonaws.services.glue.model.Table::getName) .collect(toImmutableList()); - return tableNames; + + // Store only valid tables in cache + for (software.amazon.awssdk.services.glue.model.Table table : glueTables) { + convertFromGlueIgnoringErrors(table, databaseName) + .ifPresent(cacheTable); + } + + return glueTables.stream() + .map(table -> new TableInfo( + new SchemaTableName(databaseName, table.name()), + TableInfo.ExtendedRelationType.fromTableTypeAndComment(GlueConverter.getTableType(table), table.parameters().get(TABLE_COMMENT)))) + .toList(); } - catch (EntityNotFoundException | AccessDeniedException e) { - // database does not exist or permission denied + catch (EntityNotFoundException ignore) { + // Database might have been deleted concurrently. return ImmutableList.of(); } - catch (AmazonServiceException e) { + catch (SdkException e) { throw new TrinoException(HIVE_METASTORE_ERROR, e); } } - @Override - public void createDatabase(Database database) + private static Optional
convertFromGlueIgnoringErrors(software.amazon.awssdk.services.glue.model.Table glueTable, String databaseName) { - if (database.getLocation().isEmpty() && defaultDir.isPresent()) { - String databaseLocation = new Path(defaultDir.get(), escapeSchemaName(database.getDatabaseName())).toString(); - database = Database.builder(database) - .setLocation(Optional.of(databaseLocation)) - .build(); - } - try { - DatabaseInput databaseInput = GlueInputConverter.convertDatabase(database); - stats.getCreateDatabase().call(() -> - glueClient.createDatabase(new CreateDatabaseRequest().withDatabaseInput(databaseInput))); - } - catch (AlreadyExistsException e) { - throw new SchemaAlreadyExistsException(database.getDatabaseName()); + return Optional.of(GlueConverter.fromGlueTable(glueTable, databaseName)); } - catch (AmazonServiceException e) { - throw new TrinoException(HIVE_METASTORE_ERROR, e); + catch (RuntimeException e) { + handleListingError(e, schemaTableName(glueTable.databaseName(), glueTable.name())); + return Optional.empty(); } + } - if (database.getLocation().isPresent()) { - HiveWriteUtils.createDirectory(hdfsContext, hdfsEnvironment, new Path(database.getLocation().get())); + private static void handleListingError(RuntimeException e, SchemaTableName tableName) + { + boolean silent = false; + if (e instanceof TrinoException trinoException) { + ErrorCode errorCode = trinoException.getErrorCode(); + silent = errorCode.equals(UNSUPPORTED_TABLE_TYPE.toErrorCode()) || + // e.g. table deleted concurrently + errorCode.equals(TABLE_NOT_FOUND.toErrorCode()) || + errorCode.equals(NOT_FOUND.toErrorCode()) || + // e.g. Iceberg/Delta table being deleted concurrently resulting in failure to load metadata from filesystem + errorCode.getType() == EXTERNAL; + } + if (silent) { + log.debug(e, "Failed to get metadata for table: %s", tableName); + } + else { + log.warn(e, "Failed to get metadata for table: %s", tableName); } } @Override - public void dropDatabase(String databaseName, boolean deleteData) + public Optional
getTable(String databaseName, String tableName) { - Optional location = Optional.empty(); - if (deleteData) { - location = getDatabase(databaseName) - .orElseThrow(() -> new SchemaNotFoundException(databaseName)) - .getLocation(); - } + return glueCache.getTable(databaseName, tableName, () -> getTableInternal(databaseName, tableName)); + } + private Optional
getTableInternal(String databaseName, String tableName) + { try { - stats.getDeleteDatabase().call(() -> - glueClient.deleteDatabase(new DeleteDatabaseRequest().withName(databaseName))); + GetTableResponse result = stats.getGetTable().call(() -> glueClient.getTable(builder -> builder + .applyMutation(glueContext::configureClient) + .databaseName(databaseName) + .name(tableName))); + return Optional.of(GlueConverter.fromGlueTable(result.table(), databaseName)); } catch (EntityNotFoundException e) { - throw new SchemaNotFoundException(databaseName); - } - catch (AmazonServiceException e) { - throw new TrinoException(HIVE_METASTORE_ERROR, e); - } - - if (deleteData) { - location.ifPresent(path -> deleteDir(hdfsContext, hdfsEnvironment, new Path(path), true)); + return Optional.empty(); } - } - - @Override - public void renameDatabase(String databaseName, String newDatabaseName) - { - try { - Database database = getDatabase(databaseName).orElseThrow(() -> new SchemaNotFoundException(databaseName)); - DatabaseInput renamedDatabase = GlueInputConverter.convertDatabase(database).withName(newDatabaseName); - stats.getUpdateDatabase().call(() -> - glueClient.updateDatabase(new UpdateDatabaseRequest() - .withName(databaseName) - .withDatabaseInput(renamedDatabase))); - } - catch (AmazonServiceException e) { + catch (SdkException e) { throw new TrinoException(HIVE_METASTORE_ERROR, e); } } - @Override - public void setDatabaseOwner(String databaseName, HivePrincipal principal) - { - throw new TrinoException(NOT_SUPPORTED, "setting the database owner is not supported by Glue"); - } - @Override public void createTable(Table table, PrincipalPrivileges principalPrivileges) { try { - TableInput input = GlueInputConverter.convertTable(table); - stats.getCreateTable().call(() -> - glueClient.createTable(new CreateTableRequest() - .withDatabaseName(table.getDatabaseName()) - .withTableInput(input))); + TableInput input = toGlueTableInput(table); + stats.getCreateTable().call(() -> glueClient.createTable(builder -> builder + .applyMutation(glueContext::configureClient) + .databaseName(table.getDatabaseName()) + .tableInput(input))); } catch (AlreadyExistsException e) { - throw new TableAlreadyExistsException(new SchemaTableName(table.getDatabaseName(), table.getTableName())); + // Do not throw TableAlreadyExistsException if this query has already created the table. + // This may happen when an actually successful metastore create call is retried + // because of a timeout on our side. + String expectedQueryId = table.getParameters().get(TRINO_QUERY_ID_NAME); + if (expectedQueryId != null) { + String existingQueryId = getTable(table.getDatabaseName(), table.getTableName()) + .map(Table::getParameters) + .map(parameters -> parameters.get(TRINO_QUERY_ID_NAME)) + .orElse(null); + if (expectedQueryId.equals(existingQueryId)) { + return; + } + } + throw new TableAlreadyExistsException(new SchemaTableName(table.getDatabaseName(), table.getTableName()), e); } catch (EntityNotFoundException e) { - throw new SchemaNotFoundException(table.getDatabaseName()); + throw new SchemaNotFoundException(table.getDatabaseName(), e); } - catch (AmazonServiceException e) { + catch (SdkException e) { throw new TrinoException(HIVE_METASTORE_ERROR, e); } + finally { + glueCache.invalidateTable(table.getDatabaseName(), table.getTableName(), true); + glueCache.invalidateTables(table.getDatabaseName()); + } } @Override public void dropTable(String databaseName, String tableName, boolean deleteData) { - Table table = getExistingTable(databaseName, tableName); - DeleteTableRequest deleteTableRequest = new DeleteTableRequest() - .withDatabaseName(databaseName) - .withName(tableName); + Optional deleteLocation = Optional.empty(); + if (deleteData) { + Table table = getTable(databaseName, tableName) + .orElseThrow(() -> new TableNotFoundException(new SchemaTableName(databaseName, tableName))); + if (table.getTableType().equals(MANAGED_TABLE.name())) { + deleteLocation = table.getStorage() + .getOptionalLocation() + .filter(not(String::isEmpty)) + .map(Location::of); + } + } + + DeleteTableRequest deleteTableRequest = DeleteTableRequest.builder() + .applyMutation(glueContext::configureClient) + .databaseName(databaseName) + .name(tableName) + .build(); try { - Failsafe.with(CONCURRENT_MODIFICATION_EXCEPTION_RETRY_POLICY) - .run(() -> stats.getDeleteTable().call(() -> - glueClient.deleteTable(deleteTableRequest))); + stats.getDeleteTable().call(() -> glueClient.deleteTable(deleteTableRequest)); } - catch (AmazonServiceException e) { + catch (EntityNotFoundException e) { + throw new TableNotFoundException(new SchemaTableName(databaseName, tableName)); + } + catch (SdkException e) { throw new TrinoException(HIVE_METASTORE_ERROR, e); } - - Optional location = table.getStorage().getOptionalLocation() - .filter(not(String::isEmpty)); - if (deleteData && isManagedTable(table) && location.isPresent()) { - deleteDir(hdfsContext, hdfsEnvironment, new Path(location.get()), true); + finally { + glueCache.invalidateTable(databaseName, tableName, true); + glueCache.invalidateTables(databaseName); } - } - private static boolean isManagedTable(Table table) - { - return table.getTableType().equals(MANAGED_TABLE.name()); + deleteLocation.ifPresent(this::deleteDir); } - private static void deleteDir(HdfsContext context, HdfsEnvironment hdfsEnvironment, Path path, boolean recursive) + private void deleteDir(Location path) { try { - hdfsEnvironment.getFileSystem(context, path).delete(path, recursive); + fileSystem.deleteDirectory(path); } catch (Exception e) { // don't fail if unable to delete path @@ -619,473 +593,827 @@ private static void deleteDir(HdfsContext context, HdfsEnvironment hdfsEnvironme } @Override - public void replaceTable(String databaseName, String tableName, Table newTable, PrincipalPrivileges principalPrivileges) + public void replaceTable(String databaseName, String tableName, Table newTable, PrincipalPrivileges principalPrivileges, Map environmentContext) { if (!tableName.equals(newTable.getTableName()) || !databaseName.equals(newTable.getDatabaseName())) { throw new TrinoException(NOT_SUPPORTED, "Table rename is not yet supported by Glue service"); } + + updateTable(databaseName, tableName, ignore -> newTable); + } + + private void updateTable(String databaseName, String tableName, UnaryOperator
modifier) + { + Table existingTable = getTable(databaseName, tableName) + .orElseThrow(() -> new TableNotFoundException(new SchemaTableName(databaseName, tableName))); + Table newTable = modifier.apply(existingTable); + if (!existingTable.getDatabaseName().equals(newTable.getDatabaseName()) || !existingTable.getTableName().equals(newTable.getTableName())) { + throw new TrinoException(NOT_SUPPORTED, "Update cannot be used to change rename a table"); + } + if (existingTable.getParameters().getOrDefault("table_type", "").equalsIgnoreCase("iceberg") && !Objects.equals( + existingTable.getParameters().get("metadata_location"), + newTable.getParameters().get("previous_metadata_location"))) { + throw new TrinoException(NOT_SUPPORTED, "Cannot update Iceberg table: supplied previous location does not match current location"); + } try { - TableInput newTableInput = GlueInputConverter.convertTable(newTable); - stats.getUpdateTable().call(() -> - glueClient.updateTable(new UpdateTableRequest() - .withDatabaseName(databaseName) - .withTableInput(newTableInput))); + stats.getUpdateTable().call(() -> glueClient.updateTable(builder -> builder + .applyMutation(glueContext::configureClient) + .databaseName(databaseName) + .tableInput(toGlueTableInput(newTable)))); } catch (EntityNotFoundException e) { throw new TableNotFoundException(new SchemaTableName(databaseName, tableName)); } - catch (AmazonServiceException e) { + catch (SdkException e) { throw new TrinoException(HIVE_METASTORE_ERROR, e); } + finally { + glueCache.invalidateTable(databaseName, tableName, false); + } } @Override public void renameTable(String databaseName, String tableName, String newDatabaseName, String newTableName) { - boolean newTableCreated = false; + // read the existing table + software.amazon.awssdk.services.glue.model.Table table; + try { + table = stats.getGetTable().call(() -> glueClient.getTable(builder -> builder + .applyMutation(glueContext::configureClient) + .databaseName(databaseName) + .name(tableName)) + .table()); + } + catch (EntityNotFoundException e) { + throw new TableNotFoundException(new SchemaTableName(databaseName, tableName)); + } + catch (SdkException e) { + throw new TrinoException(HIVE_METASTORE_ERROR, e); + } + + // create a new table with the same data as the old table try { - GetTableRequest getTableRequest = new GetTableRequest().withDatabaseName(databaseName) - .withName(tableName); - GetTableResult glueTable = glueClient.getTable(getTableRequest); - TableInput tableInput = convertGlueTableToTableInput(glueTable.getTable(), newTableName); - CreateTableRequest createTableRequest = new CreateTableRequest() - .withDatabaseName(newDatabaseName) - .withTableInput(tableInput); + CreateTableRequest createTableRequest = CreateTableRequest.builder() + .applyMutation(glueContext::configureClient) + .databaseName(newDatabaseName) + .tableInput(asTableInputBuilder(table) + .name(newTableName) + .build()) + .build(); stats.getCreateTable().call(() -> glueClient.createTable(createTableRequest)); - newTableCreated = true; + } + catch (AlreadyExistsException e) { + throw new TableAlreadyExistsException(new SchemaTableName(databaseName, tableName)); + } + catch (SdkException e) { + throw new TrinoException(HIVE_METASTORE_ERROR, e); + } + + // drop the old table + try { dropTable(databaseName, tableName, false); } catch (RuntimeException e) { - if (newTableCreated) { - try { - dropTable(databaseName, tableName, false); - } - catch (RuntimeException cleanupException) { - if (!cleanupException.equals(e)) { - e.addSuppressed(cleanupException); - } + // if the drop failed, try to clean up the new table + try { + dropTable(databaseName, tableName, false); + } + catch (RuntimeException cleanupException) { + if (!cleanupException.equals(e)) { + e.addSuppressed(cleanupException); } } throw e; } + finally { + glueCache.invalidateTable(databaseName, tableName, true); + glueCache.invalidateTable(newDatabaseName, newTableName, true); + glueCache.invalidateTables(databaseName); + if (!databaseName.equals(newDatabaseName)) { + glueCache.invalidateTables(newDatabaseName); + } + } } - private TableInput convertGlueTableToTableInput(com.amazonaws.services.glue.model.Table glueTable, String newTableName) + public static TableInput.Builder asTableInputBuilder(software.amazon.awssdk.services.glue.model.Table table) { - return new TableInput() - .withName(newTableName) - .withDescription(glueTable.getDescription()) - .withOwner(glueTable.getOwner()) - .withLastAccessTime(glueTable.getLastAccessTime()) - .withLastAnalyzedTime(glueTable.getLastAnalyzedTime()) - .withRetention(glueTable.getRetention()) - .withStorageDescriptor(glueTable.getStorageDescriptor()) - .withPartitionKeys(glueTable.getPartitionKeys()) - .withViewOriginalText(glueTable.getViewOriginalText()) - .withViewExpandedText(glueTable.getViewExpandedText()) - .withTableType(getTableTypeNullable(glueTable)) - .withTargetTable(glueTable.getTargetTable()) - .withParameters(getTableParameters(glueTable)); + return TableInput.builder() + .name(table.name()) + .description(table.description()) + .owner(table.owner()) + .lastAccessTime(table.lastAccessTime()) + .lastAnalyzedTime(table.lastAnalyzedTime()) + .retention(table.retention()) + .storageDescriptor(table.storageDescriptor()) + .partitionKeys(table.partitionKeys()) + .viewOriginalText(table.viewOriginalText()) + .viewExpandedText(table.viewExpandedText()) + .tableType(GlueConverter.getTableTypeNullable(table)) + .targetTable(table.targetTable()) + .parameters(table.parameters()); } @Override public void commentTable(String databaseName, String tableName, Optional comment) { - throw new TrinoException(NOT_SUPPORTED, "Table comment is not yet supported by Glue service"); + updateTable(databaseName, tableName, table -> table.withComment(comment)); } @Override public void setTableOwner(String databaseName, String tableName, HivePrincipal principal) { - // TODO Add role support https://github.com/trinodb/trino/issues/5706 - if (principal.getType() != USER) { - throw new TrinoException(NOT_SUPPORTED, "Setting table owner type as a role is not supported"); - } + updateTable(databaseName, tableName, table -> table.withOwner(principal)); + } + + @Override + public Map getTableColumnStatistics(String databaseName, String tableName, Set columnNames) + { + return glueCache.getTableColumnStatistics( + databaseName, + tableName, + columnNames, + missingColumnNames -> getTableColumnStatisticsInternal(databaseName, tableName, missingColumnNames)); + } + + private Map getTableColumnStatisticsInternal(String databaseName, String tableName, Set columnNames) + { + var columnStatsTasks = Lists.partition(ImmutableList.copyOf(columnNames), GLUE_COLUMN_READ_STAT_PAGE_SIZE).stream() + .map(partialColumns -> (Callable>) () -> stats.getGetColumnStatisticsForTable().call(() -> glueClient.getColumnStatisticsForTable(builder -> builder + .applyMutation(glueContext::configureClient) + .databaseName(databaseName) + .tableName(tableName) + .columnNames(partialColumns)) + .columnStatisticsList())) + .collect(toImmutableList()); try { - Table table = getExistingTable(databaseName, tableName); - TableInput newTableInput = GlueInputConverter.convertTable(table); - newTableInput.setOwner(principal.getName()); + return fromGlueStatistics(runParallel(columnStatsTasks)); + } + catch (ExecutionException e) { + if (e.getCause() instanceof EntityNotFoundException) { + throw new TableNotFoundException(new SchemaTableName(databaseName, tableName)); + } + throw new TrinoException(HIVE_METASTORE_ERROR, "Failed fetching column statistics for %s.%s".formatted(databaseName, tableName), e.getCause()); + } + } + + @Override + public void updateTableStatistics(String databaseName, String tableName, OptionalLong acidWriteId, StatisticsUpdateMode mode, PartitionStatistics statisticsUpdate) + { + verify(acidWriteId.isEmpty(), "Glue metastore does not support Hive Acid tables"); - stats.getUpdateTable().call(() -> - glueClient.updateTable(new UpdateTableRequest() - .withDatabaseName(databaseName) - .withTableInput(newTableInput))); + // adding zero rows does not change stats + if (mode == StatisticsUpdateMode.MERGE_INCREMENTAL && statisticsUpdate.getBasicStatistics().getRowCount().equals(OptionalLong.of(0))) { + return; } - catch (EntityNotFoundException e) { - throw new TableNotFoundException(new SchemaTableName(databaseName, tableName)); + + // existing column statistics are required, for the columns being updated when in merge incremental mode + // this is fetched before the basic statistics are updated, to avoid reloading in the case of retries updating basic statistics + Map existingColumnStatistics; + if (mode == StatisticsUpdateMode.MERGE_INCREMENTAL) { + existingColumnStatistics = getTableColumnStatistics(databaseName, tableName, statisticsUpdate.getColumnStatistics().keySet()); + } + else { + existingColumnStatistics = ImmutableMap.of(); + } + + // first update the basic statistics on the table, which requires a read-modify-write cycle + BasicTableStatisticsResult result; + try { + result = updateBasicTableStatistics(databaseName, tableName, mode, statisticsUpdate, existingColumnStatistics); } - catch (AmazonServiceException e) { + catch (SdkException e) { throw new TrinoException(HIVE_METASTORE_ERROR, e); } + finally { + glueCache.invalidateTable(databaseName, tableName, true); + } + + List> tasks = new ArrayList<>(); + Lists.partition(toGlueColumnStatistics(result.updateColumnStatistics()), GLUE_COLUMN_WRITE_STAT_PAGE_SIZE).stream() + .map(chunk -> (Callable) () -> { + stats.getUpdateColumnStatisticsForTable().call(() -> glueClient.updateColumnStatisticsForTable(builder -> builder + .applyMutation(glueContext::configureClient) + .databaseName(databaseName) + .tableName(tableName) + .columnStatisticsList(chunk))); + return null; + }) + .forEach(tasks::add); + + result.removeColumnStatistics().stream() + .map(columnName -> (Callable) () -> { + stats.getDeleteColumnStatisticsForTable().call(() -> glueClient.deleteColumnStatisticsForTable(builder -> builder + .applyMutation(glueContext::configureClient) + .databaseName(databaseName) + .tableName(tableName) + .columnName(columnName))); + return null; + }) + .forEach(tasks::add); + try { + runParallel(tasks); + } + catch (ExecutionException e) { + if (e.getCause() instanceof EntityNotFoundException) { + throw new TableNotFoundException(new SchemaTableName(databaseName, tableName)); + } + throw new TrinoException(HIVE_METASTORE_ERROR, "Failed updating column statistics for %s.%s".formatted(databaseName, tableName), e.getCause()); + } + finally { + glueCache.invalidateTableColumnStatistics(databaseName, tableName); + } } - @Override - public void commentColumn(String databaseName, String tableName, String columnName, Optional comment) + private BasicTableStatisticsResult updateBasicTableStatistics( + String databaseName, + String tableName, + StatisticsUpdateMode mode, + PartitionStatistics statisticsUpdate, + Map existingColumnStatistics) { - throw new TrinoException(NOT_SUPPORTED, "Column comment is not yet supported by Glue service"); + Table table = getTable(databaseName, tableName) + .orElseThrow(() -> new TableNotFoundException(new SchemaTableName(databaseName, tableName))); + + PartitionStatistics existingTableStats = new PartitionStatistics(getHiveBasicStatistics(table.getParameters()), existingColumnStatistics); + PartitionStatistics updatedStatistics = mode.updatePartitionStatistics(existingTableStats, statisticsUpdate); + + stats.getUpdateTable().call(() -> glueClient.updateTable(builder -> builder + .applyMutation(glueContext::configureClient) + .databaseName(databaseName) + .tableInput(toGlueTableInput(table.withParameters(updateStatisticsParameters(table.getParameters(), updatedStatistics.getBasicStatistics())))))); + + Map columns = Stream.concat(table.getDataColumns().stream(), table.getPartitionColumns().stream()) + .collect(toImmutableMap(Column::getName, identity())); + + Map updateColumnStatistics = columns.values().stream() + .filter(column -> updatedStatistics.getColumnStatistics().containsKey(column.getName())) + .collect(toImmutableMap(identity(), column -> updatedStatistics.getColumnStatistics().get(column.getName()))); + + Set removeColumnStatistics = switch (mode) { + case OVERWRITE_SOME_COLUMNS -> ImmutableSet.of(); + case OVERWRITE_ALL, MERGE_INCREMENTAL -> columns.entrySet().stream() + // all columns not being updated are removed + .filter(entry -> !updateColumnStatistics.containsKey(entry.getValue())) + .map(Entry::getKey) + .collect(Collectors.toSet()); + case UNDO_MERGE_INCREMENTAL, CLEAR_ALL -> columns.keySet(); + }; + + return new BasicTableStatisticsResult(updateColumnStatistics, removeColumnStatistics); } + private record BasicTableStatisticsResult(Map updateColumnStatistics, Set removeColumnStatistics) {} + @Override public void addColumn(String databaseName, String tableName, String columnName, HiveType columnType, String columnComment) { - Table oldTable = getExistingTable(databaseName, tableName); - Table newTable = Table.builder(oldTable) - .addDataColumn(new Column(columnName, columnType, Optional.ofNullable(columnComment))) - .build(); - replaceTable(databaseName, tableName, newTable, null); + Column newColumn = new Column(columnName, columnType, Optional.ofNullable(columnComment), ImmutableMap.of()); + updateTable(databaseName, tableName, table -> table.withAddColumn(newColumn)); } @Override public void renameColumn(String databaseName, String tableName, String oldColumnName, String newColumnName) { - Table oldTable = getExistingTable(databaseName, tableName); - if (oldTable.getPartitionColumns().stream().anyMatch(c -> c.getName().equals(oldColumnName))) { - throw new TrinoException(NOT_SUPPORTED, "Renaming partition columns is not supported"); - } - - ImmutableList.Builder newDataColumns = ImmutableList.builder(); - for (Column column : oldTable.getDataColumns()) { - if (column.getName().equals(oldColumnName)) { - newDataColumns.add(new Column(newColumnName, column.getType(), column.getComment())); - } - else { - newDataColumns.add(column); - } - } - - Table newTable = Table.builder(oldTable) - .setDataColumns(newDataColumns.build()) - .build(); - replaceTable(databaseName, tableName, newTable, null); + updateTable(databaseName, tableName, table -> table.withRenameColumn(oldColumnName, newColumnName)); } @Override public void dropColumn(String databaseName, String tableName, String columnName) + throws EntityNotFoundException { - verifyCanDropColumn(this, databaseName, tableName, columnName); - Table oldTable = getExistingTable(databaseName, tableName); - - if (oldTable.getColumn(columnName).isEmpty()) { - SchemaTableName name = new SchemaTableName(databaseName, tableName); - throw new ColumnNotFoundException(name, columnName); - } - - ImmutableList.Builder newDataColumns = ImmutableList.builder(); - oldTable.getDataColumns().stream() - .filter(fieldSchema -> !fieldSchema.getName().equals(columnName)) - .forEach(newDataColumns::add); - - Table newTable = Table.builder(oldTable) - .setDataColumns(newDataColumns.build()) - .build(); - replaceTable(databaseName, tableName, newTable, null); + updateTable(databaseName, tableName, table -> table.withDropColumn(columnName)); } @Override - public Optional getPartition(Table table, List partitionValues) + public void commentColumn(String databaseName, String tableName, String columnName, Optional comment) + throws EntityNotFoundException { - try { - GetPartitionResult result = stats.getGetPartition().call(() -> - glueClient.getPartition(new GetPartitionRequest() - .withDatabaseName(table.getDatabaseName()) - .withTableName(table.getTableName()) - .withPartitionValues(partitionValues))); - return Optional.of(new GluePartitionConverter(table).apply(result.getPartition())); - } - catch (EntityNotFoundException e) { - return Optional.empty(); - } - catch (AmazonServiceException e) { - throw new TrinoException(HIVE_METASTORE_ERROR, e); - } + updateTable(databaseName, tableName, table -> table.withColumnComment(columnName, comment)); } @Override - public Optional> getPartitionNamesByFilter( - String databaseName, - String tableName, - List columnNames, - TupleDomain partitionKeysFilter) + public Optional> getPartitionNamesByFilter(String databaseName, String tableName, List columnNames, TupleDomain partitionKeysFilter) { if (partitionKeysFilter.isNone()) { return Optional.of(ImmutableList.of()); } - String expression = GlueExpressionUtil.buildGlueExpression(columnNames, partitionKeysFilter, assumeCanonicalPartitionKeys); - List> partitionValues = getPartitionValues(databaseName, tableName, expression); - return Optional.of(buildPartitionNames(columnNames, partitionValues)); + String glueFilterExpression = GlueExpressionUtil.buildGlueExpression(columnNames, partitionKeysFilter, assumeCanonicalPartitionKeys); + Set partitionNames = glueCache.getPartitionNames( + databaseName, + tableName, + glueFilterExpression, + cachePartition -> getPartitionNames(cachePartition, databaseName, tableName, glueFilterExpression)); + + return Optional.of(partitionNames.stream() + .map(PartitionName::partitionValues) + .map(partitionValues -> toPartitionName(columnNames, partitionValues)) + .collect(toImmutableList())); } - private List> getPartitionValues(String databaseName, String tableName, String expression) + private Set getPartitionNames(Consumer cachePartition, String databaseName, String tableName, String glueExpression) + throws EntityNotFoundException { if (partitionSegments == 1) { - return getPartitionValues(databaseName, tableName, expression, null); + return getPartitionNames(cachePartition, databaseName, tableName, glueExpression, null); } - // Do parallel partition fetch. - CompletionService>> completionService = new ExecutorCompletionService<>(partitionsReadExecutor); - List> futures = new ArrayList<>(partitionSegments); - List> partitions = new ArrayList<>(); try { - for (int i = 0; i < partitionSegments; i++) { - Segment segment = new Segment().withSegmentNumber(i).withTotalSegments(partitionSegments); - futures.add(completionService.submit(() -> getPartitionValues(databaseName, tableName, expression, segment))); - } - for (int i = 0; i < partitionSegments; i++) { - Future>> futurePartitions = completionService.take(); - partitions.addAll(futurePartitions.get()); - } - // All futures completed normally - futures.clear(); - } - catch (ExecutionException | InterruptedException e) { - if (e instanceof InterruptedException) { - Thread.currentThread().interrupt(); + return runParallel(IntStream.range(0, partitionSegments) + .mapToObj(segmentNumber -> (Callable>) () -> getPartitionNames( + cachePartition, + databaseName, + tableName, + glueExpression, + Segment.builder() + .segmentNumber(segmentNumber) + .totalSegments(partitionSegments) + .build())) + .toList()) + .stream() + .flatMap(Collection::stream) + .collect(toImmutableSet()); + } + catch (ExecutionException e) { + if (e.getCause() instanceof EntityNotFoundException) { + throw new TableNotFoundException(new SchemaTableName(databaseName, tableName)); } - throw new TrinoException(HIVE_METASTORE_ERROR, "Failed to fetch partitions from Glue Data Catalog", e); - } - finally { - // Ensure any futures still running are canceled in case of failure - futures.forEach(future -> future.cancel(true)); + throw new TrinoException(HIVE_METASTORE_ERROR, e.getCause()); } - - partitions.sort(PARTITION_VALUE_COMPARATOR); - return partitions; } - private List> getPartitionValues(String databaseName, String tableName, String expression, @Nullable Segment segment) + private Set getPartitionNames(Consumer cachePartition, String databaseName, String tableName, String expression, @Nullable Segment segment) + throws EntityNotFoundException { try { - // Reuse immutable field instances opportunistically between partitions - return getPaginatedResults( - glueClient::getPartitions, - new GetPartitionsRequest() - .withDatabaseName(databaseName) - .withTableName(tableName) - .withExpression(expression) - .withSegment(segment) - // We are interested in the partition values and excluding column schema - // avoids the problem of a large response. - .withExcludeColumnSchema(true) - .withMaxResults(AWS_GLUE_GET_PARTITIONS_MAX_RESULTS), - GetPartitionsRequest::setNextToken, - GetPartitionsResult::getNextToken, - stats.getGetPartitions()) - .map(GetPartitionsResult::getPartitions) + return stats.getGetPartitionNames().call(() -> glueClient.getPartitionsPaginator(builder -> builder + .applyMutation(glueContext::configureClient) + .databaseName(databaseName) + .tableName(tableName) + .expression(expression) + .segment(segment) + .maxResults(AWS_GLUE_GET_PARTITIONS_MAX_RESULTS)).stream() + .map(GetPartitionsResponse::partitions) .flatMap(List::stream) - .map(com.amazonaws.services.glue.model.Partition::getValues) - .collect(toImmutableList()); + .map(partition -> GlueConverter.fromGluePartition(databaseName, tableName, partition)) + .peek(cachePartition) + .map(Partition::getValues) + .map(PartitionName::new) + .collect(toImmutableSet())); + } + catch (EntityNotFoundException e) { + throw new TableNotFoundException(new SchemaTableName(databaseName, tableName)); } - catch (AmazonServiceException e) { + catch (SdkException e) { throw new TrinoException(HIVE_METASTORE_ERROR, e); } } - private static List buildPartitionNames(List partitionColumns, List> partitions) + @Override + public Optional getPartition(Table table, List partitionValues) { - return mappedCopy(partitions, partition -> toPartitionName(partitionColumns, partition)); + String databaseName = table.getDatabaseName(); + String tableName = table.getTableName(); + PartitionName partitionName = new PartitionName(partitionValues); + return glueCache.getPartition(databaseName, tableName, partitionName, () -> getPartition(databaseName, tableName, partitionName)); + } + + private Optional getPartition(String databaseName, String tableName, PartitionName partitionName) + { + try { + GetPartitionResponse result = stats.getGetPartition().call(() -> glueClient.getPartition(builder -> builder + .applyMutation(glueContext::configureClient) + .databaseName(databaseName) + .tableName(tableName) + .partitionValues(partitionName.partitionValues()))); + return Optional.of(GlueConverter.fromGluePartition(databaseName, tableName, result.partition())); + } + catch (EntityNotFoundException e) { + return Optional.empty(); + } + catch (SdkException e) { + throw new TrinoException(HIVE_METASTORE_ERROR, e); + } } - /** - * 
-     * Ex: Partition keys = ['a', 'b']
-     *     Partition names = ['a=1/b=2', 'a=2/b=2']
-     *
- * - * @param partitionNames List of full partition names - * @return Mapping of partition name to partition object - */ @Override public Map> getPartitionsByNames(Table table, List partitionNames) { - return stats.getGetPartitionByName().call(() -> getPartitionsByNamesInternal(table, partitionNames)); + List names = partitionNames.stream() + .map(HivePartitionManager::extractPartitionValues) + .map(PartitionName::new) + .collect(toImmutableList()); + return getPartitionsByNames(table.getDatabaseName(), table.getTableName(), names).entrySet().stream() + .collect(toImmutableMap(entry -> makePartitionName(table.getPartitionColumns(), entry.getKey().partitionValues()), Entry::getValue)); } - private Map> getPartitionsByNamesInternal(Table table, List partitionNames) + private Map> getPartitionsByNames(String databaseName, String tableName, Collection partitionNames) { - requireNonNull(partitionNames, "partitionNames is null"); - if (partitionNames.isEmpty()) { - return ImmutableMap.of(); - } - - List partitions = batchGetPartition(table, partitionNames); - - Map> partitionNameToPartitionValuesMap = partitionNames.stream() - .collect(toMap(identity(), HiveUtil::toPartitionValues)); - Map, Partition> partitionValuesToPartitionMap = partitions.stream() - .collect(toMap(Partition::getValues, identity())); - - ImmutableMap.Builder> resultBuilder = ImmutableMap.builder(); - for (Entry> entry : partitionNameToPartitionValuesMap.entrySet()) { - Partition partition = partitionValuesToPartitionMap.get(entry.getValue()); - resultBuilder.put(entry.getKey(), Optional.ofNullable(partition)); + Collection partitions = glueCache.batchGetPartitions( + databaseName, + tableName, + partitionNames, + (cachePartition, missingPartitions) -> batchGetPartition(databaseName, tableName, missingPartitions, cachePartition)); + Map partitionValuesToPartitionMap = partitions.stream() + .collect(toMap(partition -> new PartitionName(partition.getValues()), identity())); + + ImmutableMap.Builder> resultBuilder = ImmutableMap.builder(); + for (PartitionName partitionName : partitionNames) { + Partition partition = partitionValuesToPartitionMap.get(partitionName); + resultBuilder.put(partitionName, Optional.ofNullable(partition)); } return resultBuilder.buildOrThrow(); } - private List batchGetPartition(Table table, List partitionNames) + private Collection batchGetPartition(String databaseName, String tableName, Collection partitionNames, Consumer cachePartition) + throws EntityNotFoundException { - List> batchGetPartitionFutures = new ArrayList<>(); try { List pendingPartitions = partitionNames.stream() - .map(partitionName -> new PartitionValueList().withValues(toPartitionValues(partitionName))) + .map(partitionName -> PartitionValueList.builder().values(partitionName.partitionValues()).build()) .collect(toCollection(ArrayList::new)); ImmutableList.Builder resultsBuilder = ImmutableList.builderWithExpectedSize(partitionNames.size()); - - // Reuse immutable field instances opportunistically between partitions - GluePartitionConverter converter = new GluePartitionConverter(table); - while (!pendingPartitions.isEmpty()) { - for (List partitions : Lists.partition(pendingPartitions, BATCH_GET_PARTITION_MAX_PAGE_SIZE)) { - batchGetPartitionFutures.add(glueClient.batchGetPartitionAsync(new BatchGetPartitionRequest() - .withDatabaseName(table.getDatabaseName()) - .withTableName(table.getTableName()) - .withPartitionsToGet(partitions), - new StatsRecordingAsyncHandler<>(stats.getGetPartitions()))); - } + List responses = runParallel(Lists.partition(pendingPartitions, BATCH_GET_PARTITION_MAX_PAGE_SIZE).stream() + .map(partitions -> (Callable) () -> stats.getGetPartitions().call(() -> glueClient.batchGetPartition(builder -> builder + .applyMutation(glueContext::configureClient) + .databaseName(databaseName) + .tableName(tableName) + .partitionsToGet(partitions)))) + .toList()); pendingPartitions.clear(); - for (Future future : batchGetPartitionFutures) { - BatchGetPartitionResult batchGetPartitionResult = future.get(); - List partitions = batchGetPartitionResult.getPartitions(); - List unprocessedKeys = batchGetPartitionResult.getUnprocessedKeys(); - - // In the unlikely scenario where batchGetPartition call cannot make progress on retrieving partitions, avoid infinite loop - if (partitions.isEmpty()) { - verify(!unprocessedKeys.isEmpty(), "Empty unprocessedKeys for non-empty BatchGetPartitionRequest and empty partitions result"); - throw new TrinoException(HIVE_METASTORE_ERROR, "Cannot make progress retrieving partitions. Unable to retrieve partitions: " + unprocessedKeys); + for (var response : responses) { + // In the unlikely scenario where batchGetPartition call cannot make progress on retrieving partitions, avoid infinite loop. + // We fail only in case there are still unprocessedKeys. Case with empty partitions and empty unprocessedKeys is correct in case partitions from request are not found. + if (response.partitions().isEmpty() && !response.unprocessedKeys().isEmpty()) { + throw new TrinoException(HIVE_METASTORE_ERROR, "Cannot make progress retrieving partitions. Unable to retrieve partitions: " + response.unprocessedKeys()); } - partitions.stream() - .map(converter) + response.partitions().stream() + .map(partition -> GlueConverter.fromGluePartition(databaseName, tableName, partition)) + .peek(cachePartition) .forEach(resultsBuilder::add); - pendingPartitions.addAll(unprocessedKeys); + pendingPartitions.addAll(response.unprocessedKeys()); } - batchGetPartitionFutures.clear(); } - return resultsBuilder.build(); } - catch (AmazonServiceException | InterruptedException | ExecutionException e) { - if (e instanceof InterruptedException) { - Thread.currentThread().interrupt(); + catch (ExecutionException e) { + if (e.getCause() instanceof EntityNotFoundException) { + throw new TableNotFoundException(new SchemaTableName(databaseName, tableName)); } - throw new TrinoException(HIVE_METASTORE_ERROR, e); - } - finally { - // Ensure any futures still running are canceled in case of failure - batchGetPartitionFutures.forEach(future -> future.cancel(true)); + throw new TrinoException(HIVE_METASTORE_ERROR, "Failed fetching partitions for %s.%s".formatted(databaseName, tableName), e.getCause()); } } @Override - public void addPartitions(String databaseName, String tableName, List partitions) + public void addPartitions(String databaseName, String tableName, List partitionsWithStatistics) { - try { - stats.getCreatePartitions().call(() -> { - List> futures = new ArrayList<>(); - - for (List partitionBatch : Lists.partition(partitions, BATCH_CREATE_PARTITION_MAX_PAGE_SIZE)) { - List partitionInputs = mappedCopy(partitionBatch, GlueInputConverter::convertPartition); - futures.add(glueClient.batchCreatePartitionAsync( - new BatchCreatePartitionRequest() - .withDatabaseName(databaseName) - .withTableName(tableName) - .withPartitionInputList(partitionInputs), - new StatsRecordingAsyncHandler<>(stats.getBatchCreatePartition()))); - } - - for (Future future : futures) { - try { - BatchCreatePartitionResult result = future.get(); - propagatePartitionErrorToTrinoException(databaseName, tableName, result.getErrors()); - } - catch (InterruptedException e) { - Thread.currentThread().interrupt(); - throw new TrinoException(HIVE_METASTORE_ERROR, e); + if (partitionsWithStatistics.isEmpty()) { + return; + } + partitionsWithStatistics.stream() + .map(PartitionWithStatistics::getPartition) + .collect(toImmutableList()) + .forEach(partition -> { + if (!partition.getDatabaseName().equals(databaseName) || + !partition.getTableName().equals(tableName)) { + throw new TrinoException(NOT_SUPPORTED, "All partitions must belong to the same table"); } - } - - Set updates = partitions.stream() - .map(partitionWithStatistics -> new GlueColumnStatisticsProvider.PartitionStatisticsUpdate( - partitionWithStatistics.getPartition(), - partitionWithStatistics.getStatistics().getColumnStatistics())) - .collect(toImmutableSet()); - columnStatisticsProvider.updatePartitionStatistics(updates); + }); + + List updatedPartitions = partitionsWithStatistics.stream() + .map(partitionWithStatistics -> { + Partition partition = partitionWithStatistics.getPartition(); + HiveBasicStatistics basicStatistics = partitionWithStatistics.getStatistics().getBasicStatistics(); + return partition.withParameters(updateStatisticsParameters(partition.getParameters(), basicStatistics)); + }) + .toList(); + var createPartitionTasks = Lists.partition(updatedPartitions, BATCH_UPDATE_PARTITION_MAX_PAGE_SIZE).stream() + .map(partitionBatch -> (Callable) () -> { + stats.getCreatePartitions().call(() -> glueClient.batchCreatePartition(builder -> builder + .applyMutation(glueContext::configureClient) + .databaseName(databaseName) + .tableName(tableName) + .partitionInputList(partitionBatch.stream() + .map(GlueConverter::toGluePartitionInput) + .collect(toImmutableList())))); + return null; + }) + .toList(); + try { + runParallel(createPartitionTasks); + } + catch (ExecutionException e) { + if (e.getCause() instanceof AlreadyExistsException) { + throw new TrinoException(ALREADY_EXISTS, "Partition already exists in table %s.%s: %s".formatted(databaseName, tableName, e.getCause().getMessage())); + } + if (e.getCause() instanceof EntityNotFoundException) { + throw new TableNotFoundException(new SchemaTableName(databaseName, tableName)); + } + throw new TrinoException(HIVE_METASTORE_ERROR, "Failed creating partitions for %s.%s: %s".formatted(databaseName, tableName, e.getCause().getMessage())); + } - return null; - }); + // statistics are created after partitions because it is not clear if ordering matters in Glue + var createStatisticsTasks = partitionsWithStatistics.stream() + .map(partitionWithStatistics -> createUpdatePartitionStatisticsTasks( + StatisticsUpdateMode.OVERWRITE_ALL, + partitionWithStatistics.getPartition(), + partitionWithStatistics.getStatistics().getColumnStatistics())) + .flatMap(Collection::stream) + .collect(toImmutableList()); + try { + runParallel(createStatisticsTasks); } - catch (AmazonServiceException | ExecutionException e) { - throw new TrinoException(HIVE_METASTORE_ERROR, e); + catch (ExecutionException e) { + throw new TrinoException(HIVE_METASTORE_ERROR, "Failed creating partition column statistics for %s.%s".formatted(databaseName, tableName), e.getCause()); } } - private static void propagatePartitionErrorToTrinoException(String databaseName, String tableName, List partitionErrors) + @Override + public void dropPartition(String databaseName, String tableName, List partitionValues, boolean deleteData) { - if (partitionErrors != null && !partitionErrors.isEmpty()) { - ErrorDetail errorDetail = partitionErrors.get(0).getErrorDetail(); - String glueExceptionCode = errorDetail.getErrorCode(); - - switch (glueExceptionCode) { - case "AlreadyExistsException": - throw new TrinoException(ALREADY_EXISTS, errorDetail.getErrorMessage()); - case "EntityNotFoundException": - throw new TableNotFoundException(new SchemaTableName(databaseName, tableName), errorDetail.getErrorMessage()); - default: - throw new TrinoException(HIVE_METASTORE_ERROR, errorDetail.getErrorCode() + ": " + errorDetail.getErrorMessage()); - } + PartitionName partitionName = new PartitionName(partitionValues); + Optional location = Optional.empty(); + if (deleteData) { + Partition partition = getPartition(databaseName, tableName, partitionName) + .orElseThrow(() -> new PartitionNotFoundException(new SchemaTableName(databaseName, tableName), partitionValues)); + location = Optional.of(Strings.nullToEmpty(partition.getStorage().getLocation())) + .map(Location::of); + } + try { + stats.getDeletePartition().call(() -> glueClient.deletePartition(builder -> builder + .applyMutation(glueContext::configureClient) + .databaseName(databaseName) + .tableName(tableName) + .partitionValues(partitionValues))); + } + catch (EntityNotFoundException e) { + throw new PartitionNotFoundException(new SchemaTableName(databaseName, tableName), partitionValues); + } + catch (SdkException e) { + throw new TrinoException(HIVE_METASTORE_ERROR, e); + } + finally { + glueCache.invalidatePartition(databaseName, tableName, partitionName); } + + location.ifPresent(this::deleteDir); } @Override - public void dropPartition(String databaseName, String tableName, List parts, boolean deleteData) + public void alterPartition(String databaseName, String tableName, PartitionWithStatistics partitionWithStatistics) { - Table table = getExistingTable(databaseName, tableName); - Partition partition = getPartition(table, parts) - .orElseThrow(() -> new PartitionNotFoundException(new SchemaTableName(databaseName, tableName), parts)); + if (!partitionWithStatistics.getPartition().getDatabaseName().equals(databaseName) || !partitionWithStatistics.getPartition().getTableName().equals(tableName)) { + throw new TrinoException(HIVE_METASTORE_ERROR, "Partition names must match table name"); + } + + alterPartition(partitionWithStatistics); + } + private void alterPartition(PartitionWithStatistics partitionWithStatistics) + { + Partition partition = partitionWithStatistics.getPartition(); try { - stats.getDeletePartition().call(() -> - glueClient.deletePartition(new DeletePartitionRequest() - .withDatabaseName(databaseName) - .withTableName(tableName) - .withPartitionValues(parts))); + HiveBasicStatistics basicStatistics = partitionWithStatistics.getStatistics().getBasicStatistics(); + PartitionInput newPartition = toGluePartitionInput(partition.withParameters(updateStatisticsParameters(partition.getParameters(), basicStatistics))); + stats.getUpdatePartition().call(() -> glueClient.updatePartition(builder -> builder + .applyMutation(glueContext::configureClient) + .databaseName(partition.getDatabaseName()) + .tableName(partition.getTableName()) + .partitionInput(newPartition) + .partitionValueList(partition.getValues()))); + } + catch (EntityNotFoundException e) { + throw new PartitionNotFoundException(partition.getSchemaTableName(), partition.getValues()); } - catch (AmazonServiceException e) { + catch (SdkException e) { throw new TrinoException(HIVE_METASTORE_ERROR, e); } + finally { + glueCache.invalidatePartition(partition.getDatabaseName(), partition.getTableName(), new PartitionName(partition.getValues())); + } - String partLocation = partition.getStorage().getLocation(); - if (deleteData && isManagedTable(table) && !isNullOrEmpty(partLocation)) { - deleteDir(hdfsContext, hdfsEnvironment, new Path(partLocation), true); + try { + runParallel(createUpdatePartitionStatisticsTasks(StatisticsUpdateMode.OVERWRITE_ALL, partition, partitionWithStatistics.getStatistics().getColumnStatistics())); + } + catch (ExecutionException e) { + if (e.getCause() instanceof EntityNotFoundException) { + throw new PartitionNotFoundException(new SchemaTableName(partition.getDatabaseName(), partition.getTableName()), partition.getValues()); + } + throw new TrinoException(HIVE_METASTORE_ERROR, "Failed updating partition column statistics for %s %s".formatted(partition.getSchemaTableName(), partition.getValues()), e.getCause()); + } + finally { + // this redundant, but be safe in case stats were cached before the update + glueCache.invalidatePartition(partition.getDatabaseName(), partition.getTableName(), new PartitionName(partition.getValues())); } } @Override - public void alterPartition(String databaseName, String tableName, PartitionWithStatistics partition) + public Map> getPartitionColumnStatistics( + String databaseName, + String tableName, + Set partitionNames, + Set columnNames) + { + checkArgument(!columnNames.isEmpty(), "columnNames is empty"); + Map partitionsNameMap = partitionNames.stream() + .collect(toImmutableMap(partitionName -> new PartitionName(HivePartitionManager.extractPartitionValues(partitionName)), identity())); + return getPartitionColumnStatistics(databaseName, tableName, partitionsNameMap.keySet(), columnNames).entrySet().stream() + .collect(toImmutableMap(entry -> partitionsNameMap.get(entry.getKey()), Entry::getValue)); + } + + private Map> getPartitionColumnStatistics(String databaseName, String tableName, Collection partitionNames, Set columnNames) + throws EntityNotFoundException + { + Map> columnStats; + try { + columnStats = runParallel(partitionNames.stream() + .map(partitionName -> (Callable>>) () -> entry( + partitionName, + glueCache.getPartitionColumnStatistics( + databaseName, + tableName, + partitionName, + columnNames, + missingColumnNames -> getPartitionColumnStatisticsInternal(databaseName, tableName, partitionName, missingColumnNames)))) + .collect(toImmutableList())) + .stream() + .collect(toImmutableMap(Entry::getKey, Entry::getValue)); + } + catch (ExecutionException e) { + throwIfUnchecked(e.getCause()); + throw new TrinoException(HIVE_METASTORE_ERROR, "Failed fetching partition statistics for %s.%s".formatted(databaseName, tableName), e.getCause()); + } + + // use empty statistics for missing partitions + return partitionNames.stream() + .collect(toImmutableMap(identity(), name -> columnStats.getOrDefault(name, ImmutableMap.of()))); + } + + private Map getPartitionColumnStatisticsInternal(String databaseName, String tableName, PartitionName partitionName, Set columnNames) + throws EntityNotFoundException { + var columnStatsTasks = Lists.partition(ImmutableList.copyOf(columnNames), GLUE_COLUMN_READ_STAT_PAGE_SIZE).stream() + .map(partialColumns -> (Callable>) () -> stats.getGetColumnStatisticsForPartition().call(() -> glueClient.getColumnStatisticsForPartition(builder -> builder + .applyMutation(glueContext::configureClient) + .databaseName(databaseName) + .tableName(tableName) + .partitionValues(partitionName.partitionValues()) + .columnNames(partialColumns)) + .columnStatisticsList())) + .collect(toImmutableList()); + try { - PartitionInput newPartition = convertPartition(partition); - stats.getUpdatePartition().call(() -> - glueClient.updatePartition(new UpdatePartitionRequest() - .withDatabaseName(databaseName) - .withTableName(tableName) - .withPartitionInput(newPartition) - .withPartitionValueList(partition.getPartition().getValues()))); - columnStatisticsProvider.updatePartitionStatistics( - partition.getPartition(), - partition.getStatistics().getColumnStatistics()); + var glueColumnStatistics = runParallel(columnStatsTasks); + return fromGlueStatistics(glueColumnStatistics); } - catch (EntityNotFoundException e) { - throw new PartitionNotFoundException(new SchemaTableName(databaseName, tableName), partition.getPartition().getValues()); + catch (ExecutionException e) { + if (e.getCause() instanceof EntityNotFoundException) { + return ImmutableMap.of(); + } + throw new TrinoException(HIVE_METASTORE_ERROR, "Failed fetching column statistics for %s.%s %s".formatted(databaseName, tableName, partitionName), e.getCause()); } - catch (AmazonServiceException e) { - throw new TrinoException(HIVE_METASTORE_ERROR, e); + } + + @Override + public void updatePartitionStatistics(Table table, StatisticsUpdateMode mode, Map partitionUpdates) + { + batchUpdatePartitionStatistics(table.getDatabaseName(), table.getTableName(), mode, partitionUpdates.entrySet().stream() + // filter out empty statistics for merge incremental mode + .filter(entry -> mode != StatisticsUpdateMode.MERGE_INCREMENTAL || !entry.getValue().getBasicStatistics().getRowCount().equals(OptionalLong.of(0))) + .collect(toImmutableMap(entry -> new PartitionName(HivePartitionManager.extractPartitionValues(entry.getKey())), Entry::getValue))); + } + + private void batchUpdatePartitionStatistics(String databaseName, String tableName, StatisticsUpdateMode mode, Map newStatistics) + throws EntityNotFoundException + { + // partitions are required for update + // the glue cache is not used here because we need up-to-date partitions for the read-update-write operation + // loaded partitions are not cached here since the partition values are invalidated after the update + Collection partitions = batchGetPartition(databaseName, tableName, ImmutableList.copyOf(newStatistics.keySet()), ignore -> {}); + + // existing column statistics are required for the columns being updated when in merge incremental mode + // this is fetched before the basic statistics are updated, to avoid reloading in the case of retries updating basic statistics + Map> existingColumnStatistics; + if (mode == StatisticsUpdateMode.MERGE_INCREMENTAL) { + try { + existingColumnStatistics = runParallel(newStatistics.entrySet().stream() + .map(entry -> (Callable>>) () -> { + PartitionName partitionName = entry.getKey(); + Map columnStatistics = getPartitionColumnStatisticsInternal(databaseName, tableName, partitionName, entry.getValue().getColumnStatistics().keySet()); + return entry(partitionName, columnStatistics); + }) + .collect(toImmutableList())) + .stream() + .collect(toImmutableMap(Entry::getKey, Entry::getValue)); + } + catch (ExecutionException e) { + throwIfUnchecked(e.getCause()); + throw new TrinoException(HIVE_METASTORE_ERROR, "Failed fetching current partition statistics for %s.%s".formatted(databaseName, tableName), e.getCause()); + } + } + else { + existingColumnStatistics = ImmutableMap.of(); + } + + List partitionStatisticsUpdates = new ArrayList<>(); + for (Partition partition : partitions) { + PartitionName partitionName = new PartitionName(partition.getValues()); + PartitionStatistics existingStats = new PartitionStatistics( + getHiveBasicStatistics(partition.getParameters()), + existingColumnStatistics.getOrDefault(partitionName, ImmutableMap.of())); + PartitionStatistics newStats = newStatistics.get(partitionName); + PartitionStatistics updatedStatistics = mode.updatePartitionStatistics(existingStats, newStats); + Partition updatedPartition = partition.withParameters(updateStatisticsParameters(partition.getParameters(), updatedStatistics.getBasicStatistics())); + boolean basicStatsUpdated = !updatedStatistics.getBasicStatistics().equals(existingStats.getBasicStatistics()); + partitionStatisticsUpdates.add(new PartitionStatisticsUpdate(updatedPartition, basicStatsUpdated, updatedStatistics.getColumnStatistics())); + } + + // partitions and statistics are updated in concurrently + List> tasks = new ArrayList<>(); + partitionStatisticsUpdates.stream() + .map(update -> createUpdatePartitionStatisticsTasks(mode, update.updatedPartition(), update.updatedStatistics())) + .flatMap(Collection::stream) + .forEach(tasks::add); + + List updatedPartitions = partitionStatisticsUpdates.stream() + .filter(PartitionStatisticsUpdate::basicStatsUpdated) + .map(PartitionStatisticsUpdate::updatedPartition) + .collect(toImmutableList()); + Lists.partition(updatedPartitions, BATCH_UPDATE_PARTITION_MAX_PAGE_SIZE).stream() + .map(partitionBatch -> (Callable) () -> { + stats.getBatchUpdatePartition().call(() -> glueClient.batchUpdatePartition(builder -> builder + .applyMutation(glueContext::configureClient) + .databaseName(databaseName) + .tableName(tableName) + .entries(partitionBatch.stream() + .map(partition -> BatchUpdatePartitionRequestEntry.builder() + .partitionValueList(partition.getValues()) + .partitionInput(toGluePartitionInput(partition)) + .build()) + .collect(toImmutableList())))); + return null; + }) + .forEach(tasks::add); + + try { + runParallel(tasks); } + catch (ExecutionException e) { + throw new TrinoException(HIVE_METASTORE_ERROR, "Failed updating partition column statistics for %s.%s".formatted(databaseName, tableName), e.getCause()); + } + } + + private List> createUpdatePartitionStatisticsTasks(StatisticsUpdateMode mode, Partition partition, Map updatedStatistics) + { + Map columns = partition.getColumns().stream() + .collect(toImmutableMap(Column::getName, identity())); + + Map statisticsByColumn = columns.values().stream() + .filter(column -> updatedStatistics.containsKey(column.getName())) + .collect(toImmutableMap(identity(), column -> updatedStatistics.get(column.getName()))); + + Set removeColumnStatistics = switch (mode) { + case OVERWRITE_SOME_COLUMNS -> ImmutableSet.of(); + case OVERWRITE_ALL, MERGE_INCREMENTAL -> columns.entrySet().stream() + // all columns not being updated are removed + .filter(entry -> !statisticsByColumn.containsKey(entry.getValue())) + .map(Entry::getKey) + .collect(Collectors.toSet()); + case UNDO_MERGE_INCREMENTAL, CLEAR_ALL -> columns.keySet(); + }; + + ImmutableList.Builder> tasks = ImmutableList.builder(); + Lists.partition(toGlueColumnStatistics(statisticsByColumn), GLUE_COLUMN_WRITE_STAT_PAGE_SIZE).stream() + .map(chunk -> (Callable) () -> { + stats.getUpdateColumnStatisticsForPartition().call(() -> glueClient.updateColumnStatisticsForPartition(builder -> builder + .applyMutation(glueContext::configureClient) + .databaseName(partition.getDatabaseName()) + .tableName(partition.getTableName()) + .partitionValues(partition.getValues()) + .columnStatisticsList(chunk))); + return null; + }) + .forEach(tasks::add); + removeColumnStatistics.stream() + .map(columnName -> (Callable) () -> { + stats.getDeleteColumnStatisticsForPartition().call(() -> glueClient.deleteColumnStatisticsForPartition(builder -> builder + .applyMutation(glueContext::configureClient) + .databaseName(partition.getDatabaseName()) + .tableName(partition.getTableName()) + .partitionValues(partition.getValues()) + .columnName(columnName))); + return null; + }) + .forEach(tasks::add); + return tasks.build(); } + private record PartitionStatisticsUpdate(Partition updatedPartition, boolean basicStatsUpdated, Map updatedStatistics) {} + @Override public void createRole(String role, String grantor) { @@ -1116,12 +1444,6 @@ public void revokeRoles(Set roles, Set grantees, boolean throw new TrinoException(NOT_SUPPORTED, "revokeRoles is not supported by Glue"); } - @Override - public Set listGrantedPrincipals(String role) - { - throw new TrinoException(NOT_SUPPORTED, "listPrincipals is not supported by Glue"); - } - @Override public Set listRoleGrants(HivePrincipal principal) { @@ -1155,28 +1477,62 @@ public void checkSupportsTransactions() throw new TrinoException(NOT_SUPPORTED, "Glue does not support ACID tables"); } - static class StatsRecordingAsyncHandler - implements AsyncHandler + /** + * Run all tasks on executor returning as soon as all complete or any task fails. + * Upon task execution failure, other tasks are canceled and interrupted, but not waited + * for. + * + * @return results of all tasks in any order + * @throws ExecutionException if any task fails; exception cause is the first task failure + */ + private List runParallel(Collection> tasks) + throws ExecutionException { - private final AwsApiCallStats stats; - private final Stopwatch stopwatch; + return processWithAdditionalThreads(tasks, executor); + } + + public enum TableKind + { + OTHER, DELTA, ICEBERG, HUDI + } - public StatsRecordingAsyncHandler(AwsApiCallStats stats) - { - this.stats = requireNonNull(stats, "stats is null"); - this.stopwatch = Stopwatch.createStarted(); + private static Predicate createTablePredicate(Set visibleTableKinds) + { + // nothing is visible + if (visibleTableKinds.isEmpty()) { + return ignore -> false; } - @Override - public void onError(Exception e) - { - stats.recordCall(stopwatch.elapsed(NANOSECONDS), true); + // everything is visible + if (visibleTableKinds.equals(ImmutableSet.copyOf(TableKind.values()))) { + return ignore -> true; } - @Override - public void onSuccess(AmazonWebServiceRequest request, Object o) - { - stats.recordCall(stopwatch.elapsed(NANOSECONDS), false); + // exclusion + if (visibleTableKinds.contains(TableKind.OTHER)) { + EnumSet deniedKinds = EnumSet.complementOf(EnumSet.copyOf(visibleTableKinds)); + Predicate tableVisibilityFilter = ignore -> true; + for (TableKind value : deniedKinds) { + tableVisibilityFilter = tableVisibilityFilter.and(not(tableKindPredicate(value))); + } + return tableVisibilityFilter; } + + // inclusion + Predicate tableVisibilityFilter = ignore -> false; + for (TableKind tableKind : visibleTableKinds) { + tableVisibilityFilter = tableVisibilityFilter.or(tableKindPredicate(tableKind)); + } + return tableVisibilityFilter; + } + + private static Predicate tableKindPredicate(TableKind tableKind) + { + return switch (tableKind) { + case DELTA -> table -> isDeltaLakeTable(table.parameters()); + case ICEBERG -> table -> isIcebergTable(table.parameters()); + case HUDI -> table -> table.storageDescriptor() != null && isHudiTable(table.storageDescriptor().inputFormat()); + case OTHER -> throw new IllegalArgumentException("Predicate can not be created for OTHER"); + }; } } diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/GlueHiveMetastoreConfig.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/GlueHiveMetastoreConfig.java index 9d7873f73087..9da45bc69bca 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/GlueHiveMetastoreConfig.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/GlueHiveMetastoreConfig.java @@ -42,12 +42,12 @@ public class GlueHiveMetastoreConfig private Optional awsAccessKey = Optional.empty(); private Optional awsSecretKey = Optional.empty(); private Optional awsCredentialsProvider = Optional.empty(); + private boolean useWebIdentityTokenCredentialsProvider; private Optional catalogId = Optional.empty(); private int partitionSegments = 5; - private int getPartitionThreads = 20; - private int readStatisticsThreads = 5; - private int writeStatisticsThreads = 20; + private int threads = 40; private boolean assumeCanonicalPartitionKeys; + private boolean skipArchive; public Optional getGlueRegion() { @@ -221,6 +221,20 @@ public GlueHiveMetastoreConfig setAwsSecretKey(String awsSecretKey) return this; } + public boolean isUseWebIdentityTokenCredentialsProvider() + { + return useWebIdentityTokenCredentialsProvider; + } + + @Config("hive.metastore.glue.use-web-identity-token-credentials-provider") + @ConfigDescription("If true, explicitly use the WebIdentityTokenCredentialsProvider" + + " instead of the default credential provider chain.") + public GlueHiveMetastoreConfig setUseWebIdentityTokenCredentialsProvider(boolean useWebIdentityTokenCredentialsProvider) + { + this.useWebIdentityTokenCredentialsProvider = useWebIdentityTokenCredentialsProvider; + return this; + } + public Optional getCatalogId() { return catalogId; @@ -263,16 +277,16 @@ public GlueHiveMetastoreConfig setPartitionSegments(int partitionSegments) } @Min(1) - public int getGetPartitionThreads() + public int getThreads() { - return getPartitionThreads; + return threads; } @Config("hive.metastore.glue.get-partition-threads") @ConfigDescription("Number of threads for parallel partition fetches from Glue") - public GlueHiveMetastoreConfig setGetPartitionThreads(int getPartitionThreads) + public GlueHiveMetastoreConfig setThreads(int threads) { - this.getPartitionThreads = getPartitionThreads; + this.threads = threads; return this; } @@ -289,31 +303,16 @@ public GlueHiveMetastoreConfig setAssumeCanonicalPartitionKeys(boolean assumeCan return this; } - @Min(1) - public int getReadStatisticsThreads() - { - return readStatisticsThreads; - } - - @Config("hive.metastore.glue.read-statistics-threads") - @ConfigDescription("Number of threads for parallel statistics reads from Glue") - public GlueHiveMetastoreConfig setReadStatisticsThreads(int getReadStatisticsThreads) - { - this.readStatisticsThreads = getReadStatisticsThreads; - return this; - } - - @Min(1) - public int getWriteStatisticsThreads() + public boolean isSkipArchive() { - return writeStatisticsThreads; + return skipArchive; } - @Config("hive.metastore.glue.write-statistics-threads") - @ConfigDescription("Number of threads for parallel statistics writes to Glue") - public GlueHiveMetastoreConfig setWriteStatisticsThreads(int writeStatisticsThreads) + @Config("hive.metastore.glue.skip-archive") + @ConfigDescription("Skip archiving an old table version when updating a table in the Glue metastore") + public GlueHiveMetastoreConfig setSkipArchive(boolean skipArchive) { - this.writeStatisticsThreads = writeStatisticsThreads; + this.skipArchive = skipArchive; return this; } diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/GlueMetastoreModule.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/GlueMetastoreModule.java index 77fcf7043200..5b93c209b9c2 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/GlueMetastoreModule.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/GlueMetastoreModule.java @@ -13,38 +13,55 @@ */ package io.trino.plugin.hive.metastore.glue; -import com.amazonaws.auth.AWSCredentialsProvider; -import com.amazonaws.handlers.RequestHandler2; -import com.amazonaws.services.glue.AWSGlueAsync; -import com.amazonaws.services.glue.model.Table; +import com.google.common.collect.ImmutableList; import com.google.inject.Binder; +import com.google.inject.Inject; import com.google.inject.Key; -import com.google.inject.Module; +import com.google.inject.Provider; import com.google.inject.Provides; import com.google.inject.Scopes; import com.google.inject.Singleton; -import com.google.inject.TypeLiteral; import com.google.inject.multibindings.Multibinder; -import com.google.inject.multibindings.ProvidesIntoSet; -import io.airlift.concurrent.BoundedExecutor; +import com.google.inject.multibindings.ProvidesIntoOptional; import io.airlift.configuration.AbstractConfigurationAwareModule; +import io.airlift.units.Duration; import io.opentelemetry.api.OpenTelemetry; -import io.opentelemetry.instrumentation.awssdk.v1_11.AwsSdkTelemetry; +import io.trino.plugin.base.CatalogName; import io.trino.plugin.hive.AllowHiveTableRename; -import io.trino.plugin.hive.HiveConfig; +import io.trino.plugin.hive.HideDeltaLakeTables; import io.trino.plugin.hive.metastore.HiveMetastoreFactory; import io.trino.plugin.hive.metastore.RawHiveMetastoreFactory; +import io.trino.plugin.hive.metastore.cache.CachingHiveMetastoreConfig; +import io.trino.spi.NodeManager; +import software.amazon.awssdk.auth.credentials.AwsBasicCredentials; +import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider; +import software.amazon.awssdk.core.interceptor.ExecutionInterceptor; +import software.amazon.awssdk.http.apache.ApacheHttpClient; +import software.amazon.awssdk.http.apache.ProxyConfiguration; +import software.amazon.awssdk.regions.Region; +import software.amazon.awssdk.regions.providers.DefaultAwsRegionProviderChain; +import software.amazon.awssdk.retries.api.BackoffStrategy; +import software.amazon.awssdk.services.glue.GlueClient; +import software.amazon.awssdk.services.glue.GlueClientBuilder; +import software.amazon.awssdk.services.glue.model.ConcurrentModificationException; +import software.amazon.awssdk.services.sts.StsClient; +import software.amazon.awssdk.services.sts.StsClientBuilder; +import software.amazon.awssdk.services.sts.auth.StsAssumeRoleCredentialsProvider; +import software.amazon.awssdk.services.sts.auth.StsWebIdentityTokenFileCredentialsProvider; -import java.util.concurrent.Executor; -import java.util.function.Predicate; +import java.net.URI; +import java.util.EnumSet; +import java.util.Optional; +import java.util.Set; -import static com.google.common.util.concurrent.MoreExecutors.directExecutor; +import static com.google.common.base.Preconditions.checkArgument; +import static com.google.common.base.Preconditions.checkState; import static com.google.inject.multibindings.Multibinder.newSetBinder; import static com.google.inject.multibindings.OptionalBinder.newOptionalBinder; -import static io.airlift.concurrent.Threads.daemonThreadsNamed; -import static io.airlift.configuration.ConditionalModule.conditionalModule; +import static com.google.inject.multibindings.ProvidesIntoOptional.Type.DEFAULT; import static io.airlift.configuration.ConfigBinder.configBinder; -import static java.util.concurrent.Executors.newCachedThreadPool; +import static io.trino.plugin.base.ClosingBinder.closingBinder; +import static java.util.Objects.requireNonNull; import static org.weakref.jmx.guice.ExportBinder.newExporter; public class GlueMetastoreModule @@ -53,88 +70,171 @@ public class GlueMetastoreModule @Override protected void setup(Binder binder) { - GlueHiveMetastoreConfig glueConfig = buildConfigObject(GlueHiveMetastoreConfig.class); - Multibinder requestHandlers = newSetBinder(binder, RequestHandler2.class, ForGlueHiveMetastore.class); - glueConfig.getCatalogId().ifPresent(catalogId -> requestHandlers.addBinding().toInstance(new GlueCatalogIdRequestHandler(catalogId))); - glueConfig.getGlueProxyApiId().ifPresent(glueProxyApiId -> requestHandlers.addBinding() - .toInstance(new ProxyApiRequestHandler(glueProxyApiId))); - configBinder(binder).bindConfig(HiveConfig.class); - binder.bind(AWSCredentialsProvider.class).toProvider(GlueCredentialsProvider.class).in(Scopes.SINGLETON); - - newOptionalBinder(binder, Key.get(new TypeLiteral>() {}, ForGlueHiveMetastore.class)) - .setDefault().toProvider(DefaultGlueMetastoreTableFilterProvider.class).in(Scopes.SINGLETON); + configBinder(binder).bindConfig(GlueHiveMetastoreConfig.class); + binder.bind(GlueHiveMetastoreFactory.class).in(Scopes.SINGLETON); binder.bind(GlueHiveMetastore.class).in(Scopes.SINGLETON); + binder.bind(GlueContext.class).in(Scopes.SINGLETON); + newExporter(binder).export(GlueHiveMetastore.class).withGeneratedName(); newOptionalBinder(binder, Key.get(HiveMetastoreFactory.class, RawHiveMetastoreFactory.class)) .setDefault() .to(GlueHiveMetastoreFactory.class) .in(Scopes.SINGLETON); - - // export under the old name, for backwards compatibility - binder.bind(GlueHiveMetastoreFactory.class).in(Scopes.SINGLETON); - binder.bind(Key.get(GlueMetastoreStats.class, ForGlueHiveMetastore.class)).toInstance(new GlueMetastoreStats()); - binder.bind(AWSGlueAsync.class).toProvider(HiveGlueClientProvider.class).in(Scopes.SINGLETON); - newExporter(binder).export(GlueHiveMetastoreFactory.class).as(generator -> generator.generatedNameOf(GlueHiveMetastore.class)); - binder.bind(Key.get(boolean.class, AllowHiveTableRename.class)).toInstance(false); - install(conditionalModule( - HiveConfig.class, - HiveConfig::isTableStatisticsEnabled, - getGlueStatisticsModule(DefaultGlueColumnStatisticsProviderFactory.class), - getGlueStatisticsModule(DisabledGlueColumnStatisticsProviderFactory.class))); - } + Multibinder executionInterceptorMultibinder = newSetBinder(binder, ExecutionInterceptor.class, ForGlueHiveMetastore.class); + executionInterceptorMultibinder.addBinding().toProvider(TelemetryExecutionInterceptorProvider.class).in(Scopes.SINGLETON); + executionInterceptorMultibinder.addBinding().to(GlueHiveExecutionInterceptor.class).in(Scopes.SINGLETON); - private Module getGlueStatisticsModule(Class statisticsPrividerFactoryClass) - { - return internalBinder -> newOptionalBinder(internalBinder, GlueColumnStatisticsProviderFactory.class) - .setDefault() - .to(statisticsPrividerFactoryClass) - .in(Scopes.SINGLETON); + closingBinder(binder).registerCloseable(GlueClient.class); } - @ProvidesIntoSet + @ProvidesIntoOptional(DEFAULT) @Singleton - @ForGlueHiveMetastore - public RequestHandler2 createRequestHandler(OpenTelemetry openTelemetry) + public static Set getTableKinds(@HideDeltaLakeTables boolean hideDeltaLakeTables) { - return AwsSdkTelemetry.builder(openTelemetry) - .setCaptureExperimentalSpanAttributes(true) - .build() - .newRequestHandler(); + if (hideDeltaLakeTables) { + return EnumSet.complementOf(EnumSet.of(GlueHiveMetastore.TableKind.DELTA)); + } + return EnumSet.allOf(GlueHiveMetastore.TableKind.class); } @Provides @Singleton - @ForGlueHiveMetastore - public Executor createExecutor(GlueHiveMetastoreConfig hiveConfig) + public static GlueCache createGlueCache(CachingHiveMetastoreConfig config, CatalogName catalogName, NodeManager nodeManager) { - return createExecutor("hive-glue-partitions-%s", hiveConfig.getGetPartitionThreads()); + Duration metadataCacheTtl = config.getMetastoreCacheTtl(); + Duration statsCacheTtl = config.getStatsCacheTtl(); + + // Disable caching on workers, because there currently is no way to invalidate such a cache. + // Note: while we could skip CachingHiveMetastoreModule altogether on workers, we retain it so that catalog + // configuration can remain identical for all nodes, making cluster configuration easier. + boolean enabled = nodeManager.getCurrentNode().isCoordinator() && + (metadataCacheTtl.toMillis() > 0 || statsCacheTtl.toMillis() > 0); + + checkState(config.isPartitionCacheEnabled(), "Disabling partitions cache is not supported with Glue v2"); + checkState(config.isCacheMissing(), "Disabling cache missing is not supported with Glue v2"); + checkState(config.isCacheMissingPartitions(), "Disabling cache missing partitions is not supported with Glue v2"); + checkState(config.isCacheMissingStats(), "Disabling cache missing stats is not supported with Glue v2"); + + if (enabled) { + return new InMemoryGlueCache( + catalogName, + metadataCacheTtl, + statsCacheTtl, + config.getMetastoreRefreshInterval(), + config.getMaxMetastoreRefreshThreads(), + config.getMetastoreCacheMaximumSize()); + } + return GlueCache.NOOP; } @Provides @Singleton - @ForGlueColumnStatisticsRead - public Executor createStatisticsReadExecutor(GlueHiveMetastoreConfig hiveConfig) + public static GlueClient createGlueClient(GlueHiveMetastoreConfig config, @ForGlueHiveMetastore Set executionInterceptors) { - return createExecutor("hive-glue-statistics-read-%s", hiveConfig.getReadStatisticsThreads()); + GlueClientBuilder glue = GlueClient.builder(); + + glue.overrideConfiguration(builder -> builder + .executionInterceptors(ImmutableList.copyOf(executionInterceptors)) + .retryStrategy(retryBuilder -> retryBuilder + .retryOnException(throwable -> throwable instanceof ConcurrentModificationException) + .backoffStrategy(BackoffStrategy.exponentialDelay( + java.time.Duration.ofMillis(20), + java.time.Duration.ofMillis(1500))) + .maxAttempts(config.getMaxGlueErrorRetries()))); + + Optional staticCredentialsProvider = getStaticCredentialsProvider(config); + + if (config.isUseWebIdentityTokenCredentialsProvider()) { + glue.credentialsProvider(StsWebIdentityTokenFileCredentialsProvider.builder() + .stsClient(getStsClient(config, staticCredentialsProvider)) + .asyncCredentialUpdateEnabled(true) + .build()); + } + else if (config.getIamRole().isPresent()) { + glue.credentialsProvider(StsAssumeRoleCredentialsProvider.builder() + .refreshRequest(request -> request + .roleArn(config.getIamRole().get()) + .roleSessionName("trino-session") + .externalId(config.getExternalId().orElse(null))) + .stsClient(getStsClient(config, staticCredentialsProvider)) + .asyncCredentialUpdateEnabled(true) + .build()); + } + else { + staticCredentialsProvider.ifPresent(glue::credentialsProvider); + } + + ApacheHttpClient.Builder httpClient = ApacheHttpClient.builder() + .maxConnections(config.getMaxGlueConnections()); + + if (config.getGlueEndpointUrl().isPresent()) { + checkArgument(config.getGlueRegion().isPresent(), "Glue region must be set when Glue endpoint URL is set"); + glue.region(Region.of(config.getGlueRegion().get())); + httpClient.proxyConfiguration(ProxyConfiguration.builder() + .endpoint(URI.create(config.getGlueEndpointUrl().get())) + .build()); + } + else if (config.getGlueRegion().isPresent()) { + glue.region(Region.of(config.getGlueRegion().get())); + } + else if (config.getPinGlueClientToCurrentRegion()) { + glue.region(DefaultAwsRegionProviderChain.builder().build().getRegion()); + } + + glue.httpClientBuilder(httpClient); + + return glue.build(); } - @Provides - @Singleton - @ForGlueColumnStatisticsWrite - public Executor createStatisticsWriteExecutor(GlueHiveMetastoreConfig hiveConfig) + private static Optional getStaticCredentialsProvider(GlueHiveMetastoreConfig config) + { + if (config.getAwsAccessKey().isPresent() && config.getAwsSecretKey().isPresent()) { + return Optional.of(StaticCredentialsProvider.create( + AwsBasicCredentials.create(config.getAwsAccessKey().get(), config.getAwsSecretKey().get()))); + } + return Optional.empty(); + } + + private static StsClient getStsClient(GlueHiveMetastoreConfig config, Optional staticCredentialsProvider) { - return createExecutor("hive-glue-statistics-write-%s", hiveConfig.getWriteStatisticsThreads()); + StsClientBuilder sts = StsClient.builder(); + staticCredentialsProvider.ifPresent(sts::credentialsProvider); + + if (config.getGlueStsEndpointUrl().isPresent() && config.getGlueStsRegion().isPresent()) { + sts.endpointOverride(URI.create(config.getGlueStsEndpointUrl().get())) + .region(Region.of(config.getGlueStsRegion().get())); + } + else if (config.getGlueStsRegion().isPresent()) { + sts.region(Region.of(config.getGlueStsRegion().get())); + } + else if (config.getPinGlueClientToCurrentRegion()) { + sts.region(DefaultAwsRegionProviderChain.builder().build().getRegion()); + } + + return sts.build(); } - private Executor createExecutor(String nameTemplate, int threads) + private static class TelemetryExecutionInterceptorProvider + implements Provider { - if (threads == 1) { - return directExecutor(); + private final OpenTelemetry openTelemetry; + + @Inject + public TelemetryExecutionInterceptorProvider(OpenTelemetry openTelemetry) + { + this.openTelemetry = requireNonNull(openTelemetry, "openTelemetry is null"); + } + + @Override + public ExecutionInterceptor get() + { + return io.opentelemetry.instrumentation.awssdk.v2_2.AwsSdkTelemetry.builder(openTelemetry) + .setCaptureExperimentalSpanAttributes(true) + //.setRecordIndividualHttpError(true) + .build() + .newExecutionInterceptor(); } - return new BoundedExecutor( - newCachedThreadPool(daemonThreadsNamed(nameTemplate)), - threads); } } diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/GlueMetastoreStats.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/GlueMetastoreStats.java index 7067de2efaa5..6adc1fa2420d 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/GlueMetastoreStats.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/GlueMetastoreStats.java @@ -13,10 +13,6 @@ */ package io.trino.plugin.hive.metastore.glue; -import com.amazonaws.metrics.RequestMetricCollector; -import io.trino.plugin.hive.aws.AwsApiCallStats; -import io.trino.plugin.hive.aws.AwsSdkClientCoreStats; -import org.weakref.jmx.Flatten; import org.weakref.jmx.Managed; import org.weakref.jmx.Nested; @@ -48,8 +44,6 @@ public class GlueMetastoreStats private final AwsApiCallStats updateColumnStatisticsForPartition = new AwsApiCallStats(); private final AwsApiCallStats deleteColumnStatisticsForPartition = new AwsApiCallStats(); - private final AwsSdkClientCoreStats clientCoreStats = new AwsSdkClientCoreStats(); - @Managed @Nested public AwsApiCallStats getGetDatabases() @@ -224,16 +218,4 @@ public AwsApiCallStats getDeleteColumnStatisticsForPartition() { return deleteColumnStatisticsForPartition; } - - @Managed - @Flatten - public AwsSdkClientCoreStats getClientCoreStats() - { - return clientCoreStats; - } - - public RequestMetricCollector newRequestMetricsCollector() - { - return clientCoreStats.newRequestMetricCollector(); - } } diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/HiveGlueClientProvider.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/HiveGlueClientProvider.java deleted file mode 100644 index 22bd8d7a3e3f..000000000000 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/HiveGlueClientProvider.java +++ /dev/null @@ -1,54 +0,0 @@ -/* - * 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.hive.metastore.glue; - -import com.amazonaws.auth.AWSCredentialsProvider; -import com.amazonaws.handlers.RequestHandler2; -import com.amazonaws.services.glue.AWSGlueAsync; -import com.google.common.collect.ImmutableSet; -import com.google.inject.Inject; -import com.google.inject.Provider; - -import java.util.Set; - -import static io.trino.plugin.hive.metastore.glue.GlueClientUtil.createAsyncGlueClient; -import static java.util.Objects.requireNonNull; - -public class HiveGlueClientProvider - implements Provider -{ - private final GlueMetastoreStats stats; - private final AWSCredentialsProvider credentialsProvider; - private final GlueHiveMetastoreConfig glueConfig; // TODO do not keep mutable config instance on a field - private final Set requestHandlers; - - @Inject - public HiveGlueClientProvider( - @ForGlueHiveMetastore GlueMetastoreStats stats, - AWSCredentialsProvider credentialsProvider, - @ForGlueHiveMetastore Set requestHandlers, - GlueHiveMetastoreConfig glueConfig) - { - this.stats = requireNonNull(stats, "stats is null"); - this.credentialsProvider = requireNonNull(credentialsProvider, "credentialsProvider is null"); - this.requestHandlers = ImmutableSet.copyOf(requireNonNull(requestHandlers, "requestHandlers is null")); - this.glueConfig = glueConfig; - } - - @Override - public AWSGlueAsync get() - { - return createAsyncGlueClient(glueConfig, credentialsProvider, requestHandlers, stats.newRequestMetricsCollector()); - } -} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/InMemoryGlueCache.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/InMemoryGlueCache.java new file mode 100644 index 000000000000..957ab8a70a3d --- /dev/null +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/InMemoryGlueCache.java @@ -0,0 +1,498 @@ +/* + * 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.hive.metastore.glue; + +import com.google.common.cache.CacheBuilder; +import com.google.common.cache.CacheLoader; +import com.google.common.cache.LoadingCache; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.Sets; +import io.airlift.jmx.CacheStatsMBean; +import io.airlift.units.Duration; +import io.trino.cache.SafeCaches; +import io.trino.plugin.base.CatalogName; +import io.trino.plugin.hive.metastore.Database; +import io.trino.plugin.hive.metastore.HiveColumnStatistics; +import io.trino.plugin.hive.metastore.Partition; +import io.trino.plugin.hive.metastore.Table; +import io.trino.plugin.hive.metastore.TableInfo; +import io.trino.plugin.hive.metastore.cache.ReentrantBoundedExecutor; +import io.trino.spi.connector.SchemaTableName; +import jakarta.annotation.PreDestroy; +import org.gaul.modernizer_maven_annotations.SuppressModernizer; +import org.weakref.jmx.Managed; +import org.weakref.jmx.Nested; + +import java.util.Collection; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.OptionalLong; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.Executor; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.atomic.AtomicLong; +import java.util.concurrent.locks.Lock; +import java.util.concurrent.locks.ReentrantLock; +import java.util.function.BiFunction; +import java.util.function.BooleanSupplier; +import java.util.function.Consumer; +import java.util.function.Function; +import java.util.function.Supplier; + +import static com.google.common.cache.CacheLoader.asyncReloading; +import static io.airlift.concurrent.Threads.daemonThreadsNamed; +import static io.trino.cache.CacheUtils.invalidateAllIf; +import static java.util.concurrent.Executors.newCachedThreadPool; +import static java.util.concurrent.TimeUnit.MILLISECONDS; + +class InMemoryGlueCache + implements GlueCache +{ + private enum Global + { + GLOBAL + } + + private record PartitionKey(String databaseName, String tableName, PartitionName partitionName) {} + + private record PartitionNamesKey(String databaseName, String tableName, String glueFilterExpression) {} + + private record FunctionKey(String databaseName, String functionName) {} + + private final ExecutorService refreshExecutor; + + private final LoadingCache>> databaseNamesCache; + private final LoadingCache>> databaseCache; + private final LoadingCache>> tableNamesCache; + private final LoadingCache>> tableCache; + private final LoadingCache tableColumnStatsCache; + private final LoadingCache>> partitionNamesCache; + private final LoadingCache>> partitionCache; + private final LoadingCache partitionColumnStatsCache; + + private final AtomicLong databaseInvalidationCounter = new AtomicLong(); + private final AtomicLong tableInvalidationCounter = new AtomicLong(); + private final AtomicLong partitionInvalidationCounter = new AtomicLong(); + + public InMemoryGlueCache( + CatalogName catalogName, + Duration metadataCacheTtl, + Duration statsCacheTtl, + Optional refreshInterval, + int maxMetastoreRefreshThreads, + long maximumSize) + { + this.refreshExecutor = newCachedThreadPool(daemonThreadsNamed("hive-metastore-" + catalogName + "-%s")); + Executor boundedRefreshExecutor = new ReentrantBoundedExecutor(refreshExecutor, maxMetastoreRefreshThreads); + + OptionalLong refreshMillis = refreshInterval.stream().mapToLong(Duration::toMillis).findAny(); + + OptionalLong metadataCacheTtlMillis = OptionalLong.of(metadataCacheTtl.toMillis()); + this.databaseNamesCache = buildCache(metadataCacheTtlMillis, refreshMillis, boundedRefreshExecutor, maximumSize, ValueHolder::new); + this.databaseCache = buildCache(metadataCacheTtlMillis, refreshMillis, boundedRefreshExecutor, maximumSize, ValueHolder::new); + this.tableNamesCache = buildCache(metadataCacheTtlMillis, refreshMillis, boundedRefreshExecutor, maximumSize, ValueHolder::new); + this.tableCache = buildCache(metadataCacheTtlMillis, refreshMillis, boundedRefreshExecutor, maximumSize, ValueHolder::new); + this.partitionNamesCache = buildCache(metadataCacheTtlMillis, refreshMillis, boundedRefreshExecutor, maximumSize, ValueHolder::new); + this.partitionCache = buildCache(metadataCacheTtlMillis, refreshMillis, boundedRefreshExecutor, maximumSize, ValueHolder::new); + + OptionalLong statsCacheTtlMillis = OptionalLong.of(statsCacheTtl.toMillis()); + this.tableColumnStatsCache = buildCache(statsCacheTtlMillis, refreshMillis, boundedRefreshExecutor, maximumSize, ColumnStatisticsHolder::new); + this.partitionColumnStatsCache = buildCache(statsCacheTtlMillis, refreshMillis, boundedRefreshExecutor, maximumSize, ColumnStatisticsHolder::new); + } + + @PreDestroy + public void stop() + { + refreshExecutor.shutdownNow(); + } + + @Override + public List getDatabaseNames(Function, List> loader) + { + long invalidationCounter = databaseInvalidationCounter.get(); + return databaseNamesCache.getUnchecked(Global.GLOBAL).getValue(() -> loader.apply(database -> cacheDatabase(invalidationCounter, database))); + } + + private void cacheDatabase(long invalidationCounter, Database database) + { + cacheValue(databaseCache, database.getDatabaseName(), Optional.of(database), () -> invalidationCounter == databaseInvalidationCounter.get()); + } + + @Override + public void invalidateDatabase(String databaseName) + { + databaseInvalidationCounter.incrementAndGet(); + databaseCache.invalidate(databaseName); + for (SchemaTableName schemaTableName : Sets.union(tableCache.asMap().keySet(), tableColumnStatsCache.asMap().keySet())) { + if (schemaTableName.getSchemaName().equals(databaseName)) { + invalidateTable(schemaTableName.getSchemaName(), schemaTableName.getTableName(), true); + } + } + for (PartitionKey partitionKey : Sets.union(partitionCache.asMap().keySet(), partitionColumnStatsCache.asMap().keySet())) { + if (partitionKey.databaseName().equals(databaseName)) { + invalidatePartition(partitionKey); + } + } + invalidateAllIf(partitionNamesCache, partitionNamesKey -> partitionNamesKey.databaseName().equals(databaseName)); + } + + @Override + public void invalidateDatabaseNames() + { + databaseNamesCache.invalidate(Global.GLOBAL); + } + + @Override + public Optional getDatabase(String databaseName, Supplier> loader) + { + return databaseCache.getUnchecked(databaseName).getValue(loader); + } + + @Override + public List getTables(String databaseName, Function, List> loader) + { + long invalidationCounter = tableInvalidationCounter.get(); + return tableNamesCache.getUnchecked(databaseName).getValue(() -> loader.apply(table -> cacheTable(invalidationCounter, table))); + } + + private void cacheTable(long invalidationCounter, Table table) + { + cacheValue(tableCache, table.getSchemaTableName(), Optional.of(table), () -> invalidationCounter == tableInvalidationCounter.get()); + } + + @Override + public void invalidateTables(String databaseName) + { + tableNamesCache.invalidate(databaseName); + } + + @Override + public Optional
getTable(String databaseName, String tableName, Supplier> loader) + { + return tableCache.getUnchecked(new SchemaTableName(databaseName, tableName)).getValue(loader); + } + + @Override + public void invalidateTable(String databaseName, String tableName, boolean cascade) + { + tableInvalidationCounter.incrementAndGet(); + SchemaTableName schemaTableName = new SchemaTableName(databaseName, tableName); + tableCache.invalidate(schemaTableName); + tableColumnStatsCache.invalidate(schemaTableName); + if (cascade) { + for (PartitionKey partitionKey : Sets.union(partitionCache.asMap().keySet(), partitionColumnStatsCache.asMap().keySet())) { + if (partitionKey.databaseName().equals(databaseName) && partitionKey.tableName().equals(tableName)) { + invalidatePartition(partitionKey); + } + } + invalidatePartitionNames(databaseName, tableName); + } + } + + @Override + public Map getTableColumnStatistics(String databaseName, String tableName, Set columnNames, Function, Map> loader) + { + return tableColumnStatsCache.getUnchecked(new SchemaTableName(databaseName, tableName)) + .getColumnStatistics(columnNames, loader); + } + + @Override + public void invalidateTableColumnStatistics(String databaseName, String tableName) + { + SchemaTableName schemaTableName = new SchemaTableName(databaseName, tableName); + tableColumnStatsCache.invalidate(schemaTableName); + } + + @Override + public Set getPartitionNames(String databaseName, String tableName, String glueExpression, Function, Set> loader) + { + long invalidationCounter = partitionInvalidationCounter.get(); + return partitionNamesCache.getUnchecked(new PartitionNamesKey(databaseName, tableName, glueExpression)) + .getValue(() -> loader.apply(partition -> cachePartition(invalidationCounter, partition))); + } + + private void invalidatePartitionNames(String databaseName, String tableName) + { + invalidateAllIf(partitionNamesCache, partitionNamesKey -> partitionNamesKey.databaseName().equals(databaseName) && partitionNamesKey.tableName().equals(tableName)); + } + + @Override + public Optional getPartition(String databaseName, String tableName, PartitionName partitionName, Supplier> loader) + { + return partitionCache.getUnchecked(new PartitionKey(databaseName, tableName, partitionName)).getValue(loader); + } + + @Override + public Collection batchGetPartitions( + String databaseName, + String tableName, + Collection partitionNames, + BiFunction, Collection, Collection> loader) + { + ImmutableList.Builder partitions = ImmutableList.builder(); + Set missingPartitionNames = new HashSet<>(); + for (PartitionName partitionName : partitionNames) { + ValueHolder> valueHolder = partitionCache.getIfPresent(new PartitionKey(databaseName, tableName, partitionName)); + if (valueHolder != null) { + Optional partition = valueHolder.getValueIfPresent().flatMap(Function.identity()); + if (partition.isPresent()) { + partitions.add(partition.get()); + continue; + } + } + missingPartitionNames.add(partitionName); + } + if (!missingPartitionNames.isEmpty()) { + // NOTE: loader is expected to directly insert the partitions into the cache, so there is no need to do it here + long invalidationCounter = partitionInvalidationCounter.get(); + partitions.addAll(loader.apply(partition -> cachePartition(invalidationCounter, partition), missingPartitionNames)); + } + return partitions.build(); + } + + private void cachePartition(long invalidationCounter, Partition partition) + { + PartitionKey partitionKey = new PartitionKey(partition.getDatabaseName(), partition.getTableName(), new PartitionName(partition.getValues())); + cacheValue(partitionCache, partitionKey, Optional.of(partition), () -> invalidationCounter == partitionInvalidationCounter.get()); + } + + @Override + public void invalidatePartition(String databaseName, String tableName, PartitionName partitionName) + { + invalidatePartition(new PartitionKey(databaseName, tableName, partitionName)); + } + + private void invalidatePartition(PartitionKey partitionKey) + { + partitionInvalidationCounter.incrementAndGet(); + partitionCache.invalidate(partitionKey); + partitionColumnStatsCache.invalidate(partitionKey); + } + + @Override + public Map getPartitionColumnStatistics( + String databaseName, + String tableName, + PartitionName partitionName, + Set columnNames, + Function, Map> loader) + { + return partitionColumnStatsCache.getUnchecked(new PartitionKey(databaseName, tableName, partitionName)) + .getColumnStatistics(columnNames, loader); + } + + @Override + public void flushCache() + { + databaseInvalidationCounter.incrementAndGet(); + tableInvalidationCounter.incrementAndGet(); + partitionInvalidationCounter.incrementAndGet(); + + databaseNamesCache.invalidateAll(); + databaseCache.invalidateAll(); + tableNamesCache.invalidateAll(); + tableCache.invalidateAll(); + tableColumnStatsCache.invalidateAll(); + partitionNamesCache.invalidateAll(); + partitionCache.invalidateAll(); + partitionColumnStatsCache.invalidateAll(); + } + + @Managed + @Nested + public CacheStatsMBean getDatabaseNamesCacheStats() + { + return new CacheStatsMBean(databaseNamesCache); + } + + @Managed + @Nested + public CacheStatsMBean getDatabaseCacheStats() + { + return new CacheStatsMBean(databaseCache); + } + + @Managed + @Nested + public CacheStatsMBean getTableNamesCacheStats() + { + return new CacheStatsMBean(tableNamesCache); + } + + @Managed + @Nested + public CacheStatsMBean getTableCacheStats() + { + return new CacheStatsMBean(tableCache); + } + + @Managed + @Nested + public CacheStatsMBean getTableColumnStatsCacheStats() + { + return new CacheStatsMBean(tableColumnStatsCache); + } + + @Managed + @Nested + public CacheStatsMBean getPartitionNamesCacheStats() + { + return new CacheStatsMBean(partitionNamesCache); + } + + @Managed + @Nested + public CacheStatsMBean getPartitionCacheStats() + { + return new CacheStatsMBean(partitionCache); + } + + @Managed + @Nested + public CacheStatsMBean getPartitionColumnStatsCacheStats() + { + return new CacheStatsMBean(partitionColumnStatsCache); + } + + @SuppressModernizer + private static LoadingCache buildCache( + OptionalLong expiresAfterWriteMillis, + OptionalLong refreshMillis, + Executor refreshExecutor, + long maximumSize, + Supplier loader) + { + if (expiresAfterWriteMillis.isEmpty()) { + return SafeCaches.emptyLoadingCache(CacheLoader.from(ignores -> loader.get()), true); + } + + CacheLoader cacheLoader = CacheLoader.from(loader::get); + + // this does not use EvictableCache because we want to inject values directly into the cache, + // and we want a lock per key, instead of striped locks + CacheBuilder cacheBuilder = CacheBuilder.newBuilder() + .expireAfterWrite(expiresAfterWriteMillis.getAsLong(), MILLISECONDS) + .maximumSize(maximumSize) + .recordStats(); + + if (refreshMillis.isPresent() && (expiresAfterWriteMillis.getAsLong() > refreshMillis.getAsLong())) { + cacheBuilder.refreshAfterWrite(refreshMillis.getAsLong(), MILLISECONDS); + cacheLoader = asyncReloading(cacheLoader, refreshExecutor); + } + + return cacheBuilder.build(cacheLoader); + } + + private static void cacheValue(LoadingCache> cache, K key, V value, BooleanSupplier test) + { + // get the current value before checking the invalidation counter + ValueHolder valueHolder = cache.getUnchecked(key); + if (!test.getAsBoolean()) { + return; + } + // at this point, we know our value is ok to use in the value cache we fetched before the check + valueHolder.tryOverwrite(value); + // The value is updated, but Guava does not know the update happened, so the expiration time is not extended. + // We need to replace the value in the cache to extend the expiration time iff this is still the latest value. + cache.asMap().replace(key, valueHolder, valueHolder); + } + + private static class ValueHolder + { + private final Lock writeLock = new ReentrantLock(); + private volatile V value; + + public ValueHolder() {} + + public V getValue(Supplier loader) + { + if (value == null) { + writeLock.lock(); + try { + if (value == null) { + value = loader.get(); + if (value == null) { + throw new IllegalStateException("Value loader returned null"); + } + } + } + finally { + writeLock.unlock(); + } + } + return value; + } + + public Optional getValueIfPresent() + { + return Optional.ofNullable(value); + } + + /** + * Overwrite the value unless it is currently being loaded by another thread. + */ + public void tryOverwrite(V value) + { + if (writeLock.tryLock()) { + try { + this.value = value; + } + finally { + writeLock.unlock(); + } + } + } + } + + private static class ColumnStatisticsHolder + { + private final Lock writeLock = new ReentrantLock(); + private final Map> cache = new ConcurrentHashMap<>(); + + public Map getColumnStatistics(Set columnNames, Function, Map> loader) + { + Set missingColumnNames = new HashSet<>(); + Map result = new ConcurrentHashMap<>(); + for (String columnName : columnNames) { + Optional columnStatistics = cache.get(columnName); + if (columnStatistics == null) { + missingColumnNames.add(columnName); + } + else { + columnStatistics.ifPresent(value -> result.put(columnName, value)); + } + } + if (!missingColumnNames.isEmpty()) { + writeLock.lock(); + try { + Map loadedColumnStatistics = loader.apply(missingColumnNames); + for (String missingColumnName : missingColumnNames) { + HiveColumnStatistics value = loadedColumnStatistics.get(missingColumnName); + cache.put(missingColumnName, Optional.ofNullable(value)); + if (value != null) { + result.put(missingColumnName, value); + } + } + } + finally { + writeLock.unlock(); + } + } + return result; + } + } +} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/NoopGlueCache.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/NoopGlueCache.java new file mode 100644 index 000000000000..c3eed329bd64 --- /dev/null +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/NoopGlueCache.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.hive.metastore.glue; + +import io.trino.plugin.hive.metastore.Database; +import io.trino.plugin.hive.metastore.HiveColumnStatistics; +import io.trino.plugin.hive.metastore.Partition; +import io.trino.plugin.hive.metastore.Table; +import io.trino.plugin.hive.metastore.TableInfo; + +import java.util.Collection; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.function.BiFunction; +import java.util.function.Consumer; +import java.util.function.Function; +import java.util.function.Supplier; + +class NoopGlueCache + implements GlueCache +{ + @Override + public List getDatabaseNames(Function, List> loader) + { + return loader.apply(database -> {}); + } + + @Override + public void invalidateDatabase(String databaseName) {} + + @Override + public void invalidateDatabaseNames() {} + + @Override + public Optional getDatabase(String databaseName, Supplier> loader) + { + return loader.get(); + } + + @Override + public List getTables(String databaseName, Function, List> loader) + { + return loader.apply(table -> {}); + } + + @Override + public void invalidateTables(String databaseName) {} + + @Override + public Optional
getTable(String databaseName, String tableName, Supplier> loader) + { + return loader.get(); + } + + @Override + public void invalidateTable(String databaseName, String tableName, boolean cascade) {} + + @Override + public Map getTableColumnStatistics(String databaseName, String tableName, Set columnNames, Function, Map> loader) + { + return loader.apply(columnNames); + } + + @Override + public void invalidateTableColumnStatistics(String databaseName, String tableName) {} + + @Override + public Set getPartitionNames(String databaseName, String tableName, String glueExpression, Function, Set> loader) + { + return loader.apply(partition -> {}); + } + + @Override + public Optional getPartition(String databaseName, String tableName, PartitionName partitionName, Supplier> loader) + { + return loader.get(); + } + + @Override + public Collection batchGetPartitions(String databaseName, String tableName, Collection partitionNames, BiFunction, Collection, Collection> loader) + { + return loader.apply(partition -> {}, partitionNames); + } + + @Override + public void invalidatePartition(String databaseName, String tableName, PartitionName partitionName) {} + + @Override + public Map getPartitionColumnStatistics(String databaseName, String tableName, PartitionName partitionName, Set columnNames, Function, Map> loader) + { + return loader.apply(columnNames); + } + + @Override + public void flushCache() {} +} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/PartitionName.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/PartitionName.java new file mode 100644 index 000000000000..f3d7da4b9069 --- /dev/null +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/PartitionName.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.hive.metastore.glue; + +import java.util.List; + +import static java.util.Objects.requireNonNull; + +public record PartitionName(List partitionValues) +{ + public PartitionName + { + requireNonNull(partitionValues, "partitionValues is null"); + } +} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/converter/GlueInputConverter.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/converter/GlueInputConverter.java deleted file mode 100644 index c0e0e993ba87..000000000000 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/converter/GlueInputConverter.java +++ /dev/null @@ -1,122 +0,0 @@ -/* - * 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.hive.metastore.glue.converter; - -import com.amazonaws.services.glue.model.DatabaseInput; -import com.amazonaws.services.glue.model.Order; -import com.amazonaws.services.glue.model.PartitionInput; -import com.amazonaws.services.glue.model.SerDeInfo; -import com.amazonaws.services.glue.model.StorageDescriptor; -import com.amazonaws.services.glue.model.TableInput; -import com.google.common.collect.ImmutableMap; -import io.trino.plugin.hive.HiveBucketProperty; -import io.trino.plugin.hive.PartitionStatistics; -import io.trino.plugin.hive.metastore.Column; -import io.trino.plugin.hive.metastore.Database; -import io.trino.plugin.hive.metastore.Partition; -import io.trino.plugin.hive.metastore.PartitionWithStatistics; -import io.trino.plugin.hive.metastore.Storage; -import io.trino.plugin.hive.metastore.Table; - -import java.util.List; -import java.util.Optional; - -import static com.google.common.collect.ImmutableList.toImmutableList; -import static io.trino.plugin.hive.metastore.thrift.ThriftMetastoreUtil.updateStatisticsParameters; - -public final class GlueInputConverter -{ - private GlueInputConverter() {} - - public static DatabaseInput convertDatabase(Database database) - { - DatabaseInput input = new DatabaseInput(); - input.setName(database.getDatabaseName()); - input.setParameters(database.getParameters()); - database.getComment().ifPresent(input::setDescription); - database.getLocation().ifPresent(input::setLocationUri); - return input; - } - - public static TableInput convertTable(Table table) - { - TableInput input = new TableInput(); - input.setName(table.getTableName()); - input.setOwner(table.getOwner().orElse(null)); - input.setTableType(table.getTableType()); - input.setStorageDescriptor(convertStorage(table.getStorage(), table.getDataColumns())); - input.setPartitionKeys(table.getPartitionColumns().stream().map(GlueInputConverter::convertColumn).collect(toImmutableList())); - input.setParameters(table.getParameters()); - table.getViewOriginalText().ifPresent(input::setViewOriginalText); - table.getViewExpandedText().ifPresent(input::setViewExpandedText); - return input; - } - - public static PartitionInput convertPartition(PartitionWithStatistics partitionWithStatistics) - { - PartitionInput input = convertPartition(partitionWithStatistics.getPartition()); - PartitionStatistics statistics = partitionWithStatistics.getStatistics(); - input.setParameters(updateStatisticsParameters(input.getParameters(), statistics.getBasicStatistics())); - return input; - } - - public static PartitionInput convertPartition(Partition partition) - { - PartitionInput input = new PartitionInput(); - input.setValues(partition.getValues()); - input.setStorageDescriptor(convertStorage(partition.getStorage(), partition.getColumns())); - input.setParameters(partition.getParameters()); - return input; - } - - private static StorageDescriptor convertStorage(Storage storage, List columns) - { - if (storage.isSkewed()) { - throw new IllegalArgumentException("Writing to skewed table/partition is not supported"); - } - SerDeInfo serdeInfo = new SerDeInfo() - .withSerializationLibrary(storage.getStorageFormat().getSerDeNullable()) - .withParameters(storage.getSerdeParameters()); - - StorageDescriptor sd = new StorageDescriptor(); - sd.setLocation(storage.getLocation()); - sd.setColumns(columns.stream().map(GlueInputConverter::convertColumn).collect(toImmutableList())); - sd.setSerdeInfo(serdeInfo); - sd.setInputFormat(storage.getStorageFormat().getInputFormatNullable()); - sd.setOutputFormat(storage.getStorageFormat().getOutputFormatNullable()); - sd.setParameters(ImmutableMap.of()); - - Optional bucketProperty = storage.getBucketProperty(); - if (bucketProperty.isPresent()) { - sd.setNumberOfBuckets(bucketProperty.get().getBucketCount()); - sd.setBucketColumns(bucketProperty.get().getBucketedBy()); - if (!bucketProperty.get().getSortedBy().isEmpty()) { - sd.setSortColumns(bucketProperty.get().getSortedBy().stream() - .map(column -> new Order().withColumn(column.getColumnName()).withSortOrder(column.getOrder().getHiveOrder())) - .collect(toImmutableList())); - } - } - - return sd; - } - - private static com.amazonaws.services.glue.model.Column convertColumn(Column trinoColumn) - { - return new com.amazonaws.services.glue.model.Column() - .withName(trinoColumn.getName()) - .withType(trinoColumn.getType().toString()) - .withComment(trinoColumn.getComment().orElse(null)) - .withParameters(trinoColumn.getProperties()); - } -} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/converter/GlueStatConverter.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/converter/GlueStatConverter.java deleted file mode 100644 index 1856839f96e8..000000000000 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/converter/GlueStatConverter.java +++ /dev/null @@ -1,309 +0,0 @@ -/* - * 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.hive.metastore.glue.converter; - -import com.amazonaws.services.glue.model.BinaryColumnStatisticsData; -import com.amazonaws.services.glue.model.BooleanColumnStatisticsData; -import com.amazonaws.services.glue.model.ColumnStatistics; -import com.amazonaws.services.glue.model.ColumnStatisticsData; -import com.amazonaws.services.glue.model.ColumnStatisticsType; -import com.amazonaws.services.glue.model.DateColumnStatisticsData; -import com.amazonaws.services.glue.model.DecimalColumnStatisticsData; -import com.amazonaws.services.glue.model.DecimalNumber; -import com.amazonaws.services.glue.model.DoubleColumnStatisticsData; -import com.amazonaws.services.glue.model.LongColumnStatisticsData; -import com.amazonaws.services.glue.model.StringColumnStatisticsData; -import io.trino.hive.thrift.metastore.Decimal; -import io.trino.plugin.hive.HiveType; -import io.trino.plugin.hive.metastore.Column; -import io.trino.plugin.hive.metastore.HiveColumnStatistics; -import io.trino.plugin.hive.metastore.Partition; -import io.trino.plugin.hive.metastore.Table; -import io.trino.plugin.hive.type.PrimitiveTypeInfo; -import io.trino.plugin.hive.type.TypeInfo; -import io.trino.spi.TrinoException; - -import java.math.BigDecimal; -import java.math.BigInteger; -import java.nio.ByteBuffer; -import java.time.LocalDate; -import java.util.Date; -import java.util.List; -import java.util.Map; -import java.util.Optional; -import java.util.OptionalDouble; -import java.util.OptionalLong; -import java.util.concurrent.TimeUnit; - -import static com.google.common.base.Preconditions.checkArgument; -import static com.google.common.collect.ImmutableList.toImmutableList; -import static io.trino.plugin.hive.HiveErrorCode.HIVE_INVALID_METADATA; -import static io.trino.plugin.hive.metastore.HiveColumnStatistics.createBinaryColumnStatistics; -import static io.trino.plugin.hive.metastore.HiveColumnStatistics.createBooleanColumnStatistics; -import static io.trino.plugin.hive.metastore.HiveColumnStatistics.createDateColumnStatistics; -import static io.trino.plugin.hive.metastore.HiveColumnStatistics.createDecimalColumnStatistics; -import static io.trino.plugin.hive.metastore.HiveColumnStatistics.createDoubleColumnStatistics; -import static io.trino.plugin.hive.metastore.HiveColumnStatistics.createIntegerColumnStatistics; -import static io.trino.plugin.hive.metastore.HiveColumnStatistics.createStringColumnStatistics; -import static io.trino.plugin.hive.metastore.thrift.ThriftMetastoreUtil.fromMetastoreDistinctValuesCount; -import static io.trino.plugin.hive.metastore.thrift.ThriftMetastoreUtil.fromMetastoreNullsCount; -import static io.trino.plugin.hive.metastore.thrift.ThriftMetastoreUtil.getAverageColumnLength; -import static io.trino.plugin.hive.metastore.thrift.ThriftMetastoreUtil.getTotalSizeInBytes; -import static io.trino.plugin.hive.metastore.thrift.ThriftMetastoreUtil.toMetastoreDistinctValuesCount; -import static io.trino.plugin.hive.type.Category.PRIMITIVE; - -public class GlueStatConverter -{ - private GlueStatConverter() {} - - private static final long MILLIS_PER_DAY = TimeUnit.DAYS.toMillis(1); - - public static List toGlueColumnStatistics( - Partition partition, - Map trinoColumnStats, - OptionalLong rowCount) - { - return partition.getColumns().stream() - .filter(column -> trinoColumnStats.containsKey(column.getName())) - .map(c -> toColumnStatistics(c, trinoColumnStats.get(c.getName()), rowCount)) - .collect(toImmutableList()); - } - - public static List toGlueColumnStatistics( - Table table, - Map trinoColumnStats, - OptionalLong rowCount) - { - return trinoColumnStats.entrySet().stream() - .map(e -> toColumnStatistics(table.getColumn(e.getKey()).get(), e.getValue(), rowCount)) - .collect(toImmutableList()); - } - - private static ColumnStatistics toColumnStatistics(Column column, HiveColumnStatistics statistics, OptionalLong rowCount) - { - ColumnStatistics columnStatistics = new ColumnStatistics(); - HiveType columnType = column.getType(); - columnStatistics.setColumnName(column.getName()); - columnStatistics.setColumnType(columnType.toString()); - ColumnStatisticsData catalogColumnStatisticsData = toGlueColumnStatisticsData(statistics, columnType, rowCount); - columnStatistics.setStatisticsData(catalogColumnStatisticsData); - columnStatistics.setAnalyzedTime(new Date()); - return columnStatistics; - } - - public static HiveColumnStatistics fromGlueColumnStatistics(ColumnStatisticsData catalogColumnStatisticsData, OptionalLong rowCount) - { - ColumnStatisticsType type = ColumnStatisticsType.fromValue(catalogColumnStatisticsData.getType()); - switch (type) { - case BINARY: { - BinaryColumnStatisticsData data = catalogColumnStatisticsData.getBinaryColumnStatisticsData(); - OptionalLong max = OptionalLong.of(data.getMaximumLength()); - OptionalDouble avg = OptionalDouble.of(data.getAverageLength()); - OptionalLong nulls = fromMetastoreNullsCount(data.getNumberOfNulls()); - return createBinaryColumnStatistics( - max, - getTotalSizeInBytes(avg, rowCount, nulls), - nulls); - } - case BOOLEAN: { - BooleanColumnStatisticsData catalogBooleanData = catalogColumnStatisticsData.getBooleanColumnStatisticsData(); - return createBooleanColumnStatistics( - OptionalLong.of(catalogBooleanData.getNumberOfTrues()), - OptionalLong.of(catalogBooleanData.getNumberOfFalses()), - fromMetastoreNullsCount(catalogBooleanData.getNumberOfNulls())); - } - case DATE: { - DateColumnStatisticsData data = catalogColumnStatisticsData.getDateColumnStatisticsData(); - Optional min = dateToLocalDate(data.getMinimumValue()); - Optional max = dateToLocalDate(data.getMaximumValue()); - OptionalLong nullsCount = fromMetastoreNullsCount(data.getNumberOfNulls()); - OptionalLong distinctValues = OptionalLong.of(data.getNumberOfDistinctValues()); - return createDateColumnStatistics(min, max, nullsCount, fromMetastoreDistinctValuesCount(distinctValues, nullsCount, rowCount)); - } - case DECIMAL: { - DecimalColumnStatisticsData data = catalogColumnStatisticsData.getDecimalColumnStatisticsData(); - Optional min = glueDecimalToBigDecimal(data.getMinimumValue()); - Optional max = glueDecimalToBigDecimal(data.getMaximumValue()); - OptionalLong distinctValues = OptionalLong.of(data.getNumberOfDistinctValues()); - OptionalLong nullsCount = fromMetastoreNullsCount(data.getNumberOfNulls()); - return createDecimalColumnStatistics(min, max, nullsCount, fromMetastoreDistinctValuesCount(distinctValues, nullsCount, rowCount)); - } - case DOUBLE: { - DoubleColumnStatisticsData data = catalogColumnStatisticsData.getDoubleColumnStatisticsData(); - OptionalDouble min = OptionalDouble.of(data.getMinimumValue()); - OptionalDouble max = OptionalDouble.of(data.getMaximumValue()); - OptionalLong nulls = fromMetastoreNullsCount(data.getNumberOfNulls()); - OptionalLong distinctValues = OptionalLong.of(data.getNumberOfDistinctValues()); - return createDoubleColumnStatistics(min, max, nulls, fromMetastoreDistinctValuesCount(distinctValues, nulls, rowCount)); - } - case LONG: { - LongColumnStatisticsData data = catalogColumnStatisticsData.getLongColumnStatisticsData(); - OptionalLong min = OptionalLong.of(data.getMinimumValue()); - OptionalLong max = OptionalLong.of(data.getMaximumValue()); - OptionalLong nullsCount = fromMetastoreNullsCount(data.getNumberOfNulls()); - OptionalLong distinctValues = OptionalLong.of(data.getNumberOfDistinctValues()); - return createIntegerColumnStatistics(min, max, nullsCount, fromMetastoreDistinctValuesCount(distinctValues, nullsCount, rowCount)); - } - case STRING: { - StringColumnStatisticsData data = catalogColumnStatisticsData.getStringColumnStatisticsData(); - OptionalLong max = OptionalLong.of(data.getMaximumLength()); - OptionalDouble avg = OptionalDouble.of(data.getAverageLength()); - OptionalLong nullsCount = fromMetastoreNullsCount(data.getNumberOfNulls()); - OptionalLong distinctValues = OptionalLong.of(data.getNumberOfDistinctValues()); - return createStringColumnStatistics( - max, - getTotalSizeInBytes(avg, rowCount, nullsCount), - nullsCount, - fromMetastoreDistinctValuesCount(distinctValues, nullsCount, rowCount)); - } - } - - throw new TrinoException(HIVE_INVALID_METADATA, "Invalid column statistics data: " + catalogColumnStatisticsData); - } - - private static ColumnStatisticsData toGlueColumnStatisticsData(HiveColumnStatistics statistics, HiveType columnType, OptionalLong rowCount) - { - TypeInfo typeInfo = columnType.getTypeInfo(); - checkArgument(typeInfo.getCategory() == PRIMITIVE, "Unsupported statistics type: %s", columnType); - - ColumnStatisticsData catalogColumnStatisticsData = new ColumnStatisticsData(); - - switch (((PrimitiveTypeInfo) typeInfo).getPrimitiveCategory()) { - case BOOLEAN: { - BooleanColumnStatisticsData data = new BooleanColumnStatisticsData(); - statistics.getNullsCount().ifPresent(data::setNumberOfNulls); - statistics.getBooleanStatistics().ifPresent(booleanStatistics -> { - booleanStatistics.getFalseCount().ifPresent(data::setNumberOfFalses); - booleanStatistics.getTrueCount().ifPresent(data::setNumberOfTrues); - }); - catalogColumnStatisticsData.setType(ColumnStatisticsType.BOOLEAN.toString()); - catalogColumnStatisticsData.setBooleanColumnStatisticsData(data); - break; - } - case BINARY: { - BinaryColumnStatisticsData data = new BinaryColumnStatisticsData(); - statistics.getNullsCount().ifPresent(data::setNumberOfNulls); - data.setMaximumLength(statistics.getMaxValueSizeInBytes().orElse(0)); - data.setAverageLength(getAverageColumnLength(statistics.getTotalSizeInBytes(), rowCount, statistics.getNullsCount()).orElse(0)); - catalogColumnStatisticsData.setType(ColumnStatisticsType.BINARY.toString()); - catalogColumnStatisticsData.setBinaryColumnStatisticsData(data); - break; - } - case DATE: { - DateColumnStatisticsData data = new DateColumnStatisticsData(); - statistics.getDateStatistics().ifPresent(dateStatistics -> { - dateStatistics.getMin().ifPresent(value -> data.setMinimumValue(localDateToDate(value))); - dateStatistics.getMax().ifPresent(value -> data.setMaximumValue(localDateToDate(value))); - }); - statistics.getNullsCount().ifPresent(data::setNumberOfNulls); - toMetastoreDistinctValuesCount(statistics.getDistinctValuesCount(), statistics.getNullsCount()).ifPresent(data::setNumberOfDistinctValues); - catalogColumnStatisticsData.setType(ColumnStatisticsType.DATE.toString()); - catalogColumnStatisticsData.setDateColumnStatisticsData(data); - break; - } - case DECIMAL: { - DecimalColumnStatisticsData data = new DecimalColumnStatisticsData(); - statistics.getDecimalStatistics().ifPresent(decimalStatistics -> { - decimalStatistics.getMin().ifPresent(value -> data.setMinimumValue(bigDecimalToGlueDecimal(value))); - decimalStatistics.getMax().ifPresent(value -> data.setMaximumValue(bigDecimalToGlueDecimal(value))); - }); - statistics.getNullsCount().ifPresent(data::setNumberOfNulls); - toMetastoreDistinctValuesCount(statistics.getDistinctValuesCount(), statistics.getNullsCount()).ifPresent(data::setNumberOfDistinctValues); - catalogColumnStatisticsData.setType(ColumnStatisticsType.DECIMAL.toString()); - catalogColumnStatisticsData.setDecimalColumnStatisticsData(data); - break; - } - case FLOAT: - case DOUBLE: { - DoubleColumnStatisticsData data = new DoubleColumnStatisticsData(); - statistics.getDoubleStatistics().ifPresent(doubleStatistics -> { - doubleStatistics.getMin().ifPresent(data::setMinimumValue); - doubleStatistics.getMax().ifPresent(data::setMaximumValue); - }); - statistics.getNullsCount().ifPresent(data::setNumberOfNulls); - toMetastoreDistinctValuesCount(statistics.getDistinctValuesCount(), statistics.getNullsCount()).ifPresent(data::setNumberOfDistinctValues); - catalogColumnStatisticsData.setType(ColumnStatisticsType.DOUBLE.toString()); - catalogColumnStatisticsData.setDoubleColumnStatisticsData(data); - break; - } - case BYTE: - case SHORT: - case INT: - case LONG: - case TIMESTAMP: { - LongColumnStatisticsData data = new LongColumnStatisticsData(); - statistics.getIntegerStatistics().ifPresent(stats -> { - stats.getMin().ifPresent(data::setMinimumValue); - stats.getMax().ifPresent(data::setMaximumValue); - }); - statistics.getNullsCount().ifPresent(data::setNumberOfNulls); - toMetastoreDistinctValuesCount(statistics.getDistinctValuesCount(), statistics.getNullsCount()).ifPresent(data::setNumberOfDistinctValues); - catalogColumnStatisticsData.setType(ColumnStatisticsType.LONG.toString()); - catalogColumnStatisticsData.setLongColumnStatisticsData(data); - break; - } - case VARCHAR: - case CHAR: - case STRING: { - StringColumnStatisticsData data = new StringColumnStatisticsData(); - statistics.getNullsCount().ifPresent(data::setNumberOfNulls); - toMetastoreDistinctValuesCount(statistics.getDistinctValuesCount(), statistics.getNullsCount()).ifPresent(data::setNumberOfDistinctValues); - data.setMaximumLength(statistics.getMaxValueSizeInBytes().orElse(0)); - data.setAverageLength(getAverageColumnLength(statistics.getTotalSizeInBytes(), rowCount, statistics.getNullsCount()).orElse(0)); - catalogColumnStatisticsData.setType(ColumnStatisticsType.STRING.toString()); - catalogColumnStatisticsData.setStringColumnStatisticsData(data); - break; - } - default: - throw new TrinoException(HIVE_INVALID_METADATA, "Invalid column statistics type: " + ((PrimitiveTypeInfo) typeInfo).getPrimitiveCategory()); - } - return catalogColumnStatisticsData; - } - - private static DecimalNumber bigDecimalToGlueDecimal(BigDecimal decimal) - { - Decimal hiveDecimal = new Decimal((short) decimal.scale(), ByteBuffer.wrap(decimal.unscaledValue().toByteArray())); - DecimalNumber catalogDecimal = new DecimalNumber(); - catalogDecimal.setUnscaledValue(ByteBuffer.wrap(hiveDecimal.getUnscaled())); - catalogDecimal.setScale((int) hiveDecimal.getScale()); - return catalogDecimal; - } - - private static Optional glueDecimalToBigDecimal(DecimalNumber catalogDecimal) - { - if (catalogDecimal == null) { - return Optional.empty(); - } - Decimal decimal = new Decimal(); - decimal.setUnscaled(catalogDecimal.getUnscaledValue()); - decimal.setScale(catalogDecimal.getScale().shortValue()); - return Optional.of(new BigDecimal(new BigInteger(decimal.getUnscaled()), decimal.getScale())); - } - - private static Optional dateToLocalDate(Date date) - { - if (date == null) { - return Optional.empty(); - } - long daysSinceEpoch = date.getTime() / MILLIS_PER_DAY; - return Optional.of(LocalDate.ofEpochDay(daysSinceEpoch)); - } - - private static Date localDateToDate(LocalDate date) - { - long millisecondsSinceEpoch = date.toEpochDay() * MILLIS_PER_DAY; - return new Date(millisecondsSinceEpoch); - } -} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/converter/GlueToTrinoConverter.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/converter/GlueToTrinoConverter.java deleted file mode 100644 index 8fc563af9fd0..000000000000 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/converter/GlueToTrinoConverter.java +++ /dev/null @@ -1,322 +0,0 @@ -/* - * 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.hive.metastore.glue.converter; - -import com.amazonaws.services.glue.model.SerDeInfo; -import com.amazonaws.services.glue.model.StorageDescriptor; -import com.google.common.collect.ImmutableList; -import com.google.common.collect.ImmutableMap; -import io.trino.plugin.hive.HiveBucketProperty; -import io.trino.plugin.hive.HiveStorageFormat; -import io.trino.plugin.hive.HiveType; -import io.trino.plugin.hive.metastore.Column; -import io.trino.plugin.hive.metastore.Database; -import io.trino.plugin.hive.metastore.Partition; -import io.trino.plugin.hive.metastore.SortingColumn; -import io.trino.plugin.hive.metastore.SortingColumn.Order; -import io.trino.plugin.hive.metastore.Storage; -import io.trino.plugin.hive.metastore.StorageFormat; -import io.trino.plugin.hive.metastore.Table; -import io.trino.plugin.hive.util.HiveBucketing; -import io.trino.plugin.hive.util.HiveBucketing.BucketingVersion; -import io.trino.spi.TrinoException; -import io.trino.spi.connector.SchemaTableName; -import io.trino.spi.security.PrincipalType; -import jakarta.annotation.Nullable; -import org.gaul.modernizer_maven_annotations.SuppressModernizer; - -import java.util.List; -import java.util.Locale; -import java.util.Map; -import java.util.Optional; -import java.util.function.Function; -import java.util.function.UnaryOperator; - -import static com.google.common.base.MoreObjects.firstNonNull; -import static com.google.common.base.Strings.emptyToNull; -import static com.google.common.base.Strings.nullToEmpty; -import static io.trino.plugin.hive.HiveErrorCode.HIVE_INVALID_METADATA; -import static io.trino.plugin.hive.HiveErrorCode.HIVE_UNSUPPORTED_FORMAT; -import static io.trino.plugin.hive.HiveType.HIVE_INT; -import static io.trino.plugin.hive.TableType.EXTERNAL_TABLE; -import static io.trino.plugin.hive.ViewReaderUtil.isTrinoMaterializedView; -import static io.trino.plugin.hive.metastore.util.Memoizers.memoizeLast; -import static io.trino.plugin.hive.util.HiveUtil.isDeltaLakeTable; -import static io.trino.plugin.hive.util.HiveUtil.isIcebergTable; -import static java.lang.String.format; -import static java.util.Objects.requireNonNull; - -public final class GlueToTrinoConverter -{ - private static final String PUBLIC_OWNER = "PUBLIC"; - - private GlueToTrinoConverter() {} - - @SuppressModernizer // Usage of `Column.getParameters` is not allowed. Only this method can call that. - public static Map getColumnParameters(com.amazonaws.services.glue.model.Column glueColumn) - { - return firstNonNull(glueColumn.getParameters(), ImmutableMap.of()); - } - - public static String getTableType(com.amazonaws.services.glue.model.Table glueTable) - { - // Athena treats missing table type as EXTERNAL_TABLE. - return firstNonNull(getTableTypeNullable(glueTable), EXTERNAL_TABLE.name()); - } - - @Nullable - @SuppressModernizer // Usage of `Table.getTableType` is not allowed. Only this method can call that. - public static String getTableTypeNullable(com.amazonaws.services.glue.model.Table glueTable) - { - return glueTable.getTableType(); - } - - @SuppressModernizer // Usage of `Table.getParameters` is not allowed. Only this method can call that. - public static Map getTableParameters(com.amazonaws.services.glue.model.Table glueTable) - { - return firstNonNull(glueTable.getParameters(), ImmutableMap.of()); - } - - @SuppressModernizer // Usage of `Partition.getParameters` is not allowed. Only this method can call that. - public static Map getPartitionParameters(com.amazonaws.services.glue.model.Partition gluePartition) - { - return firstNonNull(gluePartition.getParameters(), ImmutableMap.of()); - } - - @SuppressModernizer // Usage of `SerDeInfo.getParameters` is not allowed. Only this method can call that. - public static Map getSerDeInfoParameters(com.amazonaws.services.glue.model.SerDeInfo glueSerDeInfo) - { - return firstNonNull(glueSerDeInfo.getParameters(), ImmutableMap.of()); - } - - public static Database convertDatabase(com.amazonaws.services.glue.model.Database glueDb) - { - return Database.builder() - .setDatabaseName(glueDb.getName()) - // Currently it's not possible to create a Glue database with empty location string "" - // (validation error detected: Value '' at 'database.locationUri' failed to satisfy constraint: Member must have length greater than or equal to 1) - // However, it has been observed that Glue databases with empty location do exist in the wild. - .setLocation(Optional.ofNullable(emptyToNull(glueDb.getLocationUri()))) - .setComment(Optional.ofNullable(glueDb.getDescription())) - .setParameters(firstNonNull(glueDb.getParameters(), ImmutableMap.of())) - .setOwnerName(Optional.of(PUBLIC_OWNER)) - .setOwnerType(Optional.of(PrincipalType.ROLE)) - .build(); - } - - public static Table convertTable(com.amazonaws.services.glue.model.Table glueTable, String dbName) - { - SchemaTableName table = new SchemaTableName(dbName, glueTable.getName()); - - String tableType = getTableType(glueTable); - Map tableParameters = ImmutableMap.copyOf(getTableParameters(glueTable)); - Table.Builder tableBuilder = Table.builder() - .setDatabaseName(table.getSchemaName()) - .setTableName(table.getTableName()) - .setOwner(Optional.ofNullable(glueTable.getOwner())) - .setTableType(tableType) - .setParameters(tableParameters) - .setViewOriginalText(Optional.ofNullable(glueTable.getViewOriginalText())) - .setViewExpandedText(Optional.ofNullable(glueTable.getViewExpandedText())); - - StorageDescriptor sd = glueTable.getStorageDescriptor(); - - if (isIcebergTable(tableParameters) || - (sd == null && isDeltaLakeTable(tableParameters)) || - (sd == null && isTrinoMaterializedView(tableType, tableParameters))) { - // Iceberg tables do not need to read the StorageDescriptor field, but we still need to return dummy properties for compatibility - // Delta Lake tables only need to provide a dummy properties if a StorageDescriptor was not explicitly configured. - // Materialized views do not need to read the StorageDescriptor, but we still need to return dummy properties for compatibility - tableBuilder.setDataColumns(ImmutableList.of(new Column("dummy", HIVE_INT, Optional.empty(), ImmutableMap.of()))); - tableBuilder.getStorageBuilder().setStorageFormat(StorageFormat.fromHiveStorageFormat(HiveStorageFormat.PARQUET)); - } - else { - if (sd == null) { - throw new TrinoException(HIVE_UNSUPPORTED_FORMAT, "Table StorageDescriptor is null for table '%s' %s".formatted(table, glueTable)); - } - tableBuilder.setDataColumns(convertColumns(table, sd.getColumns(), sd.getSerdeInfo().getSerializationLibrary())); - if (glueTable.getPartitionKeys() != null) { - tableBuilder.setPartitionColumns(convertColumns(table, glueTable.getPartitionKeys(), sd.getSerdeInfo().getSerializationLibrary())); - } - else { - tableBuilder.setPartitionColumns(ImmutableList.of()); - } - // No benefit to memoizing here, just reusing the implementation - new StorageConverter().setStorageBuilder(sd, tableBuilder.getStorageBuilder(), tableParameters); - } - - return tableBuilder.build(); - } - - private static Column convertColumn(SchemaTableName table, com.amazonaws.services.glue.model.Column glueColumn, String serde) - { - // OpenCSVSerde deserializes columns from csv file into strings, so we set the column type from the metastore - // to string to avoid cast exceptions. - if (HiveStorageFormat.CSV.getSerde().equals(serde)) { - //TODO(https://github.com/trinodb/trino/issues/7240) Add tests - return new Column(glueColumn.getName(), HiveType.HIVE_STRING, Optional.ofNullable(glueColumn.getComment()), getColumnParameters(glueColumn)); - } - return new Column(glueColumn.getName(), convertType(table, glueColumn), Optional.ofNullable(glueColumn.getComment()), getColumnParameters(glueColumn)); - } - - private static HiveType convertType(SchemaTableName table, com.amazonaws.services.glue.model.Column column) - { - try { - return HiveType.valueOf(column.getType().toLowerCase(Locale.ENGLISH)); - } - catch (IllegalArgumentException e) { - throw new TrinoException(HIVE_INVALID_METADATA, "Glue table '%s' column '%s' has invalid data type: %s".formatted(table, column.getName(), column.getType())); - } - } - - private static List convertColumns(SchemaTableName table, List glueColumns, String serde) - { - return mappedCopy(glueColumns, glueColumn -> convertColumn(table, glueColumn, serde)); - } - - private static Function, Map> parametersConverter() - { - return memoizeLast(ImmutableMap::copyOf); - } - - private static boolean isNullOrEmpty(List list) - { - return list == null || list.isEmpty(); - } - - public static final class GluePartitionConverter - implements Function - { - private final Function, List> columnsConverter; - private final Function, Map> parametersConverter = parametersConverter(); - private final StorageConverter storageConverter = new StorageConverter(); - private final String databaseName; - private final String tableName; - private final Map tableParameters; - - public GluePartitionConverter(Table table) - { - requireNonNull(table, "table is null"); - this.databaseName = requireNonNull(table.getDatabaseName(), "databaseName is null"); - this.tableName = requireNonNull(table.getTableName(), "tableName is null"); - this.tableParameters = table.getParameters(); - this.columnsConverter = memoizeLast(glueColumns -> convertColumns( - table.getSchemaTableName(), - glueColumns, - table.getStorage().getStorageFormat().getSerde())); - } - - @Override - public Partition apply(com.amazonaws.services.glue.model.Partition gluePartition) - { - requireNonNull(gluePartition.getStorageDescriptor(), "Partition StorageDescriptor is null"); - StorageDescriptor sd = gluePartition.getStorageDescriptor(); - - if (!databaseName.equals(gluePartition.getDatabaseName())) { - throw new IllegalArgumentException(format("Unexpected databaseName, expected: %s, but found: %s", databaseName, gluePartition.getDatabaseName())); - } - if (!tableName.equals(gluePartition.getTableName())) { - throw new IllegalArgumentException(format("Unexpected tableName, expected: %s, but found: %s", tableName, gluePartition.getTableName())); - } - Partition.Builder partitionBuilder = Partition.builder() - .setDatabaseName(databaseName) - .setTableName(tableName) - .setValues(gluePartition.getValues()) // No memoization benefit - .setColumns(columnsConverter.apply(sd.getColumns())) - .setParameters(parametersConverter.apply(getPartitionParameters(gluePartition))); - - storageConverter.setStorageBuilder(sd, partitionBuilder.getStorageBuilder(), tableParameters); - - return partitionBuilder.build(); - } - } - - private static final class StorageConverter - { - private final Function, List> bucketColumns = memoizeLast(ImmutableList::copyOf); - private final Function, List> sortColumns = memoizeLast(StorageConverter::createSortingColumns); - private final UnaryOperator> bucketProperty = memoizeLast(); - private final Function, Map> serdeParametersConverter = parametersConverter(); - private final StorageFormatConverter storageFormatConverter = new StorageFormatConverter(); - - public void setStorageBuilder(StorageDescriptor sd, Storage.Builder storageBuilder, Map tableParameters) - { - requireNonNull(sd.getSerdeInfo(), "StorageDescriptor SerDeInfo is null"); - SerDeInfo serdeInfo = sd.getSerdeInfo(); - - storageBuilder.setStorageFormat(storageFormatConverter.createStorageFormat(serdeInfo, sd)) - .setLocation(nullToEmpty(sd.getLocation())) - .setBucketProperty(convertToBucketProperty(tableParameters, sd)) - .setSkewed(sd.getSkewedInfo() != null && !isNullOrEmpty(sd.getSkewedInfo().getSkewedColumnNames())) - .setSerdeParameters(serdeParametersConverter.apply(getSerDeInfoParameters(serdeInfo))) - .build(); - } - - private Optional convertToBucketProperty(Map tableParameters, StorageDescriptor sd) - { - if (sd.getNumberOfBuckets() > 0) { - if (isNullOrEmpty(sd.getBucketColumns())) { - throw new TrinoException(HIVE_INVALID_METADATA, "Table/partition metadata has 'numBuckets' set, but 'bucketCols' is not set"); - } - List bucketColumns = this.bucketColumns.apply(sd.getBucketColumns()); - List sortedBy = this.sortColumns.apply(sd.getSortColumns()); - BucketingVersion bucketingVersion = HiveBucketing.getBucketingVersion(tableParameters); - return bucketProperty.apply(Optional.of(new HiveBucketProperty(bucketColumns, bucketingVersion, sd.getNumberOfBuckets(), sortedBy))); - } - return Optional.empty(); - } - - private static List createSortingColumns(List sortColumns) - { - if (isNullOrEmpty(sortColumns)) { - return ImmutableList.of(); - } - return mappedCopy(sortColumns, column -> new SortingColumn(column.getColumn(), Order.fromMetastoreApiOrder(column.getSortOrder(), "unknown"))); - } - } - - private static final class StorageFormatConverter - { - private static final StorageFormat ALL_NULLS = StorageFormat.createNullable(null, null, null); - private final UnaryOperator serializationLib = memoizeLast(); - private final UnaryOperator inputFormat = memoizeLast(); - private final UnaryOperator outputFormat = memoizeLast(); - // Second phase to attempt memoization on the entire instance beyond just the fields - private final UnaryOperator storageFormat = memoizeLast(); - - public StorageFormat createStorageFormat(SerDeInfo serdeInfo, StorageDescriptor storageDescriptor) - { - String serializationLib = this.serializationLib.apply(serdeInfo.getSerializationLibrary()); - String inputFormat = this.inputFormat.apply(storageDescriptor.getInputFormat()); - String outputFormat = this.outputFormat.apply(storageDescriptor.getOutputFormat()); - if (serializationLib == null && inputFormat == null && outputFormat == null) { - return ALL_NULLS; - } - return this.storageFormat.apply(StorageFormat.createNullable(serializationLib, inputFormat, outputFormat)); - } - } - - public static List mappedCopy(List list, Function mapper) - { - requireNonNull(list, "list is null"); - requireNonNull(mapper, "mapper is null"); - // Uses a pre-sized builder to avoid intermediate allocations and copies, which is especially significant when the - // number of elements is large and the size of the resulting list can be known in advance - ImmutableList.Builder builder = ImmutableList.builderWithExpectedSize(list.size()); - for (T item : list) { - builder.add(mapper.apply(item)); - } - return builder.build(); - } -} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/recording/HiveMetastoreRecording.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/recording/HiveMetastoreRecording.java index d30430ba292b..4535afa1c543 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/recording/HiveMetastoreRecording.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/recording/HiveMetastoreRecording.java @@ -33,6 +33,7 @@ import io.trino.plugin.hive.metastore.Partition; import io.trino.plugin.hive.metastore.PartitionFilter; import io.trino.plugin.hive.metastore.Table; +import io.trino.plugin.hive.metastore.TableInfo; import io.trino.plugin.hive.metastore.TablesWithParameterCacheKey; import io.trino.plugin.hive.metastore.UserTableKey; import io.trino.spi.TrinoException; @@ -71,7 +72,7 @@ public class HiveMetastoreRecording private final NonEvictableCache> tableCache; private final NonEvictableCache tableStatisticsCache; private final NonEvictableCache partitionStatisticsCache; - private final NonEvictableCache> tableNamesCache; + private final NonEvictableCache> tableNamesCache; private final NonEvictableCache>> allTableNamesCache; private final NonEvictableCache> tablesWithParameterCache; private final NonEvictableCache> viewNamesCache; @@ -179,7 +180,7 @@ public Map getPartitionStatistics(Set getAllTables(String databaseName, Supplier> valueSupplier) + public List getAllTables(String databaseName, Supplier> valueSupplier) { return loadValue(tableNamesCache, databaseName, valueSupplier); } @@ -338,7 +339,7 @@ public static class Recording private final List>> tables; private final List> tableStatistics; private final List> partitionStatistics; - private final List>> allTables; + private final List>> allTables; private final List>> tablesWithParameter; private final List>> allViews; private final List>> partitions; @@ -357,7 +358,7 @@ public Recording( @JsonProperty("tables") List>> tables, @JsonProperty("tableStatistics") List> tableStatistics, @JsonProperty("partitionStatistics") List> partitionStatistics, - @JsonProperty("allTables") List>> allTables, + @JsonProperty("allTables") List>> allTables, @JsonProperty("tablesWithParameter") List>> tablesWithParameter, @JsonProperty("allViews") List>> allViews, @JsonProperty("partitions") List>> partitions, @@ -429,7 +430,7 @@ public List> getPartitionStatistics } @JsonProperty - public List>> getAllTables() + public List>> getAllTables() { return allTables; } diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/recording/RecordingHiveMetastore.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/recording/RecordingHiveMetastore.java index 0048d4e32ea9..fb8a0eebec83 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/recording/RecordingHiveMetastore.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/recording/RecordingHiveMetastore.java @@ -13,11 +13,11 @@ */ package io.trino.plugin.hive.metastore.recording; -import io.trino.plugin.hive.HiveColumnStatisticType; +import com.google.common.collect.ImmutableSet; import io.trino.plugin.hive.HiveType; import io.trino.plugin.hive.PartitionStatistics; -import io.trino.plugin.hive.acid.AcidTransaction; import io.trino.plugin.hive.metastore.Database; +import io.trino.plugin.hive.metastore.HiveColumnStatistics; import io.trino.plugin.hive.metastore.HiveMetastore; import io.trino.plugin.hive.metastore.HivePrincipal; import io.trino.plugin.hive.metastore.HivePrivilegeInfo; @@ -25,24 +25,22 @@ import io.trino.plugin.hive.metastore.Partition; import io.trino.plugin.hive.metastore.PartitionWithStatistics; import io.trino.plugin.hive.metastore.PrincipalPrivileges; +import io.trino.plugin.hive.metastore.StatisticsUpdateMode; import io.trino.plugin.hive.metastore.Table; -import io.trino.plugin.hive.metastore.TablesWithParameterCacheKey; +import io.trino.plugin.hive.metastore.TableInfo; import io.trino.plugin.hive.metastore.UserTableKey; -import io.trino.spi.connector.SchemaTableName; import io.trino.spi.predicate.TupleDomain; import io.trino.spi.security.RoleGrant; -import io.trino.spi.type.Type; import java.util.List; import java.util.Map; import java.util.Optional; +import java.util.OptionalLong; import java.util.Set; -import java.util.function.Function; import static com.google.common.collect.ImmutableSet.toImmutableSet; import static io.trino.plugin.hive.metastore.HivePartitionName.hivePartitionName; import static io.trino.plugin.hive.metastore.HiveTableName.hiveTableName; -import static io.trino.plugin.hive.metastore.MetastoreUtil.makePartitionName; import static io.trino.plugin.hive.metastore.PartitionFilter.partitionFilter; import static java.util.Objects.requireNonNull; @@ -77,83 +75,39 @@ public Optional
getTable(String databaseName, String tableName) } @Override - public Set getSupportedColumnStatistics(Type type) + public Map getTableColumnStatistics(String databaseName, String tableName, Set columnNames) { - // No need to record that, since it's a pure local operation. - return delegate.getSupportedColumnStatistics(type); + return delegate.getTableColumnStatistics(databaseName, tableName, columnNames); } @Override - public PartitionStatistics getTableStatistics(Table table) + public Map> getPartitionColumnStatistics(String databaseName, String tableName, Set partitionNames, Set columnNames) { - return recording.getTableStatistics( - hiveTableName(table.getDatabaseName(), table.getTableName()), - () -> delegate.getTableStatistics(table)); + return delegate.getPartitionColumnStatistics(databaseName, tableName, partitionNames, columnNames); } @Override - public Map getPartitionStatistics(Table table, List partitions) + public void updateTableStatistics(String databaseName, String tableName, OptionalLong acidWriteId, StatisticsUpdateMode mode, PartitionStatistics statisticsUpdate) { - return recording.getPartitionStatistics( - partitions.stream() - .map(partition -> hivePartitionName(hiveTableName(table.getDatabaseName(), table.getTableName()), makePartitionName(table, partition))) - .collect(toImmutableSet()), - () -> delegate.getPartitionStatistics(table, partitions)); - } - - @Override - public void updateTableStatistics(String databaseName, - String tableName, - AcidTransaction transaction, - Function update) - { - verifyRecordingMode(); - delegate.updateTableStatistics(databaseName, tableName, transaction, update); - } - - @Override - public void updatePartitionStatistics(Table table, String partitionName, Function update) - { - verifyRecordingMode(); - delegate.updatePartitionStatistics(table, partitionName, update); - } - - @Override - public void updatePartitionStatistics(Table table, Map> updates) - { - verifyRecordingMode(); - delegate.updatePartitionStatistics(table, updates); - } - - @Override - public List getAllTables(String databaseName) - { - return recording.getAllTables(databaseName, () -> delegate.getAllTables(databaseName)); - } - - @Override - public List getTablesWithParameter(String databaseName, String parameterKey, String parameterValue) - { - TablesWithParameterCacheKey key = new TablesWithParameterCacheKey(databaseName, parameterKey, parameterValue); - return recording.getTablesWithParameter(key, () -> delegate.getTablesWithParameter(databaseName, parameterKey, parameterValue)); + delegate.updateTableStatistics(databaseName, tableName, acidWriteId, mode, statisticsUpdate); } @Override - public List getAllViews(String databaseName) + public void updatePartitionStatistics(Table table, StatisticsUpdateMode mode, Map partitionUpdates) { - return recording.getAllViews(databaseName, () -> delegate.getAllViews(databaseName)); + delegate.updatePartitionStatistics(table, mode, partitionUpdates); } @Override - public Optional> getAllTables() + public List getTables(String databaseName) { - return recording.getAllTables(delegate::getAllTables); + return recording.getAllTables(databaseName, () -> delegate.getTables(databaseName)); } @Override - public Optional> getAllViews() + public List getTableNamesWithParameters(String databaseName, String parameterKey, ImmutableSet parameterValues) { - return recording.getAllViews(delegate::getAllViews); + return delegate.getTableNamesWithParameters(databaseName, parameterKey, parameterValues); } @Override @@ -199,10 +153,10 @@ public void dropTable(String databaseName, String tableName, boolean deleteData) } @Override - public void replaceTable(String databaseName, String tableName, Table newTable, PrincipalPrivileges principalPrivileges) + public void replaceTable(String databaseName, String tableName, Table newTable, PrincipalPrivileges principalPrivileges, Map environmentContext) { verifyRecordingMode(); - delegate.replaceTable(databaseName, tableName, newTable, principalPrivileges); + delegate.replaceTable(databaseName, tableName, newTable, principalPrivileges, environmentContext); } @Override @@ -357,14 +311,6 @@ public void revokeRoles(Set roles, Set grantees, boolean delegate.revokeRoles(roles, grantees, adminOption, grantor); } - @Override - public Set listGrantedPrincipals(String role) - { - return recording.listGrantedPrincipals( - role, - () -> delegate.listGrantedPrincipals(role)); - } - @Override public Set listRoleGrants(HivePrincipal principal) { diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/BridgingHiveMetastore.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/BridgingHiveMetastore.java index fac11fd21d51..b702b3766057 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/BridgingHiveMetastore.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/BridgingHiveMetastore.java @@ -13,17 +13,19 @@ */ package io.trino.plugin.hive.metastore.thrift; +import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; -import io.trino.hive.thrift.metastore.DataOperationType; +import com.google.common.collect.ImmutableSet; import io.trino.hive.thrift.metastore.FieldSchema; -import io.trino.plugin.hive.HiveColumnStatisticType; import io.trino.plugin.hive.HivePartition; import io.trino.plugin.hive.HiveType; import io.trino.plugin.hive.PartitionStatistics; +import io.trino.plugin.hive.SchemaAlreadyExistsException; +import io.trino.plugin.hive.TableAlreadyExistsException; import io.trino.plugin.hive.acid.AcidOperation; -import io.trino.plugin.hive.acid.AcidTransaction; import io.trino.plugin.hive.metastore.AcidTransactionOwner; import io.trino.plugin.hive.metastore.Database; +import io.trino.plugin.hive.metastore.HiveColumnStatistics; import io.trino.plugin.hive.metastore.HiveMetastore; import io.trino.plugin.hive.metastore.HivePrincipal; import io.trino.plugin.hive.metastore.HivePrivilegeInfo; @@ -31,7 +33,9 @@ import io.trino.plugin.hive.metastore.Partition; import io.trino.plugin.hive.metastore.PartitionWithStatistics; import io.trino.plugin.hive.metastore.PrincipalPrivileges; +import io.trino.plugin.hive.metastore.StatisticsUpdateMode; import io.trino.plugin.hive.metastore.Table; +import io.trino.plugin.hive.metastore.TableInfo; import io.trino.plugin.hive.util.HiveUtil; import io.trino.spi.TrinoException; import io.trino.spi.connector.SchemaNotFoundException; @@ -39,25 +43,25 @@ import io.trino.spi.connector.TableNotFoundException; import io.trino.spi.predicate.TupleDomain; import io.trino.spi.security.RoleGrant; -import io.trino.spi.type.Type; import java.util.List; import java.util.Map; import java.util.Optional; import java.util.OptionalLong; import java.util.Set; -import java.util.function.Function; import java.util.stream.Collectors; +import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.collect.ImmutableList.toImmutableList; -import static io.trino.plugin.hive.HiveMetadata.TABLE_COMMENT; -import static io.trino.plugin.hive.metastore.MetastoreUtil.isAvroTableWithSchemaSet; +import static io.trino.plugin.hive.HiveMetadata.TRINO_QUERY_ID_NAME; import static io.trino.plugin.hive.metastore.MetastoreUtil.verifyCanDropColumn; +import static io.trino.plugin.hive.metastore.Table.TABLE_COMMENT; import static io.trino.plugin.hive.metastore.thrift.ThriftMetastoreUtil.csvSchemaFields; import static io.trino.plugin.hive.metastore.thrift.ThriftMetastoreUtil.fromMetastoreApiDatabase; import static io.trino.plugin.hive.metastore.thrift.ThriftMetastoreUtil.fromMetastoreApiTable; import static io.trino.plugin.hive.metastore.thrift.ThriftMetastoreUtil.isAvroTableWithSchemaSet; import static io.trino.plugin.hive.metastore.thrift.ThriftMetastoreUtil.isCsvTable; +import static io.trino.plugin.hive.metastore.thrift.ThriftMetastoreUtil.toDataOperationType; import static io.trino.plugin.hive.metastore.thrift.ThriftMetastoreUtil.toMetastoreApiDatabase; import static io.trino.plugin.hive.metastore.thrift.ThriftMetastoreUtil.toMetastoreApiTable; import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED; @@ -102,74 +106,75 @@ public Optional
getTable(String databaseName, String tableName) } @Override - public Set getSupportedColumnStatistics(Type type) + public Map getTableColumnStatistics(String databaseName, String tableName, Set columnNames) { - return delegate.getSupportedColumnStatistics(type); + checkArgument(!columnNames.isEmpty(), "columnNames is empty"); + return delegate.getTableColumnStatistics(databaseName, tableName, columnNames); } @Override - public PartitionStatistics getTableStatistics(Table table) + public Map> getPartitionColumnStatistics(String databaseName, String tableName, Set partitionNames, Set columnNames) { - return delegate.getTableStatistics(toMetastoreApiTable(table)); + checkArgument(!columnNames.isEmpty(), "columnNames is empty"); + return delegate.getPartitionColumnStatistics(databaseName, tableName, partitionNames, columnNames); } @Override - public Map getPartitionStatistics(Table table, List partitions) + public boolean useSparkTableStatistics() { - return delegate.getPartitionStatistics( - toMetastoreApiTable(table), - partitions.stream() - .map(ThriftMetastoreUtil::toMetastoreApiPartition) - .collect(toImmutableList())); + return delegate.useSparkTableStatistics(); } @Override - public void updateTableStatistics(String databaseName, String tableName, AcidTransaction transaction, Function update) + public void updateTableStatistics(String databaseName, String tableName, OptionalLong acidWriteId, StatisticsUpdateMode mode, PartitionStatistics statisticsUpdate) { - delegate.updateTableStatistics(databaseName, tableName, transaction, update); + delegate.updateTableStatistics(databaseName, tableName, acidWriteId, mode, statisticsUpdate); } @Override - public void updatePartitionStatistics(Table table, Map> updates) + public void updatePartitionStatistics(Table table, StatisticsUpdateMode mode, Map partitionUpdates) { io.trino.hive.thrift.metastore.Table metastoreTable = toMetastoreApiTable(table); - updates.forEach((partitionName, update) -> delegate.updatePartitionStatistics(metastoreTable, partitionName, update)); + partitionUpdates.forEach((partitionName, update) -> delegate.updatePartitionStatistics(metastoreTable, partitionName, mode, update)); } @Override - public List getAllTables(String databaseName) + public List getTables(String databaseName) { - return delegate.getAllTables(databaseName); - } - - @Override - public List getTablesWithParameter(String databaseName, String parameterKey, String parameterValue) - { - return delegate.getTablesWithParameter(databaseName, parameterKey, parameterValue); - } - - @Override - public List getAllViews(String databaseName) - { - return delegate.getAllViews(databaseName); - } - - @Override - public Optional> getAllTables() - { - return delegate.getAllTables(); + return delegate.getTables(databaseName).stream() + .map(table -> new TableInfo( + new SchemaTableName(table.getDbName(), table.getTableName()), + TableInfo.ExtendedRelationType.fromTableTypeAndComment(table.getTableType(), table.getComments()))) + .collect(toImmutableList()); } @Override - public Optional> getAllViews() + public List getTableNamesWithParameters(String databaseName, String parameterKey, ImmutableSet parameterValues) { - return delegate.getAllViews(); + return delegate.getTableNamesWithParameters(databaseName, parameterKey, parameterValues); } @Override public void createDatabase(Database database) { - delegate.createDatabase(toMetastoreApiDatabase(database)); + try { + delegate.createDatabase(toMetastoreApiDatabase(database)); + } + catch (SchemaAlreadyExistsException e) { + // Ignore SchemaAlreadyExistsException when this query has already created the database. + // This may happen when an actually successful metastore create call is retried + // because of a timeout on our side. + String expectedQueryId = database.getParameters().get(TRINO_QUERY_ID_NAME); + if (expectedQueryId != null) { + String existingQueryId = getDatabase(database.getDatabaseName()) + .map(Database::getParameters) + .map(parameters -> parameters.get(TRINO_QUERY_ID_NAME)) + .orElse(null); + if (!expectedQueryId.equals(existingQueryId)) { + throw e; + } + } + } } @Override @@ -210,7 +215,24 @@ public void setDatabaseOwner(String databaseName, HivePrincipal principal) @Override public void createTable(Table table, PrincipalPrivileges principalPrivileges) { - delegate.createTable(toMetastoreApiTable(table, principalPrivileges)); + try { + delegate.createTable(toMetastoreApiTable(table, principalPrivileges)); + } + catch (TableAlreadyExistsException e) { + // Ignore TableAlreadyExistsException when this query has already created the table. + // This may happen when an actually successful metastore create call is retried + // because of a timeout on our side. + String expectedQueryId = table.getParameters().get(TRINO_QUERY_ID_NAME); + if (expectedQueryId != null) { + String existingQueryId = getTable(table.getDatabaseName(), table.getTableName()) + .map(Table::getParameters) + .map(parameters -> parameters.get(TRINO_QUERY_ID_NAME)) + .orElse(null); + if (!expectedQueryId.equals(existingQueryId)) { + throw e; + } + } + } } @Override @@ -220,9 +242,9 @@ public void dropTable(String databaseName, String tableName, boolean deleteData) } @Override - public void replaceTable(String databaseName, String tableName, Table newTable, PrincipalPrivileges principalPrivileges) + public void replaceTable(String databaseName, String tableName, Table newTable, PrincipalPrivileges principalPrivileges, Map environmentContext) { - alterTable(databaseName, tableName, toMetastoreApiTable(newTable, principalPrivileges)); + alterTable(databaseName, tableName, toMetastoreApiTable(newTable, principalPrivileges), environmentContext); } @Override @@ -265,7 +287,7 @@ public void setTableOwner(String databaseName, String tableName, HivePrincipal p .setOwner(Optional.of(principal.getName())) .build(); - delegate.alterTable(databaseName, tableName, toMetastoreApiTable(newTable)); + delegate.alterTable(databaseName, tableName, toMetastoreApiTable(newTable), ImmutableMap.of()); } @Override @@ -274,7 +296,12 @@ public void commentColumn(String databaseName, String tableName, String columnNa io.trino.hive.thrift.metastore.Table table = delegate.getTable(databaseName, tableName) .orElseThrow(() -> new TableNotFoundException(new SchemaTableName(databaseName, tableName))); - for (FieldSchema fieldSchema : table.getSd().getCols()) { + List fieldSchemas = ImmutableList.builder() + .addAll(table.getSd().getCols()) + .addAll(table.getPartitionKeys()) + .build(); + + for (FieldSchema fieldSchema : fieldSchemas) { if (fieldSchema.getName().equals(columnName)) { if (comment.isPresent()) { fieldSchema.setComment(comment.get()); @@ -328,7 +355,12 @@ public void dropColumn(String databaseName, String tableName, String columnName) private void alterTable(String databaseName, String tableName, io.trino.hive.thrift.metastore.Table table) { - delegate.alterTable(databaseName, tableName, table); + delegate.alterTable(databaseName, tableName, table, ImmutableMap.of()); + } + + private void alterTable(String databaseName, String tableName, io.trino.hive.thrift.metastore.Table table, Map context) + { + delegate.alterTable(databaseName, tableName, table, context); } @Override @@ -428,12 +460,6 @@ public void revokeRoles(Set roles, Set grantees, boolean delegate.revokeRoles(roles, grantees, adminOption, grantor); } - @Override - public Set listGrantedPrincipals(String role) - { - return delegate.listGrantedPrincipals(role); - } - @Override public Set listRoleGrants(HivePrincipal principal) { @@ -524,10 +550,10 @@ public void acquireTableWriteLock( long transactionId, String dbName, String tableName, - DataOperationType operation, + AcidOperation acidOperation, boolean isDynamicPartitionWrite) { - delegate.acquireTableWriteLock(transactionOwner, queryId, transactionId, dbName, tableName, operation, isDynamicPartitionWrite); + delegate.acquireTableWriteLock(transactionOwner, queryId, transactionId, dbName, tableName, toDataOperationType(acidOperation), isDynamicPartitionWrite); } @Override @@ -536,20 +562,10 @@ public void updateTableWriteId(String dbName, String tableName, long transaction delegate.updateTableWriteId(dbName, tableName, transactionId, writeId, rowCountChange); } - @Override - public void alterPartitions(String dbName, String tableName, List partitions, long writeId) - { - List hadoopPartitions = partitions.stream() - .map(ThriftMetastoreUtil::toMetastoreApiPartition) - .peek(partition -> partition.setWriteId(writeId)) - .collect(toImmutableList()); - delegate.alterPartitions(dbName, tableName, hadoopPartitions, writeId); - } - @Override public void addDynamicPartitions(String dbName, String tableName, List partitionNames, long transactionId, long writeId, AcidOperation operation) { - delegate.addDynamicPartitions(dbName, tableName, partitionNames, transactionId, writeId, operation); + delegate.addDynamicPartitions(dbName, tableName, partitionNames, transactionId, writeId, toDataOperationType(operation)); } @Override diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/DefaultThriftMetastoreClientFactory.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/DefaultThriftMetastoreClientFactory.java index ab63be8ca9ef..ea6967d92610 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/DefaultThriftMetastoreClientFactory.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/DefaultThriftMetastoreClientFactory.java @@ -34,6 +34,7 @@ import java.io.FileInputStream; import java.io.IOException; import java.io.InputStream; +import java.net.URI; import java.security.GeneralSecurityException; import java.security.KeyStore; import java.security.cert.Certificate; @@ -54,31 +55,34 @@ public class DefaultThriftMetastoreClientFactory { private final Optional sslContext; private final Optional socksProxy; - private final int timeoutMillis; + private final int connectTimeoutMillis; + private final int readTimeoutMillis; private final HiveMetastoreAuthentication metastoreAuthentication; private final String hostname; + private final Optional catalogName; private final MetastoreSupportsDateStatistics metastoreSupportsDateStatistics = new MetastoreSupportsDateStatistics(); private final AtomicInteger chosenGetTableAlternative = new AtomicInteger(Integer.MAX_VALUE); private final AtomicInteger chosenTableParamAlternative = new AtomicInteger(Integer.MAX_VALUE); - private final AtomicInteger chosenGetAllViewsPerDatabaseAlternative = new AtomicInteger(Integer.MAX_VALUE); private final AtomicInteger chosenAlterTransactionalTableAlternative = new AtomicInteger(Integer.MAX_VALUE); private final AtomicInteger chosenAlterPartitionsAlternative = new AtomicInteger(Integer.MAX_VALUE); - private final AtomicInteger chosenGetAllTablesAlternative = new AtomicInteger(Integer.MAX_VALUE); - private final AtomicInteger chosenGetAllViewsAlternative = new AtomicInteger(Integer.MAX_VALUE); public DefaultThriftMetastoreClientFactory( Optional sslContext, Optional socksProxy, - Duration timeout, + Duration connectTimeout, + Duration readTimeout, HiveMetastoreAuthentication metastoreAuthentication, - String hostname) + String hostname, + Optional catalogName) { this.sslContext = requireNonNull(sslContext, "sslContext is null"); this.socksProxy = requireNonNull(socksProxy, "socksProxy is null"); - this.timeoutMillis = toIntExact(timeout.toMillis()); + this.connectTimeoutMillis = toIntExact(connectTimeout.toMillis()); + this.readTimeoutMillis = toIntExact(readTimeout.toMillis()); this.metastoreAuthentication = requireNonNull(metastoreAuthentication, "metastoreAuthentication is null"); this.hostname = requireNonNull(hostname, "hostname is null"); + this.catalogName = requireNonNull(catalogName, "catalogName is null"); } @Inject @@ -92,16 +96,18 @@ public DefaultThriftMetastoreClientFactory( config.isTlsEnabled(), Optional.ofNullable(config.getKeystorePath()), Optional.ofNullable(config.getKeystorePassword()), - config.getTruststorePath(), + Optional.ofNullable(config.getTruststorePath()), Optional.ofNullable(config.getTruststorePassword())), Optional.ofNullable(config.getSocksProxy()), - config.getMetastoreTimeout(), + config.getConnectTimeout(), + config.getReadTimeout(), metastoreAuthentication, - nodeManager.getCurrentNode().getHost()); + nodeManager.getCurrentNode().getHost(), + config.getCatalogName()); } @Override - public ThriftMetastoreClient create(HostAndPort address, Optional delegationToken) + public ThriftMetastoreClient create(URI address, Optional delegationToken) throws TTransportException { return create(() -> createTransport(address, delegationToken), hostname); @@ -113,27 +119,26 @@ protected ThriftMetastoreClient create(TransportSupplier transportSupplier, Stri return new ThriftHiveMetastoreClient( transportSupplier, hostname, + catalogName, metastoreSupportsDateStatistics, + true, chosenGetTableAlternative, chosenTableParamAlternative, - chosenGetAllTablesAlternative, - chosenGetAllViewsPerDatabaseAlternative, - chosenGetAllViewsAlternative, chosenAlterTransactionalTableAlternative, chosenAlterPartitionsAlternative); } - private TTransport createTransport(HostAndPort address, Optional delegationToken) + private TTransport createTransport(URI uri, Optional delegationToken) throws TTransportException { - return Transport.create(address, sslContext, socksProxy, timeoutMillis, metastoreAuthentication, delegationToken); + return Transport.create(HostAndPort.fromParts(uri.getHost(), uri.getPort()), sslContext, socksProxy, connectTimeoutMillis, readTimeoutMillis, metastoreAuthentication, delegationToken); } private static Optional buildSslContext( boolean tlsEnabled, Optional keyStorePath, Optional keyStorePassword, - File trustStorePath, + Optional trustStorePath, Optional trustStorePassword) { if (!tlsEnabled) { @@ -163,7 +168,7 @@ private static Optional buildSslContext( } // load TrustStore - KeyStore trustStore = loadTrustStore(trustStorePath, trustStorePassword); + KeyStore trustStore = loadTrustStore(trustStorePath.get(), trustStorePassword); // create TrustManagerFactory TrustManagerFactory trustManagerFactory = TrustManagerFactory.getInstance(TrustManagerFactory.getDefaultAlgorithm()); diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/FailureAwareThriftMetastoreClient.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/FailureAwareThriftMetastoreClient.java index 20e697383a6c..1b4a7146186a 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/FailureAwareThriftMetastoreClient.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/FailureAwareThriftMetastoreClient.java @@ -15,6 +15,7 @@ import com.google.common.annotations.VisibleForTesting; import io.trino.hive.thrift.metastore.ColumnStatisticsObj; +import io.trino.hive.thrift.metastore.DataOperationType; import io.trino.hive.thrift.metastore.Database; import io.trino.hive.thrift.metastore.EnvironmentContext; import io.trino.hive.thrift.metastore.FieldSchema; @@ -28,14 +29,13 @@ import io.trino.hive.thrift.metastore.Role; import io.trino.hive.thrift.metastore.RolePrincipalGrant; import io.trino.hive.thrift.metastore.Table; +import io.trino.hive.thrift.metastore.TableMeta; import io.trino.hive.thrift.metastore.TxnToWriteId; -import io.trino.plugin.hive.acid.AcidOperation; -import io.trino.spi.connector.SchemaTableName; import org.apache.thrift.TException; import java.util.List; import java.util.Map; -import java.util.Optional; +import java.util.Set; import static java.util.Objects.requireNonNull; @@ -85,38 +85,17 @@ public Database getDatabase(String databaseName) } @Override - public List getAllTables(String databaseName) + public List getTableMeta(String databaseName) throws TException { - return runWithHandle(() -> delegate.getAllTables(databaseName)); + return runWithHandle(() -> delegate.getTableMeta(databaseName)); } @Override - public Optional> getAllTables() + public List getTableNamesWithParameters(String databaseName, String parameterKey, Set parameterValues) throws TException { - return runWithHandle(() -> delegate.getAllTables()); - } - - @Override - public List getAllViews(String databaseName) - throws TException - { - return runWithHandle(() -> delegate.getAllViews(databaseName)); - } - - @Override - public Optional> getAllViews() - throws TException - { - return runWithHandle(() -> delegate.getAllViews()); - } - - @Override - public List getTablesWithParameter(String databaseName, String parameterKey, String parameterValue) - throws TException - { - return runWithHandle(() -> delegate.getTablesWithParameter(databaseName, parameterKey, parameterValue)); + return runWithHandle(() -> delegate.getTableNamesWithParameters(databaseName, parameterKey, parameterValues)); } @Override @@ -329,13 +308,6 @@ public void revokeRole(String role, String granteeName, PrincipalType granteeTyp runWithHandle(() -> delegate.revokeRole(role, granteeName, granteeType, grantOption)); } - @Override - public List listGrantedPrincipals(String role) - throws TException - { - return runWithHandle(() -> delegate.listGrantedPrincipals(role)); - } - @Override public List listRoleGrants(String name, PrincipalType principalType) throws TException @@ -435,7 +407,7 @@ public void alterPartitions(String dbName, String tableName, List par } @Override - public void addDynamicPartitions(String dbName, String tableName, List partitionNames, long transactionId, long writeId, AcidOperation operation) + public void addDynamicPartitions(String dbName, String tableName, List partitionNames, long transactionId, long writeId, DataOperationType operation) throws TException { runWithHandle(() -> delegate.addDynamicPartitions(dbName, tableName, partitionNames, transactionId, writeId, operation)); diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/StaticTokenAwareMetastoreClientFactory.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/StaticTokenAwareMetastoreClientFactory.java index bd881a37080c..2899f8b9ccdb 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/StaticTokenAwareMetastoreClientFactory.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/StaticTokenAwareMetastoreClientFactory.java @@ -76,7 +76,6 @@ private StaticTokenAwareMetastoreClientFactory(List metastoreUris, @Nullabl checkArgument(!metastoreUris.isEmpty(), "metastoreUris must specify at least one URI"); this.backoffs = metastoreUris.stream() .map(StaticTokenAwareMetastoreClientFactory::checkMetastoreUri) - .map(uri -> HostAndPort.fromParts(uri.getHost(), uri.getPort())) .map(address -> new Backoff(address, ticker)) .collect(toImmutableList()); @@ -105,7 +104,7 @@ public ThriftMetastoreClient createMetastoreClient(Optional delegationTo TException lastException = null; for (Backoff backoff : backoffsSorted) { try { - return getClient(backoff.getAddress(), backoff, delegationToken); + return getClient(backoff.getUri(), backoff, delegationToken); } catch (TException e) { lastException = e; @@ -116,7 +115,7 @@ public ThriftMetastoreClient createMetastoreClient(Optional delegationTo throw new TException("Failed connecting to Hive metastore: " + addresses, lastException); } - private ThriftMetastoreClient getClient(HostAndPort address, Backoff backoff, Optional delegationToken) + private ThriftMetastoreClient getClient(URI address, Backoff backoff, Optional delegationToken) throws TException { ThriftMetastoreClient client = new FailureAwareThriftMetastoreClient(clientFactory.create(address, delegationToken), new Callback() @@ -156,20 +155,25 @@ static class Backoff static final long MIN_BACKOFF = new Duration(50, MILLISECONDS).roundTo(NANOSECONDS); static final long MAX_BACKOFF = new Duration(60, SECONDS).roundTo(NANOSECONDS); - private final HostAndPort address; + private final URI uri; private final Ticker ticker; private long backoffDuration = MIN_BACKOFF; private OptionalLong lastFailureTimestamp = OptionalLong.empty(); - Backoff(HostAndPort address, Ticker ticker) + Backoff(URI address, Ticker ticker) { - this.address = requireNonNull(address, "address is null"); + this.uri = requireNonNull(address, "address is null"); this.ticker = requireNonNull(ticker, "ticker is null"); } + public URI getUri() + { + return uri; + } + public HostAndPort getAddress() { - return address; + return HostAndPort.fromParts(uri.getHost(), uri.getPort()); } synchronized void fail() diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/ThriftHiveMetastore.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/ThriftHiveMetastore.java index 6a06d94fef7b..2c62feb472e8 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/ThriftHiveMetastore.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/ThriftHiveMetastore.java @@ -22,8 +22,9 @@ import io.airlift.concurrent.MoreFutures; import io.airlift.log.Logger; import io.airlift.units.Duration; -import io.trino.hdfs.HdfsContext; -import io.trino.hdfs.HdfsEnvironment; +import io.trino.filesystem.Location; +import io.trino.filesystem.TrinoFileSystem; +import io.trino.filesystem.TrinoFileSystemFactory; import io.trino.hive.thrift.metastore.AlreadyExistsException; import io.trino.hive.thrift.metastore.ColumnStatisticsObj; import io.trino.hive.thrift.metastore.ConfigValSecurityException; @@ -51,20 +52,18 @@ import io.trino.hive.thrift.metastore.PrivilegeBag; import io.trino.hive.thrift.metastore.PrivilegeGrantInfo; import io.trino.hive.thrift.metastore.Table; +import io.trino.hive.thrift.metastore.TableMeta; import io.trino.hive.thrift.metastore.TxnAbortedException; import io.trino.hive.thrift.metastore.TxnToWriteId; import io.trino.hive.thrift.metastore.UnknownDBException; import io.trino.hive.thrift.metastore.UnknownTableException; import io.trino.plugin.hive.HiveBasicStatistics; -import io.trino.plugin.hive.HiveColumnStatisticType; import io.trino.plugin.hive.HivePartition; import io.trino.plugin.hive.HiveType; import io.trino.plugin.hive.PartitionNotFoundException; import io.trino.plugin.hive.PartitionStatistics; import io.trino.plugin.hive.SchemaAlreadyExistsException; import io.trino.plugin.hive.TableAlreadyExistsException; -import io.trino.plugin.hive.acid.AcidOperation; -import io.trino.plugin.hive.acid.AcidTransaction; import io.trino.plugin.hive.metastore.AcidTransactionOwner; import io.trino.plugin.hive.metastore.Column; import io.trino.plugin.hive.metastore.HiveColumnStatistics; @@ -72,6 +71,7 @@ import io.trino.plugin.hive.metastore.HivePrivilegeInfo; import io.trino.plugin.hive.metastore.HivePrivilegeInfo.HivePrivilege; import io.trino.plugin.hive.metastore.PartitionWithStatistics; +import io.trino.plugin.hive.metastore.StatisticsUpdateMode; import io.trino.plugin.hive.util.RetryDriver; import io.trino.spi.TrinoException; import io.trino.spi.connector.SchemaNotFoundException; @@ -80,14 +80,11 @@ import io.trino.spi.predicate.TupleDomain; import io.trino.spi.security.ConnectorIdentity; import io.trino.spi.security.RoleGrant; -import io.trino.spi.type.Type; -import org.apache.hadoop.fs.Path; import org.apache.thrift.TException; import java.io.IOException; import java.net.InetAddress; import java.net.UnknownHostException; -import java.util.AbstractMap.SimpleEntry; import java.util.ArrayList; import java.util.Collections; import java.util.HashMap; @@ -101,7 +98,7 @@ import java.util.Set; import java.util.concurrent.ExecutorService; import java.util.concurrent.Future; -import java.util.function.Function; +import java.util.concurrent.atomic.AtomicInteger; import java.util.stream.Stream; import static com.google.common.base.Preconditions.checkArgument; @@ -116,22 +113,20 @@ import static io.trino.plugin.hive.HiveErrorCode.HIVE_METASTORE_ERROR; import static io.trino.plugin.hive.HiveErrorCode.HIVE_TABLE_LOCK_NOT_ACQUIRED; import static io.trino.plugin.hive.TableType.MANAGED_TABLE; -import static io.trino.plugin.hive.ViewReaderUtil.PRESTO_VIEW_FLAG; import static io.trino.plugin.hive.metastore.HivePrivilegeInfo.HivePrivilege.OWNERSHIP; import static io.trino.plugin.hive.metastore.MetastoreUtil.adjustRowCount; +import static io.trino.plugin.hive.metastore.MetastoreUtil.getHiveBasicStatistics; import static io.trino.plugin.hive.metastore.MetastoreUtil.partitionKeyFilterToStringList; +import static io.trino.plugin.hive.metastore.MetastoreUtil.updateStatisticsParameters; import static io.trino.plugin.hive.metastore.thrift.ThriftMetastoreUtil.createMetastoreColumnStatistics; import static io.trino.plugin.hive.metastore.thrift.ThriftMetastoreUtil.fromMetastoreApiPrincipalType; import static io.trino.plugin.hive.metastore.thrift.ThriftMetastoreUtil.fromMetastoreApiTable; import static io.trino.plugin.hive.metastore.thrift.ThriftMetastoreUtil.fromRolePrincipalGrants; import static io.trino.plugin.hive.metastore.thrift.ThriftMetastoreUtil.fromTrinoPrincipalType; -import static io.trino.plugin.hive.metastore.thrift.ThriftMetastoreUtil.getBasicStatisticsWithSparkFallback; -import static io.trino.plugin.hive.metastore.thrift.ThriftMetastoreUtil.getHiveBasicStatistics; import static io.trino.plugin.hive.metastore.thrift.ThriftMetastoreUtil.isAvroTableWithSchemaSet; import static io.trino.plugin.hive.metastore.thrift.ThriftMetastoreUtil.parsePrivilege; import static io.trino.plugin.hive.metastore.thrift.ThriftMetastoreUtil.toMetastoreApiPartition; -import static io.trino.plugin.hive.metastore.thrift.ThriftMetastoreUtil.updateStatisticsParameters; -import static io.trino.plugin.hive.util.HiveUtil.makePartName; +import static io.trino.plugin.hive.metastore.thrift.ThriftSparkMetastoreUtil.getSparkTableStatistics; import static io.trino.spi.StandardErrorCode.ALREADY_EXISTS; import static io.trino.spi.StandardErrorCode.GENERIC_USER_ERROR; import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED; @@ -141,7 +136,7 @@ import static java.util.Objects.requireNonNull; @ThreadSafe -public class ThriftHiveMetastore +public final class ThriftHiveMetastore implements ThriftMetastore { private static final Logger log = Logger.get(ThriftHiveMetastore.class); @@ -150,7 +145,7 @@ public class ThriftHiveMetastore private static final CharMatcher DOT_MATCHER = CharMatcher.is('.'); private final Optional identity; - private final HdfsEnvironment hdfsEnvironment; + private final TrinoFileSystemFactory fileSystemFactory; private final IdentityAwareMetastoreClientFactory metastoreClientFactory; private final double backoffScaleFactor; private final Duration minBackoffDelay; @@ -159,7 +154,6 @@ public class ThriftHiveMetastore private final Duration maxWaitForLock; private final int maxRetries; private final boolean deleteFilesOnDrop; - private final boolean translateHiveViews; private final boolean assumeCanonicalPartitionKeys; private final boolean useSparkTableStatisticsFallback; private final ThriftMetastoreStats stats; @@ -167,7 +161,7 @@ public class ThriftHiveMetastore public ThriftHiveMetastore( Optional identity, - HdfsEnvironment hdfsEnvironment, + TrinoFileSystemFactory fileSystemFactory, IdentityAwareMetastoreClientFactory metastoreClientFactory, double backoffScaleFactor, Duration minBackoffDelay, @@ -176,14 +170,13 @@ public ThriftHiveMetastore( Duration maxWaitForLock, int maxRetries, boolean deleteFilesOnDrop, - boolean translateHiveViews, boolean assumeCanonicalPartitionKeys, boolean useSparkTableStatisticsFallback, ThriftMetastoreStats stats, ExecutorService writeStatisticsExecutor) { this.identity = requireNonNull(identity, "identity is null"); - this.hdfsEnvironment = requireNonNull(hdfsEnvironment, "hdfsEnvironment is null"); + this.fileSystemFactory = requireNonNull(fileSystemFactory, "fileSystemFactory is null"); this.metastoreClientFactory = requireNonNull(metastoreClientFactory, "metastoreClientFactory is null"); this.backoffScaleFactor = backoffScaleFactor; this.minBackoffDelay = requireNonNull(minBackoffDelay, "minBackoffDelay is null"); @@ -192,7 +185,6 @@ public ThriftHiveMetastore( this.maxWaitForLock = requireNonNull(maxWaitForLock, "maxWaitForLock is null"); this.maxRetries = maxRetries; this.deleteFilesOnDrop = deleteFilesOnDrop; - this.translateHiveViews = translateHiveViews; this.assumeCanonicalPartitionKeys = assumeCanonicalPartitionKeys; this.useSparkTableStatisticsFallback = useSparkTableStatisticsFallback; this.stats = requireNonNull(stats, "stats is null"); @@ -250,15 +242,15 @@ public Optional getDatabase(String databaseName) } @Override - public List getAllTables(String databaseName) + public List getTables(String databaseName) { try { return retry() .stopOn(NoSuchObjectException.class) .stopOnIllegalExceptions() - .run("getAllTables", () -> { + .run("getTables", () -> { try (ThriftMetastoreClient client = createMetastoreClient()) { - return client.getAllTables(databaseName); + return client.getTableMeta(databaseName); } }); } @@ -274,19 +266,19 @@ public List getAllTables(String databaseName) } @Override - public List getTablesWithParameter(String databaseName, String parameterKey, String parameterValue) + public List getTableNamesWithParameters(String databaseName, String parameterKey, Set parameterValues) { try { return retry() - .stopOn(UnknownDBException.class) + .stopOn(NoSuchObjectException.class) .stopOnIllegalExceptions() - .run("getTablesWithParameter", stats.getGetTablesWithParameter().wrap(() -> { + .run("getTableNamesWithParameters", () -> { try (ThriftMetastoreClient client = createMetastoreClient()) { - return client.getTablesWithParameter(databaseName, parameterKey, parameterValue); + return client.getTableNamesWithParameters(databaseName, parameterKey, parameterValues); } - })); + }); } - catch (UnknownDBException e) { + catch (NoSuchObjectException e) { return ImmutableList.of(); } catch (TException e) { @@ -322,32 +314,7 @@ public Optional
getTable(String databaseName, String tableName) } @Override - public Set getSupportedColumnStatistics(Type type) - { - return ThriftMetastoreUtil.getSupportedColumnStatistics(type); - } - - @Override - public PartitionStatistics getTableStatistics(Table table) - { - List dataColumns = table.getSd().getCols().stream() - .map(FieldSchema::getName) - .collect(toImmutableList()); - Map parameters = table.getParameters(); - HiveBasicStatistics basicStatistics = getHiveBasicStatistics(parameters); - - if (useSparkTableStatisticsFallback && basicStatistics.getRowCount().isEmpty()) { - PartitionStatistics sparkTableStatistics = ThriftSparkMetastoreUtil.getTableStatistics(table); - if (sparkTableStatistics.getBasicStatistics().getRowCount().isPresent()) { - return sparkTableStatistics; - } - } - - Map columnStatistics = getTableColumnStatistics(table.getDbName(), table.getTableName(), dataColumns, basicStatistics.getRowCount()); - return new PartitionStatistics(basicStatistics, columnStatistics); - } - - private Map getTableColumnStatistics(String databaseName, String tableName, List columns, OptionalLong rowCount) + public Map getTableColumnStatistics(String databaseName, String tableName, Set columnNames) { try { return retry() @@ -355,12 +322,12 @@ private Map getTableColumnStatistics(String databa .stopOnIllegalExceptions() .run("getTableColumnStatistics", stats.getGetTableColumnStatistics().wrap(() -> { try (ThriftMetastoreClient client = createMetastoreClient()) { - return groupStatisticsByColumn(client.getTableColumnStatistics(databaseName, tableName, columns), rowCount); + return groupStatisticsByColumn(client.getTableColumnStatistics(databaseName, tableName, ImmutableList.copyOf(columnNames))); } })); } catch (NoSuchObjectException e) { - throw new TableNotFoundException(new SchemaTableName(databaseName, tableName)); + throw new TableNotFoundException(new SchemaTableName(databaseName, tableName), e); } catch (TException e) { throw new TrinoException(HIVE_METASTORE_ERROR, e); @@ -371,53 +338,19 @@ private Map getTableColumnStatistics(String databa } @Override - public Map getPartitionStatistics(Table table, List partitions) + public Map> getPartitionColumnStatistics(String databaseName, String tableName, Set partitionNames, Set columnNames) { - List dataColumns = table.getSd().getCols().stream() - .map(FieldSchema::getName) - .collect(toImmutableList()); - List partitionColumns = table.getPartitionKeys().stream() - .map(FieldSchema::getName) - .collect(toImmutableList()); - - Map partitionBasicStatistics = partitions.stream() + return getPartitionColumnStatistics(databaseName, tableName, partitionNames, ImmutableList.copyOf(columnNames)).entrySet().stream() + .filter(entry -> !entry.getValue().isEmpty()) .collect(toImmutableMap( - partition -> makePartName(partitionColumns, partition.getValues()), - partition -> { - if (useSparkTableStatisticsFallback) { - return getBasicStatisticsWithSparkFallback(partition.getParameters()); - } - return getHiveBasicStatistics(partition.getParameters()); - })); - Map partitionRowCounts = partitionBasicStatistics.entrySet().stream() - .collect(toImmutableMap(Map.Entry::getKey, entry -> entry.getValue().getRowCount())); - - long tableRowCount = partitionRowCounts.values().stream() - .mapToLong(count -> count.orElse(0)) - .sum(); - if (!partitionRowCounts.isEmpty() && tableRowCount == 0) { - // When the table has partitions, but row count statistics are set to zero, we treat this case as empty - // statistics to avoid underestimation in the CBO. This scenario may be caused when other engines are - // used to ingest data into partitioned hive tables. - partitionBasicStatistics = partitionBasicStatistics.keySet().stream() - .map(partitionName -> new SimpleEntry<>(partitionName, HiveBasicStatistics.createEmptyStatistics())) - .collect(toImmutableMap(SimpleEntry::getKey, SimpleEntry::getValue)); - } - - Map> partitionColumnStatistics = getPartitionColumnStatistics( - table.getDbName(), - table.getTableName(), - partitionBasicStatistics.keySet(), - dataColumns, - partitionRowCounts); - ImmutableMap.Builder result = ImmutableMap.builder(); - for (String partitionName : partitionBasicStatistics.keySet()) { - HiveBasicStatistics basicStatistics = partitionBasicStatistics.get(partitionName); - Map columnStatistics = partitionColumnStatistics.getOrDefault(partitionName, ImmutableMap.of()); - result.put(partitionName, new PartitionStatistics(basicStatistics, columnStatistics)); - } + Map.Entry::getKey, + entry -> groupStatisticsByColumn(entry.getValue()))); + } - return result.buildOrThrow(); + @Override + public boolean useSparkTableStatistics() + { + return useSparkTableStatisticsFallback; } @Override @@ -444,21 +377,7 @@ public Optional> getFields(String databaseName, String tableNa } } - private Map> getPartitionColumnStatistics( - String databaseName, - String tableName, - Set partitionNames, - List columnNames, - Map partitionRowCounts) - { - return getMetastorePartitionColumnStatistics(databaseName, tableName, partitionNames, columnNames).entrySet().stream() - .filter(entry -> !entry.getValue().isEmpty()) - .collect(toImmutableMap( - Map.Entry::getKey, - entry -> groupStatisticsByColumn(entry.getValue(), partitionRowCounts.getOrDefault(entry.getKey(), OptionalLong.empty())))); - } - - private Map> getMetastorePartitionColumnStatistics(String databaseName, String tableName, Set partitionNames, List columnNames) + private Map> getPartitionColumnStatistics(String databaseName, String tableName, Set partitionNames, List columnNames) { try { return retry() @@ -471,7 +390,7 @@ private Map> getMetastorePartitionColumnStatis })); } catch (NoSuchObjectException e) { - throw new TableNotFoundException(new SchemaTableName(databaseName, tableName)); + throw new TableNotFoundException(new SchemaTableName(databaseName, tableName), e); } catch (TException e) { throw new TrinoException(HIVE_METASTORE_ERROR, e); @@ -481,11 +400,11 @@ private Map> getMetastorePartitionColumnStatis } } - private static Map groupStatisticsByColumn(List statistics, OptionalLong rowCount) + private static Map groupStatisticsByColumn(List statistics) { Map statisticsByColumn = new HashMap<>(); for (ColumnStatisticsObj stats : statistics) { - HiveColumnStatistics newColumnStatistics = ThriftMetastoreUtil.fromMetastoreApiColumnStatistics(stats, rowCount); + HiveColumnStatistics newColumnStatistics = ThriftMetastoreUtil.fromMetastoreApiColumnStatistics(stats); if (statisticsByColumn.containsKey(stats.getColName())) { HiveColumnStatistics existingColumnStatistics = statisticsByColumn.get(stats.getColName()); if (!newColumnStatistics.equals(existingColumnStatistics)) { @@ -500,35 +419,33 @@ private static Map groupStatisticsByColumn(List update) + public void updateTableStatistics(String databaseName, String tableName, OptionalLong acidWriteId, StatisticsUpdateMode mode, PartitionStatistics statisticsUpdate) { Table originalTable = getTable(databaseName, tableName) .orElseThrow(() -> new TableNotFoundException(new SchemaTableName(databaseName, tableName))); - PartitionStatistics currentStatistics = getTableStatistics(originalTable); - PartitionStatistics updatedStatistics = update.apply(currentStatistics); + PartitionStatistics currentStatistics = getCurrentTableStatistics(originalTable); + PartitionStatistics updatedStatistics = mode.updatePartitionStatistics(currentStatistics, statisticsUpdate); Table modifiedTable = originalTable.deepCopy(); - HiveBasicStatistics basicStatistics = updatedStatistics.getBasicStatistics(); - modifiedTable.setParameters(updateStatisticsParameters(modifiedTable.getParameters(), basicStatistics)); - if (transaction.isAcidTransactionRunning()) { - modifiedTable.setWriteId(transaction.getWriteId()); + modifiedTable.setParameters(updateStatisticsParameters(modifiedTable.getParameters(), updatedStatistics.getBasicStatistics())); + if (acidWriteId.isPresent()) { + modifiedTable.setWriteId(acidWriteId.getAsLong()); } - alterTable(databaseName, tableName, modifiedTable); + alterTable(databaseName, tableName, modifiedTable, ImmutableMap.of()); io.trino.plugin.hive.metastore.Table table = fromMetastoreApiTable(modifiedTable); - OptionalLong rowCount = basicStatistics.getRowCount(); List metastoreColumnStatistics = updatedStatistics.getColumnStatistics().entrySet().stream() .flatMap(entry -> { Optional column = table.getColumn(entry.getKey()); if (column.isEmpty() && isAvroTableWithSchemaSet(modifiedTable)) { // Avro table can have different effective schema than declared in metastore. Still, metastore does not allow - // to store statistics for a column it does not know about. + // storing statistics for a column it does not know about. return Stream.of(); } HiveType type = column.orElseThrow(() -> new IllegalStateException("Column not found: " + entry.getKey())).getType(); - return Stream.of(createMetastoreColumnStatistics(entry.getKey(), type, entry.getValue(), rowCount)); + return Stream.of(createMetastoreColumnStatistics(entry.getKey(), type, entry.getValue())); }) .collect(toImmutableList()); if (!metastoreColumnStatistics.isEmpty()) { @@ -538,6 +455,23 @@ public void updateTableStatistics(String databaseName, String tableName, AcidTra removedColumnStatistics.forEach(column -> deleteTableColumnStatistics(databaseName, tableName, column)); } + private PartitionStatistics getCurrentTableStatistics(Table table) + { + Map columns = table.getSd().getCols().stream() + .collect(toImmutableMap(FieldSchema::getName, fieldSchema -> HiveType.valueOf(fieldSchema.getType()))); + + if (useSparkTableStatisticsFallback) { + Optional sparkTableStatistics = getSparkTableStatistics(table.getParameters(), columns); + if (sparkTableStatistics.isPresent()) { + return sparkTableStatistics.get(); + } + } + + HiveBasicStatistics basicStatistics = getHiveBasicStatistics(table.getParameters()); + Map columnStatistics = getTableColumnStatistics(table.getDbName(), table.getTableName(), columns.keySet()); + return new PartitionStatistics(basicStatistics, columnStatistics); + } + private void setTableColumnStatistics(String databaseName, String tableName, List statistics) { try { @@ -552,7 +486,7 @@ private void setTableColumnStatistics(String databaseName, String tableName, Lis })); } catch (NoSuchObjectException e) { - throw new TableNotFoundException(new SchemaTableName(databaseName, tableName)); + throw new TableNotFoundException(new SchemaTableName(databaseName, tableName), e); } catch (TException e) { throw new TrinoException(HIVE_METASTORE_ERROR, e); @@ -576,7 +510,7 @@ private void deleteTableColumnStatistics(String databaseName, String tableName, })); } catch (NoSuchObjectException e) { - throw new TableNotFoundException(new SchemaTableName(databaseName, tableName)); + throw new TableNotFoundException(new SchemaTableName(databaseName, tableName), e); } catch (TException e) { throw new TrinoException(HIVE_METASTORE_ERROR, e); @@ -587,7 +521,7 @@ private void deleteTableColumnStatistics(String databaseName, String tableName, } @Override - public void updatePartitionStatistics(Table table, String partitionName, Function update) + public void updatePartitionStatistics(Table table, String partitionName, StatisticsUpdateMode mode, PartitionStatistics statisticsUpdate) { List partitions = getPartitionsByNames(table.getDbName(), table.getTableName(), ImmutableList.of(partitionName)); if (partitions.isEmpty()) { @@ -598,9 +532,16 @@ public void updatePartitionStatistics(Table table, String partitionName, Functio } Partition originalPartition = getOnlyElement(partitions); - PartitionStatistics currentStatistics = requireNonNull( - getPartitionStatistics(table, partitions).get(partitionName), "getPartitionStatistics() did not return statistics for partition"); - PartitionStatistics updatedStatistics = update.apply(currentStatistics); + HiveBasicStatistics currentBasicStats = getHiveBasicStatistics(originalPartition.getParameters()); + Map currentColumnStats = getPartitionColumnStatistics( + table.getDbName(), + table.getTableName(), + ImmutableSet.of(partitionName), + table.getSd().getCols().stream() + .map(FieldSchema::getName) + .collect(toImmutableSet())) + .getOrDefault(partitionName, ImmutableMap.of()); + PartitionStatistics updatedStatistics = mode.updatePartitionStatistics(new PartitionStatistics(currentBasicStats, currentColumnStats), statisticsUpdate); Partition modifiedPartition = originalPartition.deepCopy(); HiveBasicStatistics basicStatistics = updatedStatistics.getBasicStatistics(); @@ -609,9 +550,9 @@ public void updatePartitionStatistics(Table table, String partitionName, Functio Map columns = modifiedPartition.getSd().getCols().stream() .collect(toImmutableMap(FieldSchema::getName, schema -> HiveType.valueOf(schema.getType()))); - setPartitionColumnStatistics(table.getDbName(), table.getTableName(), partitionName, columns, updatedStatistics.getColumnStatistics(), basicStatistics.getRowCount()); + setPartitionColumnStatistics(table.getDbName(), table.getTableName(), partitionName, columns, updatedStatistics.getColumnStatistics()); - Set removedStatistics = difference(currentStatistics.getColumnStatistics().keySet(), updatedStatistics.getColumnStatistics().keySet()); + Set removedStatistics = difference(currentColumnStats.keySet(), updatedStatistics.getColumnStatistics().keySet()); removedStatistics.forEach(column -> deletePartitionColumnStatistics(table.getDbName(), table.getTableName(), partitionName, column)); } @@ -620,12 +561,11 @@ private void setPartitionColumnStatistics( String tableName, String partitionName, Map columns, - Map columnStatistics, - OptionalLong rowCount) + Map columnStatistics) { List metastoreColumnStatistics = columnStatistics.entrySet().stream() .filter(entry -> columns.containsKey(entry.getKey())) - .map(entry -> createMetastoreColumnStatistics(entry.getKey(), columns.get(entry.getKey()), entry.getValue(), rowCount)) + .map(entry -> createMetastoreColumnStatistics(entry.getKey(), columns.get(entry.getKey()), entry.getValue())) .collect(toImmutableList()); if (!metastoreColumnStatistics.isEmpty()) { setPartitionColumnStatistics(databaseName, tableName, partitionName, metastoreColumnStatistics); @@ -646,7 +586,7 @@ private void setPartitionColumnStatistics(String databaseName, String tableName, })); } catch (NoSuchObjectException e) { - throw new TableNotFoundException(new SchemaTableName(databaseName, tableName)); + throw new TableNotFoundException(new SchemaTableName(databaseName, tableName), e); } catch (TException e) { throw new TrinoException(HIVE_METASTORE_ERROR, e); @@ -670,7 +610,7 @@ private void deletePartitionColumnStatistics(String databaseName, String tableNa })); } catch (NoSuchObjectException e) { - throw new TableNotFoundException(new SchemaTableName(databaseName, tableName)); + throw new TableNotFoundException(new SchemaTableName(databaseName, tableName), e); } catch (TException e) { throw new TrinoException(HIVE_METASTORE_ERROR, e); @@ -814,27 +754,6 @@ private void revokeRole(String role, String granteeName, PrincipalType granteeTy } } - @Override - public Set listGrantedPrincipals(String role) - { - try { - return retry() - .stopOn(MetaException.class) - .stopOnIllegalExceptions() - .run("listPrincipals", stats.getListGrantedPrincipals().wrap(() -> { - try (ThriftMetastoreClient client = createMetastoreClient()) { - return fromRolePrincipalGrants(client.listGrantedPrincipals(role)); - } - })); - } - catch (TException e) { - throw new TrinoException(HIVE_METASTORE_ERROR, e); - } - catch (Exception e) { - throw propagate(e); - } - } - @Override public Set listRoleGrants(HivePrincipal principal) { @@ -856,81 +775,6 @@ public Set listRoleGrants(HivePrincipal principal) } } - @Override - public List getAllViews(String databaseName) - { - try { - return retry() - .stopOn(UnknownDBException.class) - .stopOnIllegalExceptions() - .run("getAllViews", stats.getGetAllViews().wrap(() -> { - try (ThriftMetastoreClient client = createMetastoreClient()) { - if (translateHiveViews) { - return client.getAllViews(databaseName); - } - return client.getTablesWithParameter(databaseName, PRESTO_VIEW_FLAG, "true"); - } - })); - } - catch (UnknownDBException e) { - return ImmutableList.of(); - } - catch (TException e) { - throw new TrinoException(HIVE_METASTORE_ERROR, e); - } - catch (Exception e) { - throw propagate(e); - } - } - - @Override - public Optional> getAllTables() - { - try { - return retry() - .stopOn(UnknownDBException.class) - .stopOnIllegalExceptions() - .run("getAllTables", stats.getGetAllTables().wrap(() -> { - try (ThriftMetastoreClient client = createMetastoreClient()) { - return client.getAllTables(); - } - })); - } - catch (TException e) { - throw new TrinoException(HIVE_METASTORE_ERROR, e); - } - catch (Exception e) { - throw propagate(e); - } - } - - @Override - public Optional> getAllViews() - { - // Without translateHiveViews, Hive views are represented as tables in Trino, - // and they should not be returned from ThriftHiveMetastore.getAllViews() call - if (!translateHiveViews) { - return Optional.empty(); - } - - try { - return retry() - .stopOn(UnknownDBException.class) - .stopOnIllegalExceptions() - .run("getAllViews", stats.getGetAllViews().wrap(() -> { - try (ThriftMetastoreClient client = createMetastoreClient()) { - return client.getAllViews(); - } - })); - } - catch (TException e) { - throw new TrinoException(HIVE_METASTORE_ERROR, e); - } - catch (Exception e) { - throw propagate(e); - } - } - @Override public void createDatabase(Database database) { @@ -947,7 +791,7 @@ public void createDatabase(Database database) })); } catch (AlreadyExistsException e) { - throw new SchemaAlreadyExistsException(database.getName()); + throw new SchemaAlreadyExistsException(database.getName(), e.getMessage()); } catch (TException e) { throw new TrinoException(HIVE_METASTORE_ERROR, e); @@ -972,7 +816,7 @@ public void dropDatabase(String databaseName, boolean deleteData) })); } catch (NoSuchObjectException e) { - throw new SchemaNotFoundException(databaseName); + throw new SchemaNotFoundException(databaseName, e); } catch (TException e) { throw new TrinoException(HIVE_METASTORE_ERROR, e); @@ -1000,7 +844,7 @@ public void alterDatabase(String databaseName, Database database) })); } catch (NoSuchObjectException e) { - throw new SchemaNotFoundException(databaseName); + throw new SchemaNotFoundException(databaseName, e); } catch (TException e) { throw new TrinoException(HIVE_METASTORE_ERROR, e); @@ -1026,10 +870,10 @@ public void createTable(Table table) })); } catch (AlreadyExistsException e) { - throw new TableAlreadyExistsException(new SchemaTableName(table.getDbName(), table.getTableName())); + throw new TableAlreadyExistsException(new SchemaTableName(table.getDbName(), table.getTableName()), e.getMessage()); } catch (NoSuchObjectException e) { - throw new SchemaNotFoundException(table.getDbName()); + throw new SchemaNotFoundException(table.getDbName(), e.getMessage()); } catch (InvalidObjectException e) { boolean databaseMissing; @@ -1041,7 +885,7 @@ public void createTable(Table table) databaseMissing = false; // we don't know, assume it exists for the purpose of error reporting } if (databaseMissing) { - throw new SchemaNotFoundException(table.getDbName()); + throw new SchemaNotFoundException(table.getDbName(), e); } throw new TrinoException(HIVE_METASTORE_ERROR, e); } @@ -1056,24 +900,38 @@ public void createTable(Table table) @Override public void dropTable(String databaseName, String tableName, boolean deleteData) { + AtomicInteger attemptCount = new AtomicInteger(); try { retry() .stopOn(NoSuchObjectException.class) .stopOnIllegalExceptions() .run("dropTable", stats.getDropTable().wrap(() -> { try (ThriftMetastoreClient client = createMetastoreClient()) { - Table table = client.getTable(databaseName, tableName); + attemptCount.incrementAndGet(); + Table table; + try { + table = client.getTable(databaseName, tableName); + } + catch (NoSuchObjectException e) { + if (attemptCount.get() == 1) { + // Throw exception only on the first attempt. + throw e; + } + // If the table is not found on consecutive attempts, it was probably dropped on the first attempt and timeout occurred. + // Exception in such a case can be safely ignored and dropping table is finished. + return null; + } client.dropTable(databaseName, tableName, deleteData); String tableLocation = table.getSd().getLocation(); if (deleteFilesOnDrop && deleteData && isManagedTable(table) && !isNullOrEmpty(tableLocation)) { - deleteDirRecursive(new Path(tableLocation)); + deleteDirRecursive(Location.of(tableLocation)); } } return null; })); } catch (NoSuchObjectException e) { - throw new TableNotFoundException(new SchemaTableName(databaseName, tableName)); + throw new TableNotFoundException(new SchemaTableName(databaseName, tableName), e); } catch (TException e) { throw new TrinoException(HIVE_METASTORE_ERROR, e); @@ -1083,12 +941,12 @@ public void dropTable(String databaseName, String tableName, boolean deleteData) } } - private void deleteDirRecursive(Path path) + private void deleteDirRecursive(Location path) { try { - HdfsContext context = new HdfsContext(identity.orElseGet(() -> - ConnectorIdentity.ofUser(DEFAULT_METASTORE_USER))); - hdfsEnvironment.getFileSystem(context, path).delete(path, true); + TrinoFileSystem fileSystem = fileSystemFactory.create( + identity.orElseGet(() -> ConnectorIdentity.ofUser(DEFAULT_METASTORE_USER))); + fileSystem.deleteDirectory(path); } catch (IOException | RuntimeException e) { // don't fail if unable to delete path @@ -1102,7 +960,7 @@ private static boolean isManagedTable(Table table) } @Override - public void alterTable(String databaseName, String tableName, Table table) + public void alterTable(String databaseName, String tableName, Table table, Map environmentContext) { if (!Objects.equals(databaseName, table.getDbName())) { validateObjectName(table.getDbName()); @@ -1118,9 +976,12 @@ public void alterTable(String databaseName, String tableName, Table table) try (ThriftMetastoreClient client = createMetastoreClient()) { EnvironmentContext context = new EnvironmentContext(); // This prevents Hive 3.x from collecting basic table stats at table creation time. - // These stats are not useful by themselves and can take very long time to collect when creating an - // external table over large data set. - context.setProperties(ImmutableMap.of("DO_NOT_UPDATE_STATS", "true")); + // These stats are not useful by themselves and can take a very long time to collect when creating an + // external table over a large data set. + context.setProperties(ImmutableMap.builder() + .put("DO_NOT_UPDATE_STATS", "true") + .putAll(environmentContext) + .buildOrThrow()); client.alterTableWithEnvironmentContext(databaseName, tableName, table, context); } return null; @@ -1284,7 +1145,7 @@ public void dropPartition(String databaseName, String tableName, List pa client.dropPartition(databaseName, tableName, parts, deleteData); String partitionLocation = partition.getSd().getLocation(); if (deleteFilesOnDrop && deleteData && !isNullOrEmpty(partitionLocation) && isManagedTable(client.getTable(databaseName, tableName))) { - deleteDirRecursive(new Path(partitionLocation)); + deleteDirRecursive(Location.of(partitionLocation)); } } return null; @@ -1342,7 +1203,7 @@ private void storePartitionColumnStatistics(String databaseName, String tableNam } Map columnTypes = partitionWithStatistics.getPartition().getColumns().stream() .collect(toImmutableMap(Column::getName, Column::getType)); - setPartitionColumnStatistics(databaseName, tableName, partitionName, columnTypes, columnStatistics, statistics.getBasicStatistics().getRowCount()); + setPartitionColumnStatistics(databaseName, tableName, partitionName, columnTypes, columnStatistics); } /* @@ -1351,7 +1212,7 @@ private void storePartitionColumnStatistics(String databaseName, String tableNam * The old statistics are supposed to be replaced by storing the new partition statistics. * * In case when the new statistics are not present for some columns, or if the table schema has changed - * if is needed to explicitly remove the statistics from the metastore for that columns. + * if is needed to explicitly remove the statistics from the metastore for those columns. */ private void dropExtraColumnStatisticsAfterAlterPartition( String databaseName, @@ -1374,7 +1235,7 @@ private void dropExtraColumnStatisticsAfterAlterPartition( // check if statistics for the columnsWithMissingStatistics are actually stored in the metastore // when trying to remove any missing statistics the metastore throws NoSuchObjectException String partitionName = partitionWithStatistics.getPartitionName(); - List statisticsToBeRemoved = getMetastorePartitionColumnStatistics( + List statisticsToBeRemoved = getPartitionColumnStatistics( databaseName, tableName, ImmutableSet.of(partitionName), @@ -1461,7 +1322,7 @@ public void grantTablePrivileges(String databaseName, String tableName, String t HivePrivilegeInfo requestedPrivilege = getOnlyElement(parsePrivilege(iterator.next(), Optional.empty())); for (HivePrivilegeInfo existingPrivilege : existingPrivileges) { - if ((requestedPrivilege.isContainedIn(existingPrivilege))) { + if (requestedPrivilege.isContainedIn(existingPrivilege)) { iterator.remove(); } else if (existingPrivilege.isContainedIn(requestedPrivilege)) { @@ -1644,12 +1505,12 @@ public void sendTransactionHeartbeat(long transactionId) try { retry() .stopOnIllegalExceptions() - .run("sendTransactionHeartbeat", (() -> { + .run("sendTransactionHeartbeat", () -> { try (ThriftMetastoreClient metastoreClient = createMetastoreClient()) { metastoreClient.sendTransactionHeartbeat(transactionId); } return null; - })); + }); } catch (TException e) { throw new TrinoException(HIVE_METASTORE_ERROR, e); @@ -1826,7 +1687,7 @@ private static LockComponent createLockComponentForOperation(SchemaTableName tab .setOperationType(operation) .setDbname(table.getSchemaName()) .setTablename(table.getTableName()) - // acquire locks is called only for TransactionalTable + // acquire lock is called only for TransactionalTable .setIsTransactional(true) .setIsDynamicPartitionWrite(isDynamicPartitionWrite) .setLevel(LockLevel.TABLE); @@ -1949,29 +1810,7 @@ public void updateTableWriteId(String dbName, String tableName, long transaction } @Override - public void alterPartitions(String dbName, String tableName, List partitions, long writeId) - { - checkArgument(writeId > 0, "writeId should be a positive integer, but was %s", writeId); - try { - retry() - .stopOnIllegalExceptions() - .run("alterPartitions", stats.getAlterPartitions().wrap(() -> { - try (ThriftMetastoreClient metastoreClient = createMetastoreClient()) { - metastoreClient.alterPartitions(dbName, tableName, partitions, writeId); - } - return null; - })); - } - catch (TException e) { - throw new TrinoException(HIVE_METASTORE_ERROR, e); - } - catch (Exception e) { - throw propagate(e); - } - } - - @Override - public void addDynamicPartitions(String dbName, String tableName, List partitionNames, long transactionId, long writeId, AcidOperation operation) + public void addDynamicPartitions(String dbName, String tableName, List partitionNames, long transactionId, long writeId, DataOperationType operation) { checkArgument(writeId > 0, "writeId should be a positive integer, but was %s", writeId); requireNonNull(partitionNames, "partitionNames is null"); @@ -2052,7 +1891,7 @@ private static void validateObjectName(String objectName) throw new TrinoException(GENERIC_USER_ERROR, format("Invalid object name: '%s'", objectName)); } if (objectName.contains("/")) { - // Older HMS instances may allow names like 'foo/bar' which can cause managed tables to be + // Older HMS instances may allow names like 'foo/bar', which can cause managed tables to be // saved in a different location than its intended schema directory throw new TrinoException(GENERIC_USER_ERROR, format("Invalid object name: '%s'", objectName)); } diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/ThriftHiveMetastoreClient.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/ThriftHiveMetastoreClient.java index 0df715380dd0..fc179fc60c77 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/ThriftHiveMetastoreClient.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/ThriftHiveMetastoreClient.java @@ -29,11 +29,10 @@ import io.trino.hive.thrift.metastore.ColumnStatisticsDesc; import io.trino.hive.thrift.metastore.ColumnStatisticsObj; import io.trino.hive.thrift.metastore.CommitTxnRequest; +import io.trino.hive.thrift.metastore.DataOperationType; import io.trino.hive.thrift.metastore.Database; import io.trino.hive.thrift.metastore.EnvironmentContext; import io.trino.hive.thrift.metastore.FieldSchema; -import io.trino.hive.thrift.metastore.GetPrincipalsInRoleRequest; -import io.trino.hive.thrift.metastore.GetPrincipalsInRoleResponse; import io.trino.hive.thrift.metastore.GetRoleGrantsForPrincipalRequest; import io.trino.hive.thrift.metastore.GetRoleGrantsForPrincipalResponse; import io.trino.hive.thrift.metastore.GetTableRequest; @@ -63,9 +62,10 @@ import io.trino.hive.thrift.metastore.TxnToWriteId; import io.trino.hive.thrift.metastore.UnlockRequest; import io.trino.plugin.base.util.LoggingInvocationHandler; -import io.trino.plugin.hive.acid.AcidOperation; import io.trino.plugin.hive.metastore.thrift.MetastoreSupportsDateStatistics.DateStatisticsSupport; +import io.trino.spi.connector.RelationType; import io.trino.spi.connector.SchemaTableName; +import jakarta.annotation.Nullable; import org.apache.thrift.TApplicationException; import org.apache.thrift.TException; import org.apache.thrift.protocol.TBinaryProtocol; @@ -73,9 +73,11 @@ import org.apache.thrift.transport.TTransportException; import java.util.ArrayList; +import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.Optional; +import java.util.Set; import java.util.concurrent.atomic.AtomicInteger; import java.util.function.Predicate; import java.util.regex.Pattern; @@ -91,7 +93,7 @@ import static io.trino.hive.thrift.metastore.GrantRevokeType.REVOKE; import static io.trino.hive.thrift.metastore.hive_metastoreConstants.HIVE_FILTER_FIELD_PARAMS; import static io.trino.plugin.hive.TableType.VIRTUAL_VIEW; -import static io.trino.plugin.hive.ViewReaderUtil.PRESTO_VIEW_FLAG; +import static io.trino.plugin.hive.metastore.TableInfo.PRESTO_VIEW_COMMENT; import static io.trino.plugin.hive.metastore.thrift.MetastoreSupportsDateStatistics.DateStatisticsSupport.NOT_SUPPORTED; import static io.trino.plugin.hive.metastore.thrift.MetastoreSupportsDateStatistics.DateStatisticsSupport.SUPPORTED; import static io.trino.plugin.hive.metastore.thrift.MetastoreSupportsDateStatistics.DateStatisticsSupport.UNKNOWN; @@ -99,6 +101,10 @@ import static io.trino.plugin.hive.metastore.thrift.TxnUtils.createValidTxnWriteIdList; import static java.lang.String.format; import static java.util.Objects.requireNonNull; +import static java.util.stream.Collectors.joining; +import static org.apache.hadoop.hive.metastore.utils.MetaStoreUtils.CATALOG_DB_SEPARATOR; +import static org.apache.hadoop.hive.metastore.utils.MetaStoreUtils.CATALOG_DB_THRIFT_NAME_MARKER; +import static org.apache.hadoop.hive.metastore.utils.MetaStoreUtils.DB_EMPTY_MARKER; import static org.apache.thrift.TApplicationException.UNKNOWN_METHOD; public class ThriftHiveMetastoreClient @@ -115,23 +121,21 @@ public class ThriftHiveMetastoreClient private final String hostname; private final MetastoreSupportsDateStatistics metastoreSupportsDateStatistics; + private final boolean metastoreSupportsTableMeta; private final AtomicInteger chosenGetTableAlternative; private final AtomicInteger chosenTableParamAlternative; - private final AtomicInteger chosenGetAllTablesAlternative; - private final AtomicInteger chosenGetAllViewsPerDatabaseAlternative; - private final AtomicInteger chosenGetAllViewsAlternative; private final AtomicInteger chosenAlterTransactionalTableAlternative; private final AtomicInteger chosenAlterPartitionsAlternative; + private final Optional catalogName; public ThriftHiveMetastoreClient( TransportSupplier transportSupplier, String hostname, + Optional catalogName, MetastoreSupportsDateStatistics metastoreSupportsDateStatistics, + boolean metastoreSupportsTableMeta, AtomicInteger chosenGetTableAlternative, AtomicInteger chosenTableParamAlternative, - AtomicInteger chosenGetAllTablesAlternative, - AtomicInteger chosenGetAllViewsPerDatabaseAlternative, - AtomicInteger chosenGetAllViewsAlternative, AtomicInteger chosenAlterTransactionalTableAlternative, AtomicInteger chosenAlterPartitionsAlternative) throws TTransportException @@ -139,13 +143,12 @@ public ThriftHiveMetastoreClient( this.transportSupplier = requireNonNull(transportSupplier, "transportSupplier is null"); this.hostname = requireNonNull(hostname, "hostname is null"); this.metastoreSupportsDateStatistics = requireNonNull(metastoreSupportsDateStatistics, "metastoreSupportsDateStatistics is null"); + this.metastoreSupportsTableMeta = metastoreSupportsTableMeta; this.chosenGetTableAlternative = requireNonNull(chosenGetTableAlternative, "chosenGetTableAlternative is null"); this.chosenTableParamAlternative = requireNonNull(chosenTableParamAlternative, "chosenTableParamAlternative is null"); - this.chosenGetAllViewsPerDatabaseAlternative = requireNonNull(chosenGetAllViewsPerDatabaseAlternative, "chosenGetAllViewsPerDatabaseAlternative is null"); this.chosenAlterTransactionalTableAlternative = requireNonNull(chosenAlterTransactionalTableAlternative, "chosenAlterTransactionalTableAlternative is null"); this.chosenAlterPartitionsAlternative = requireNonNull(chosenAlterPartitionsAlternative, "chosenAlterPartitionsAlternative is null"); - this.chosenGetAllTablesAlternative = requireNonNull(chosenGetAllTablesAlternative, "chosenGetAllTablesAlternative is null"); - this.chosenGetAllViewsAlternative = requireNonNull(chosenGetAllViewsAlternative, "chosenGetAllViewsAlternative is null"); + this.catalogName = requireNonNull(catalogName, "catalogName is null"); connect(); } @@ -187,56 +190,34 @@ public Database getDatabase(String dbName) } @Override - public List getAllTables(String databaseName) + public List getTableMeta(String databaseName) throws TException { - return client.getAllTables(databaseName); - } - - @Override - public Optional> getAllTables() - throws TException - { - return alternativeCall( - exception -> !isUnknownMethodExceptionalResponse(exception), - chosenGetAllTablesAlternative, - // Empty table types argument (the 3rd one) means all types of tables - () -> getSchemaTableNames(client.getTableMeta("*", "*", ImmutableList.of())), - Optional::empty); - } - - @Override - public List getAllViews(String databaseName) - throws TException - { - return alternativeCall( - exception -> !isUnknownMethodExceptionalResponse(exception), - chosenGetAllViewsPerDatabaseAlternative, - () -> client.getTablesByType(databaseName, ".*", VIRTUAL_VIEW.name()), - // fallback to enumerating Presto views only (Hive views can still be executed, but will be listed as tables and not views) - () -> getTablesWithParameter(databaseName, PRESTO_VIEW_FLAG, "true")); - } - - @Override - public Optional> getAllViews() - throws TException - { - return alternativeCall( - exception -> !isUnknownMethodExceptionalResponse(exception), - chosenGetAllViewsAlternative, - () -> getSchemaTableNames(client.getTableMeta("*", "*", ImmutableList.of(VIRTUAL_VIEW.name()))), - Optional::empty); - } + // TODO: remove this once Unity adds support for getTableMeta + if (!metastoreSupportsTableMeta) { + String catalogDatabaseName = prependCatalogToDbName(catalogName, databaseName); + Map tables = new HashMap<>(); + client.getTables(catalogDatabaseName, ".*").forEach(name -> tables.put(name, new TableMeta(databaseName, name, RelationType.TABLE.toString()))); + client.getTablesByType(catalogDatabaseName, ".*", VIRTUAL_VIEW.name()).forEach(name -> { + TableMeta tableMeta = new TableMeta(databaseName, name, VIRTUAL_VIEW.name()); + // This makes all views look like a Trino view, so that they are not filtered out during SHOW VIEWS + tableMeta.setComments(PRESTO_VIEW_COMMENT); + tables.put(name, tableMeta); + }); + return ImmutableList.copyOf(tables.values()); + } - private static Optional> getSchemaTableNames(List tablesMetadata) - { - return Optional.of(tablesMetadata.stream() - .map(metadata -> new SchemaTableName(metadata.getDbName(), metadata.getTableName())) - .collect(toImmutableList())); + if (databaseName.indexOf('*') >= 0 || databaseName.indexOf('|') >= 0) { + // in this case we replace any pipes with a glob and then filter the output + return client.getTableMeta(prependCatalogToDbName(catalogName, databaseName.replace('|', '*')), "*", ImmutableList.of()).stream() + .filter(tableMeta -> tableMeta.getDbName().equals(databaseName)) + .collect(toImmutableList()); + } + return client.getTableMeta(prependCatalogToDbName(catalogName, databaseName), "*", ImmutableList.of()); } @Override - public List getTablesWithParameter(String databaseName, String parameterKey, String parameterValue) + public List getTableNamesWithParameters(String databaseName, String parameterKey, Set parameterValues) throws TException { checkArgument(TABLE_PARAMETER_SAFE_KEY_PATTERN.matcher(parameterKey).matches(), "Parameter key contains invalid characters: '%s'", parameterKey); @@ -247,14 +228,21 @@ public List getTablesWithParameter(String databaseName, String parameter * HMS's behavior from outside. Also, by restricting parameter values, we avoid the problem * of how to quote them when passing within the filter string. */ - checkArgument(TABLE_PARAMETER_SAFE_VALUE_PATTERN.matcher(parameterValue).matches(), "Parameter value contains invalid characters: '%s'", parameterValue); + for (String parameterValue : parameterValues) { + checkArgument(TABLE_PARAMETER_SAFE_VALUE_PATTERN.matcher(parameterValue).matches(), "Parameter value contains invalid characters: '%s'", parameterValue); + } /* * Thrift call `get_table_names_by_filter` may be translated by Metastore to an SQL query against Metastore database. * Hive 2.3 on some databases uses CLOB for table parameter value column and some databases disallow `=` predicate over * CLOB values. At the same time, they allow `LIKE` predicates over them. */ - String filterWithEquals = HIVE_FILTER_FIELD_PARAMS + parameterKey + " = \"" + parameterValue + "\""; - String filterWithLike = HIVE_FILTER_FIELD_PARAMS + parameterKey + " LIKE \"" + parameterValue + "\""; + String filterWithEquals = parameterValues.stream() + .map(parameterValue -> HIVE_FILTER_FIELD_PARAMS + parameterKey + " = \"" + parameterValue + "\"") + .collect(joining(" or ")); + + String filterWithLike = parameterValues.stream() + .map(parameterValue -> HIVE_FILTER_FIELD_PARAMS + parameterKey + " LIKE \"" + parameterValue + "\"") + .collect(joining(" or ")); return alternativeCall( ThriftHiveMetastoreClient::defaultIsValidExceptionalResponse, @@ -263,6 +251,13 @@ public List getTablesWithParameter(String databaseName, String parameter () -> client.getTableNamesByFilter(databaseName, filterWithLike, (short) -1)); } + private static Optional> getSchemaTableNames(List tablesMetadata) + { + return Optional.of(tablesMetadata.stream() + .map(metadata -> new SchemaTableName(metadata.getDbName(), metadata.getTableName())) + .collect(toImmutableList())); + } + @Override public void createDatabase(Database database) throws TException @@ -610,15 +605,6 @@ private void removeGrant(String role, String granteeName, PrincipalType granteeT } } - @Override - public List listGrantedPrincipals(String role) - throws TException - { - GetPrincipalsInRoleRequest request = new GetPrincipalsInRoleRequest(role); - GetPrincipalsInRoleResponse response = client.getPrincipalsInRole(request); - return ImmutableList.copyOf(response.getPrincipalGrants()); - } - @Override public List listRoleGrants(String principalName, PrincipalType principalType) throws TException @@ -743,11 +729,11 @@ public void alterPartitions(String dbName, String tableName, List par } @Override - public void addDynamicPartitions(String dbName, String tableName, List partitionNames, long transactionId, long writeId, AcidOperation operation) + public void addDynamicPartitions(String dbName, String tableName, List partitionNames, long transactionId, long writeId, DataOperationType operation) throws TException { AddDynamicPartitions request = new AddDynamicPartitions(transactionId, writeId, dbName, tableName, partitionNames); - request.setOperationType(operation.getMetastoreOperationType()); + request.setOperationType(operation); client.addDynamicPartitions(request); } @@ -879,4 +865,33 @@ public interface TransportSupplier TTransport createTransport() throws TTransportException; } + + /** + * To construct a pattern using database and catalog name that Hive Thrift Server. + * Based on the Hive's implementation. + * + * @param catalogName hive catalog name + * @param databaseName database name + * @return string pattern that Hive Thrift Server understands + */ + private static String prependCatalogToDbName(Optional catalogName, @Nullable String databaseName) + { + if (catalogName.isEmpty()) { + return databaseName; + } + + StringBuilder catalogDatabaseName = new StringBuilder() + .append(CATALOG_DB_THRIFT_NAME_MARKER) + .append(catalogName.orElseThrow()) + .append(CATALOG_DB_SEPARATOR); + if (databaseName != null) { + if (databaseName.isEmpty()) { + catalogDatabaseName.append(DB_EMPTY_MARKER); + } + else { + catalogDatabaseName.append(databaseName); + } + } + return catalogDatabaseName.toString(); + } } diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/ThriftHiveMetastoreFactory.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/ThriftHiveMetastoreFactory.java index e0b300226542..58d8f9e14790 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/ThriftHiveMetastoreFactory.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/ThriftHiveMetastoreFactory.java @@ -15,7 +15,7 @@ import com.google.inject.Inject; import io.airlift.units.Duration; -import io.trino.hdfs.HdfsEnvironment; +import io.trino.filesystem.TrinoFileSystemFactory; import io.trino.plugin.hive.HideDeltaLakeTables; import io.trino.spi.security.ConnectorIdentity; import org.weakref.jmx.Flatten; @@ -30,7 +30,7 @@ public class ThriftHiveMetastoreFactory implements ThriftMetastoreFactory { - private final HdfsEnvironment hdfsEnvironment; + private final TrinoFileSystemFactory fileSystemFactory; private final IdentityAwareMetastoreClientFactory metastoreClientFactory; private final double backoffScaleFactor; private final Duration minBackoffDelay; @@ -40,7 +40,6 @@ public class ThriftHiveMetastoreFactory private final int maxRetries; private final boolean impersonationEnabled; private final boolean deleteFilesOnDrop; - private final boolean translateHiveViews; private final boolean assumeCanonicalPartitionKeys; private final boolean useSparkTableStatisticsFallback; private final ExecutorService writeStatisticsExecutor; @@ -50,13 +49,12 @@ public class ThriftHiveMetastoreFactory public ThriftHiveMetastoreFactory( IdentityAwareMetastoreClientFactory metastoreClientFactory, @HideDeltaLakeTables boolean hideDeltaLakeTables, - @TranslateHiveViews boolean translateHiveViews, ThriftMetastoreConfig thriftConfig, - HdfsEnvironment hdfsEnvironment, + TrinoFileSystemFactory fileSystemFactory, @ThriftHiveWriteStatisticsExecutor ExecutorService writeStatisticsExecutor) { this.metastoreClientFactory = requireNonNull(metastoreClientFactory, "metastoreClientFactory is null"); - this.hdfsEnvironment = requireNonNull(hdfsEnvironment, "hdfsEnvironment is null"); + this.fileSystemFactory = requireNonNull(fileSystemFactory, "fileSystemFactory is null"); this.backoffScaleFactor = thriftConfig.getBackoffScaleFactor(); this.minBackoffDelay = thriftConfig.getMinBackoffDelay(); this.maxBackoffDelay = thriftConfig.getMaxBackoffDelay(); @@ -64,7 +62,6 @@ public ThriftHiveMetastoreFactory( this.maxRetries = thriftConfig.getMaxRetries(); this.impersonationEnabled = thriftConfig.isImpersonationEnabled(); this.deleteFilesOnDrop = thriftConfig.isDeleteFilesOnDrop(); - this.translateHiveViews = translateHiveViews; checkArgument(!hideDeltaLakeTables, "Hiding Delta Lake tables is not supported"); // TODO this.maxWaitForLock = thriftConfig.getMaxWaitForTransactionLock(); @@ -91,7 +88,7 @@ public ThriftMetastore createMetastore(Optional identity) { return new ThriftHiveMetastore( identity, - hdfsEnvironment, + fileSystemFactory, metastoreClientFactory, backoffScaleFactor, minBackoffDelay, @@ -100,7 +97,6 @@ public ThriftMetastore createMetastore(Optional identity) maxWaitForLock, maxRetries, deleteFilesOnDrop, - translateHiveViews, assumeCanonicalPartitionKeys, useSparkTableStatisticsFallback, stats, diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/ThriftMetastore.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/ThriftMetastore.java index 57bf2f1e9288..4aa8f05869b4 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/ThriftMetastore.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/ThriftMetastore.java @@ -18,34 +18,33 @@ import io.trino.hive.thrift.metastore.FieldSchema; import io.trino.hive.thrift.metastore.Partition; import io.trino.hive.thrift.metastore.Table; -import io.trino.plugin.hive.HiveColumnStatisticType; +import io.trino.hive.thrift.metastore.TableMeta; import io.trino.plugin.hive.HivePartition; import io.trino.plugin.hive.PartitionStatistics; -import io.trino.plugin.hive.acid.AcidOperation; -import io.trino.plugin.hive.acid.AcidTransaction; import io.trino.plugin.hive.metastore.AcidTransactionOwner; +import io.trino.plugin.hive.metastore.HiveColumnStatistics; import io.trino.plugin.hive.metastore.HivePrincipal; import io.trino.plugin.hive.metastore.HivePrivilegeInfo; import io.trino.plugin.hive.metastore.HivePrivilegeInfo.HivePrivilege; import io.trino.plugin.hive.metastore.PartitionWithStatistics; +import io.trino.plugin.hive.metastore.StatisticsUpdateMode; import io.trino.spi.TrinoException; import io.trino.spi.connector.SchemaTableName; import io.trino.spi.connector.TableNotFoundException; import io.trino.spi.predicate.TupleDomain; import io.trino.spi.security.RoleGrant; -import io.trino.spi.type.Type; import java.util.List; import java.util.Map; import java.util.Optional; import java.util.OptionalLong; import java.util.Set; -import java.util.function.Function; import static io.trino.plugin.hive.HiveErrorCode.HIVE_INVALID_METADATA; import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED; -public interface ThriftMetastore +public sealed interface ThriftMetastore + permits ThriftHiveMetastore { void createDatabase(Database database); @@ -57,21 +56,15 @@ public interface ThriftMetastore void dropTable(String databaseName, String tableName, boolean deleteData); - void alterTable(String databaseName, String tableName, Table table); + void alterTable(String databaseName, String tableName, Table table, Map environmentContext); void alterTransactionalTable(Table table, long transactionId, long writeId); List getAllDatabases(); - List getAllTables(String databaseName); + List getTables(String databaseName); - Optional> getAllTables(); - - List getTablesWithParameter(String databaseName, String parameterKey, String parameterValue); - - List getAllViews(String databaseName); - - Optional> getAllViews(); + List getTableNamesWithParameters(String databaseName, String parameterKey, Set parameterValues); Optional getDatabase(String databaseName); @@ -89,15 +82,15 @@ public interface ThriftMetastore Optional
getTable(String databaseName, String tableName); - Set getSupportedColumnStatistics(Type type); + Map getTableColumnStatistics(String databaseName, String tableName, Set columnNames); - PartitionStatistics getTableStatistics(Table table); + Map> getPartitionColumnStatistics(String databaseName, String tableName, Set partitionNames, Set columnNames); - Map getPartitionStatistics(Table table, List partitions); + boolean useSparkTableStatistics(); - void updateTableStatistics(String databaseName, String tableName, AcidTransaction transaction, Function update); + void updateTableStatistics(String databaseName, String tableName, OptionalLong acidWriteId, StatisticsUpdateMode mode, PartitionStatistics statisticsUpdate); - void updatePartitionStatistics(Table table, String partitionName, Function update); + void updatePartitionStatistics(Table table, String partitionName, StatisticsUpdateMode mode, PartitionStatistics statisticsUpdate); void createRole(String role, String grantor); @@ -109,8 +102,6 @@ public interface ThriftMetastore void revokeRoles(Set roles, Set grantees, boolean adminOption, HivePrincipal grantor); - Set listGrantedPrincipals(String role); - Set listRoleGrants(HivePrincipal principal); void grantTablePrivileges(String databaseName, String tableName, String tableOwner, HivePrincipal grantee, HivePrincipal grantor, Set privileges, boolean grantOption); @@ -215,12 +206,7 @@ default void updateTableWriteId(String dbName, String tableName, long transactio throw new UnsupportedOperationException(); } - default void alterPartitions(String dbName, String tableName, List partitions, long writeId) - { - throw new UnsupportedOperationException(); - } - - default void addDynamicPartitions(String dbName, String tableName, List partitionNames, long transactionId, long writeId, AcidOperation operation) + default void addDynamicPartitions(String dbName, String tableName, List partitionNames, long transactionId, long writeId, DataOperationType operation) { throw new UnsupportedOperationException(); } diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/ThriftMetastoreClient.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/ThriftMetastoreClient.java index a66be01a404d..ddbe6502368b 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/ThriftMetastoreClient.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/ThriftMetastoreClient.java @@ -14,6 +14,7 @@ package io.trino.plugin.hive.metastore.thrift; import io.trino.hive.thrift.metastore.ColumnStatisticsObj; +import io.trino.hive.thrift.metastore.DataOperationType; import io.trino.hive.thrift.metastore.Database; import io.trino.hive.thrift.metastore.EnvironmentContext; import io.trino.hive.thrift.metastore.FieldSchema; @@ -27,15 +28,14 @@ import io.trino.hive.thrift.metastore.Role; import io.trino.hive.thrift.metastore.RolePrincipalGrant; import io.trino.hive.thrift.metastore.Table; +import io.trino.hive.thrift.metastore.TableMeta; import io.trino.hive.thrift.metastore.TxnToWriteId; -import io.trino.plugin.hive.acid.AcidOperation; -import io.trino.spi.connector.SchemaTableName; import org.apache.thrift.TException; import java.io.Closeable; import java.util.List; import java.util.Map; -import java.util.Optional; +import java.util.Set; public interface ThriftMetastoreClient extends Closeable @@ -49,19 +49,10 @@ List getAllDatabases() Database getDatabase(String databaseName) throws TException; - List getAllTables(String databaseName) + List getTableMeta(String databaseName) throws TException; - Optional> getAllTables() - throws TException; - - List getAllViews(String databaseName) - throws TException; - - Optional> getAllViews() - throws TException; - - List getTablesWithParameter(String databaseName, String parameterKey, String parameterValue) + List getTableNamesWithParameters(String databaseName, String parameterKey, Set parameterValues) throws TException; void createDatabase(Database database) @@ -154,9 +145,6 @@ void grantRole(String role, String granteeName, PrincipalType granteeType, Strin void revokeRole(String role, String granteeName, PrincipalType granteeType, boolean grantOption) throws TException; - List listGrantedPrincipals(String role) - throws TException; - List listRoleGrants(String name, PrincipalType principalType) throws TException; @@ -205,7 +193,7 @@ default List allocateTableWriteIds(String database, String tableNa void alterPartitions(String dbName, String tableName, List partitions, long writeId) throws TException; - void addDynamicPartitions(String dbName, String tableName, List partitionNames, long transactionId, long writeId, AcidOperation operation) + void addDynamicPartitions(String dbName, String tableName, List partitionNames, long transactionId, long writeId, DataOperationType operation) throws TException; void alterTransactionalTable(Table table, long transactionId, long writeId, EnvironmentContext context) diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/ThriftMetastoreClientFactory.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/ThriftMetastoreClientFactory.java index d2c16b955a17..694f29b0dcb7 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/ThriftMetastoreClientFactory.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/ThriftMetastoreClientFactory.java @@ -13,13 +13,13 @@ */ package io.trino.plugin.hive.metastore.thrift; -import com.google.common.net.HostAndPort; import org.apache.thrift.transport.TTransportException; +import java.net.URI; import java.util.Optional; public interface ThriftMetastoreClientFactory { - ThriftMetastoreClient create(HostAndPort address, Optional delegationToken) + ThriftMetastoreClient create(URI address, Optional delegationToken) throws TTransportException; } diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/ThriftMetastoreConfig.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/ThriftMetastoreConfig.java index c53a59c04032..62edfbb9f975 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/ThriftMetastoreConfig.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/ThriftMetastoreConfig.java @@ -26,11 +26,13 @@ import jakarta.validation.constraints.NotNull; import java.io.File; +import java.util.Optional; import java.util.concurrent.TimeUnit; public class ThriftMetastoreConfig { - private Duration metastoreTimeout = new Duration(10, TimeUnit.SECONDS); + private Duration connectTimeout = new Duration(10, TimeUnit.SECONDS); + private Duration readTimeout = new Duration(10, TimeUnit.SECONDS); private HostAndPort socksProxy; private int maxRetries = RetryDriver.DEFAULT_MAX_ATTEMPTS - 1; private double backoffScaleFactor = RetryDriver.DEFAULT_SCALE_FACTOR; @@ -43,6 +45,7 @@ public class ThriftMetastoreConfig private long delegationTokenCacheMaximumSize = 1000; private boolean deleteFilesOnDrop; private Duration maxWaitForTransactionLock = new Duration(10, TimeUnit.MINUTES); + private String catalogName; private boolean tlsEnabled; private File keystorePath; @@ -53,15 +56,32 @@ public class ThriftMetastoreConfig private int writeStatisticsThreads = 20; @NotNull - public Duration getMetastoreTimeout() + public Duration getConnectTimeout() { - return metastoreTimeout; + return connectTimeout; } - @Config("hive.metastore-timeout") - public ThriftMetastoreConfig setMetastoreTimeout(Duration metastoreTimeout) + @Config("hive.metastore.thrift.client.connect-timeout") + @LegacyConfig("hive.metastore-timeout") + @ConfigDescription("Socket connect timeout for metastore client") + public ThriftMetastoreConfig setConnectTimeout(Duration connectTimeout) { - this.metastoreTimeout = metastoreTimeout; + this.connectTimeout = connectTimeout; + return this; + } + + @NotNull + public Duration getReadTimeout() + { + return readTimeout; + } + + @Config("hive.metastore.thrift.client.read-timeout") + @LegacyConfig("hive.metastore-timeout") + @ConfigDescription("Socket read timeout for metastore client") + public ThriftMetastoreConfig setReadTimeout(Duration readTimeout) + { + this.readTimeout = readTimeout; return this; } @@ -325,4 +345,17 @@ public ThriftMetastoreConfig setWriteStatisticsThreads(int writeStatisticsThread this.writeStatisticsThreads = writeStatisticsThreads; return this; } + + public Optional getCatalogName() + { + return Optional.ofNullable(catalogName); + } + + @Config("hive.metastore.thrift.catalog-name") + @ConfigDescription("Hive metastore thrift catalog name") + public ThriftMetastoreConfig setCatalogName(String catalogName) + { + this.catalogName = catalogName; + return this; + } } diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/ThriftMetastoreUtil.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/ThriftMetastoreUtil.java index 00723cde574f..46187e68158f 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/ThriftMetastoreUtil.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/ThriftMetastoreUtil.java @@ -21,6 +21,7 @@ import io.trino.hive.thrift.metastore.BinaryColumnStatsData; import io.trino.hive.thrift.metastore.BooleanColumnStatsData; import io.trino.hive.thrift.metastore.ColumnStatisticsObj; +import io.trino.hive.thrift.metastore.DataOperationType; import io.trino.hive.thrift.metastore.Date; import io.trino.hive.thrift.metastore.DateColumnStatsData; import io.trino.hive.thrift.metastore.Decimal; @@ -39,6 +40,7 @@ import io.trino.plugin.hive.HiveBucketProperty; import io.trino.plugin.hive.HiveColumnStatisticType; import io.trino.plugin.hive.HiveType; +import io.trino.plugin.hive.acid.AcidOperation; import io.trino.plugin.hive.metastore.Column; import io.trino.plugin.hive.metastore.Database; import io.trino.plugin.hive.metastore.HiveColumnStatistics; @@ -486,39 +488,39 @@ public static Partition fromMetastoreApiPartition(io.trino.hive.thrift.metastore return partitionBuilder.build(); } - public static HiveColumnStatistics fromMetastoreApiColumnStatistics(ColumnStatisticsObj columnStatistics, OptionalLong rowCount) + public static HiveColumnStatistics fromMetastoreApiColumnStatistics(ColumnStatisticsObj columnStatistics) { if (columnStatistics.getStatsData().isSetLongStats()) { LongColumnStatsData longStatsData = columnStatistics.getStatsData().getLongStats(); OptionalLong min = longStatsData.isSetLowValue() ? OptionalLong.of(longStatsData.getLowValue()) : OptionalLong.empty(); OptionalLong max = longStatsData.isSetHighValue() ? OptionalLong.of(longStatsData.getHighValue()) : OptionalLong.empty(); OptionalLong nullsCount = longStatsData.isSetNumNulls() ? fromMetastoreNullsCount(longStatsData.getNumNulls()) : OptionalLong.empty(); - OptionalLong distinctValuesCount = longStatsData.isSetNumDVs() ? OptionalLong.of(longStatsData.getNumDVs()) : OptionalLong.empty(); - return createIntegerColumnStatistics(min, max, nullsCount, fromMetastoreDistinctValuesCount(distinctValuesCount, nullsCount, rowCount)); + OptionalLong distinctValuesWithNullCount = longStatsData.isSetNumDVs() ? OptionalLong.of(longStatsData.getNumDVs()) : OptionalLong.empty(); + return createIntegerColumnStatistics(min, max, nullsCount, distinctValuesWithNullCount); } if (columnStatistics.getStatsData().isSetDoubleStats()) { DoubleColumnStatsData doubleStatsData = columnStatistics.getStatsData().getDoubleStats(); OptionalDouble min = doubleStatsData.isSetLowValue() ? OptionalDouble.of(doubleStatsData.getLowValue()) : OptionalDouble.empty(); OptionalDouble max = doubleStatsData.isSetHighValue() ? OptionalDouble.of(doubleStatsData.getHighValue()) : OptionalDouble.empty(); OptionalLong nullsCount = doubleStatsData.isSetNumNulls() ? fromMetastoreNullsCount(doubleStatsData.getNumNulls()) : OptionalLong.empty(); - OptionalLong distinctValuesCount = doubleStatsData.isSetNumDVs() ? OptionalLong.of(doubleStatsData.getNumDVs()) : OptionalLong.empty(); - return createDoubleColumnStatistics(min, max, nullsCount, fromMetastoreDistinctValuesCount(distinctValuesCount, nullsCount, rowCount)); + OptionalLong distinctValuesWithNullCount = doubleStatsData.isSetNumDVs() ? OptionalLong.of(doubleStatsData.getNumDVs()) : OptionalLong.empty(); + return createDoubleColumnStatistics(min, max, nullsCount, distinctValuesWithNullCount); } if (columnStatistics.getStatsData().isSetDecimalStats()) { DecimalColumnStatsData decimalStatsData = columnStatistics.getStatsData().getDecimalStats(); Optional min = decimalStatsData.isSetLowValue() ? fromMetastoreDecimal(decimalStatsData.getLowValue()) : Optional.empty(); Optional max = decimalStatsData.isSetHighValue() ? fromMetastoreDecimal(decimalStatsData.getHighValue()) : Optional.empty(); OptionalLong nullsCount = decimalStatsData.isSetNumNulls() ? fromMetastoreNullsCount(decimalStatsData.getNumNulls()) : OptionalLong.empty(); - OptionalLong distinctValuesCount = decimalStatsData.isSetNumDVs() ? OptionalLong.of(decimalStatsData.getNumDVs()) : OptionalLong.empty(); - return createDecimalColumnStatistics(min, max, nullsCount, fromMetastoreDistinctValuesCount(distinctValuesCount, nullsCount, rowCount)); + OptionalLong distinctValuesWithNullCount = decimalStatsData.isSetNumDVs() ? OptionalLong.of(decimalStatsData.getNumDVs()) : OptionalLong.empty(); + return createDecimalColumnStatistics(min, max, nullsCount, distinctValuesWithNullCount); } if (columnStatistics.getStatsData().isSetDateStats()) { DateColumnStatsData dateStatsData = columnStatistics.getStatsData().getDateStats(); Optional min = dateStatsData.isSetLowValue() ? fromMetastoreDate(dateStatsData.getLowValue()) : Optional.empty(); Optional max = dateStatsData.isSetHighValue() ? fromMetastoreDate(dateStatsData.getHighValue()) : Optional.empty(); OptionalLong nullsCount = dateStatsData.isSetNumNulls() ? fromMetastoreNullsCount(dateStatsData.getNumNulls()) : OptionalLong.empty(); - OptionalLong distinctValuesCount = dateStatsData.isSetNumDVs() ? OptionalLong.of(dateStatsData.getNumDVs()) : OptionalLong.empty(); - return createDateColumnStatistics(min, max, nullsCount, fromMetastoreDistinctValuesCount(distinctValuesCount, nullsCount, rowCount)); + OptionalLong distinctValuesWithNullCount = dateStatsData.isSetNumDVs() ? OptionalLong.of(dateStatsData.getNumDVs()) : OptionalLong.empty(); + return createDateColumnStatistics(min, max, nullsCount, distinctValuesWithNullCount); } if (columnStatistics.getStatsData().isSetBooleanStats()) { BooleanColumnStatsData booleanStatsData = columnStatistics.getStatsData().getBooleanStats(); @@ -539,12 +541,12 @@ public static HiveColumnStatistics fromMetastoreApiColumnStatistics(ColumnStatis OptionalLong maxColumnLength = stringStatsData.isSetMaxColLen() ? OptionalLong.of(stringStatsData.getMaxColLen()) : OptionalLong.empty(); OptionalDouble averageColumnLength = stringStatsData.isSetAvgColLen() ? OptionalDouble.of(stringStatsData.getAvgColLen()) : OptionalDouble.empty(); OptionalLong nullsCount = stringStatsData.isSetNumNulls() ? fromMetastoreNullsCount(stringStatsData.getNumNulls()) : OptionalLong.empty(); - OptionalLong distinctValuesCount = stringStatsData.isSetNumDVs() ? OptionalLong.of(stringStatsData.getNumDVs()) : OptionalLong.empty(); + OptionalLong distinctValuesWithNullCount = stringStatsData.isSetNumDVs() ? OptionalLong.of(stringStatsData.getNumDVs()) : OptionalLong.empty(); return createStringColumnStatistics( maxColumnLength, - getTotalSizeInBytes(averageColumnLength, rowCount, nullsCount), + averageColumnLength, nullsCount, - fromMetastoreDistinctValuesCount(distinctValuesCount, nullsCount, rowCount)); + distinctValuesWithNullCount); } if (columnStatistics.getStatsData().isSetBinaryStats()) { BinaryColumnStatsData binaryStatsData = columnStatistics.getStatsData().getBinaryStats(); @@ -553,7 +555,7 @@ public static HiveColumnStatistics fromMetastoreApiColumnStatistics(ColumnStatis OptionalLong nullsCount = binaryStatsData.isSetNumNulls() ? fromMetastoreNullsCount(binaryStatsData.getNumNulls()) : OptionalLong.empty(); return createBinaryColumnStatistics( maxColumnLength, - getTotalSizeInBytes(averageColumnLength, rowCount, nullsCount), + averageColumnLength, nullsCount); } throw new TrinoException(HIVE_INVALID_METADATA, "Invalid column statistics data: " + columnStatistics); @@ -671,7 +673,7 @@ public static FieldSchema toMetastoreApiFieldSchema(Column column) private static Column fromMetastoreApiFieldSchema(FieldSchema fieldSchema) { - return new Column(fieldSchema.getName(), HiveType.valueOf(fieldSchema.getType()), Optional.ofNullable(fieldSchema.getComment())); + return new Column(fieldSchema.getName(), HiveType.valueOf(fieldSchema.getType()), Optional.ofNullable(fieldSchema.getComment()), ImmutableMap.of()); } private static void fromMetastoreApiStorageDescriptor( @@ -789,7 +791,7 @@ public static Map updateStatisticsParameters(Map return result.buildOrThrow(); } - public static ColumnStatisticsObj createMetastoreColumnStatistics(String columnName, HiveType columnType, HiveColumnStatistics statistics, OptionalLong rowCount) + public static ColumnStatisticsObj createMetastoreColumnStatistics(String columnName, HiveType columnType, HiveColumnStatistics statistics) { TypeInfo typeInfo = columnType.getTypeInfo(); checkArgument(typeInfo.getCategory() == PRIMITIVE, "unsupported type: %s", columnType); @@ -808,11 +810,11 @@ public static ColumnStatisticsObj createMetastoreColumnStatistics(String columnN case STRING: case VARCHAR: case CHAR: - return createStringStatistics(columnName, columnType, statistics, rowCount); + return createStringStatistics(columnName, columnType, statistics); case DATE: return createDateStatistics(columnName, columnType, statistics); case BINARY: - return createBinaryStatistics(columnName, columnType, statistics, rowCount); + return createBinaryStatistics(columnName, columnType, statistics); case DECIMAL: return createDecimalStatistics(columnName, columnType, statistics); @@ -846,7 +848,7 @@ private static ColumnStatisticsObj createLongStatistics(String columnName, HiveT integerStatistics.getMax().ifPresent(data::setHighValue); }); statistics.getNullsCount().ifPresent(data::setNumNulls); - toMetastoreDistinctValuesCount(statistics.getDistinctValuesCount(), statistics.getNullsCount()).ifPresent(data::setNumDVs); + toMetastoreDistinctValuesCount(statistics.getDistinctValuesWithNullCount(), statistics.getNullsCount()).ifPresent(data::setNumDVs); return new ColumnStatisticsObj(columnName, columnType.toString(), longStats(data)); } @@ -858,17 +860,17 @@ private static ColumnStatisticsObj createDoubleStatistics(String columnName, Hiv doubleStatistics.getMax().ifPresent(data::setHighValue); }); statistics.getNullsCount().ifPresent(data::setNumNulls); - toMetastoreDistinctValuesCount(statistics.getDistinctValuesCount(), statistics.getNullsCount()).ifPresent(data::setNumDVs); + toMetastoreDistinctValuesCount(statistics.getDistinctValuesWithNullCount(), statistics.getNullsCount()).ifPresent(data::setNumDVs); return new ColumnStatisticsObj(columnName, columnType.toString(), doubleStats(data)); } - private static ColumnStatisticsObj createStringStatistics(String columnName, HiveType columnType, HiveColumnStatistics statistics, OptionalLong rowCount) + private static ColumnStatisticsObj createStringStatistics(String columnName, HiveType columnType, HiveColumnStatistics statistics) { StringColumnStatsData data = new StringColumnStatsData(); statistics.getNullsCount().ifPresent(data::setNumNulls); - toMetastoreDistinctValuesCount(statistics.getDistinctValuesCount(), statistics.getNullsCount()).ifPresent(data::setNumDVs); + statistics.getDistinctValuesWithNullCount().ifPresent(data::setNumDVs); data.setMaxColLen(statistics.getMaxValueSizeInBytes().orElse(0)); - data.setAvgColLen(getAverageColumnLength(statistics.getTotalSizeInBytes(), rowCount, statistics.getNullsCount()).orElse(0)); + data.setAvgColLen(statistics.getAverageColumnLength().orElse(0)); return new ColumnStatisticsObj(columnName, columnType.toString(), stringStats(data)); } @@ -880,16 +882,16 @@ private static ColumnStatisticsObj createDateStatistics(String columnName, HiveT dateStatistics.getMax().ifPresent(value -> data.setHighValue(toMetastoreDate(value))); }); statistics.getNullsCount().ifPresent(data::setNumNulls); - toMetastoreDistinctValuesCount(statistics.getDistinctValuesCount(), statistics.getNullsCount()).ifPresent(data::setNumDVs); + statistics.getDistinctValuesWithNullCount().ifPresent(data::setNumDVs); return new ColumnStatisticsObj(columnName, columnType.toString(), dateStats(data)); } - private static ColumnStatisticsObj createBinaryStatistics(String columnName, HiveType columnType, HiveColumnStatistics statistics, OptionalLong rowCount) + private static ColumnStatisticsObj createBinaryStatistics(String columnName, HiveType columnType, HiveColumnStatistics statistics) { BinaryColumnStatsData data = new BinaryColumnStatsData(); statistics.getNullsCount().ifPresent(data::setNumNulls); data.setMaxColLen(statistics.getMaxValueSizeInBytes().orElse(0)); - data.setAvgColLen(getAverageColumnLength(statistics.getTotalSizeInBytes(), rowCount, statistics.getNullsCount()).orElse(0)); + data.setAvgColLen(statistics.getAverageColumnLength().orElse(0)); return new ColumnStatisticsObj(columnName, columnType.toString(), binaryStats(data)); } @@ -901,7 +903,7 @@ private static ColumnStatisticsObj createDecimalStatistics(String columnName, Hi decimalStatistics.getMax().ifPresent(value -> data.setHighValue(toMetastoreDecimal(value))); }); statistics.getNullsCount().ifPresent(data::setNumNulls); - toMetastoreDistinctValuesCount(statistics.getDistinctValuesCount(), statistics.getNullsCount()).ifPresent(data::setNumDVs); + statistics.getDistinctValuesWithNullCount().ifPresent(data::setNumDVs); return new ColumnStatisticsObj(columnName, columnType.toString(), decimalStats(data)); } @@ -968,4 +970,22 @@ private static boolean isNumericType(Type type) type.equals(DOUBLE) || type.equals(REAL) || type instanceof DecimalType; } + + public static DataOperationType toDataOperationType(AcidOperation acidOperation) + { + return switch (acidOperation) { + case INSERT -> DataOperationType.INSERT; + case MERGE -> DataOperationType.UPDATE; + default -> throw new IllegalStateException("No metastore operation for ACID operation " + acidOperation); + }; + } + + public static boolean isAvroTableWithSchemaSet(Table table) + { + return AVRO.getSerde().equals(table.getStorage().getStorageFormat().getSerDeNullable()) && + ((table.getParameters().get(AVRO_SCHEMA_URL_KEY) != null || + (table.getStorage().getSerdeParameters().get(AVRO_SCHEMA_URL_KEY) != null)) || + (table.getParameters().get(AVRO_SCHEMA_LITERAL_KEY) != null || + (table.getStorage().getSerdeParameters().get(AVRO_SCHEMA_LITERAL_KEY) != null))); + } } diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/ThriftSparkMetastoreUtil.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/ThriftSparkMetastoreUtil.java index 5b75c4837073..cae0d82b4e57 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/ThriftSparkMetastoreUtil.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/ThriftSparkMetastoreUtil.java @@ -20,11 +20,13 @@ import io.trino.plugin.hive.HiveType; import io.trino.plugin.hive.PartitionStatistics; import io.trino.plugin.hive.metastore.HiveColumnStatistics; +import io.trino.plugin.hive.metastore.MetastoreUtil; import io.trino.plugin.hive.type.PrimitiveTypeInfo; import io.trino.plugin.hive.type.TypeInfo; import java.util.AbstractMap; import java.util.Map; +import java.util.Optional; import java.util.OptionalDouble; import java.util.OptionalLong; @@ -41,10 +43,9 @@ import static io.trino.plugin.hive.metastore.thrift.ThriftMetastoreParameterParserUtils.toDouble; import static io.trino.plugin.hive.metastore.thrift.ThriftMetastoreParameterParserUtils.toLong; import static io.trino.plugin.hive.metastore.thrift.ThriftMetastoreUtil.NUM_ROWS; -import static io.trino.plugin.hive.metastore.thrift.ThriftMetastoreUtil.getTotalSizeInBytes; import static io.trino.plugin.hive.type.Category.PRIMITIVE; -final class ThriftSparkMetastoreUtil +public final class ThriftSparkMetastoreUtil { private static final String SPARK_SQL_STATS_PREFIX = "spark.sql.statistics."; private static final String COLUMN_STATS_PREFIX = SPARK_SQL_STATS_PREFIX + "colStats."; @@ -72,6 +73,23 @@ public static PartitionStatistics getTableStatistics(Table table) return new PartitionStatistics(sparkBasicStatistics, columnStatistics); } + public static Optional getSparkTableStatistics(Map parameters, Map columns) + { + if (toLong(parameters.get(MetastoreUtil.NUM_ROWS)).isPresent()) { + return Optional.empty(); + } + + HiveBasicStatistics sparkBasicStatistics = getSparkBasicStatistics(parameters); + if (sparkBasicStatistics.getRowCount().isEmpty()) { + return Optional.empty(); + } + + Map columnStatistics = columns.entrySet().stream() + .map(entry -> Map.entry(entry.getKey(), fromMetastoreColumnStatistics(entry.getKey(), entry.getValue(), parameters))) + .collect(toImmutableMap(Map.Entry::getKey, Map.Entry::getValue)); + return Optional.of(new PartitionStatistics(sparkBasicStatistics, columnStatistics)); + } + public static HiveBasicStatistics getSparkBasicStatistics(Map parameters) { OptionalLong rowCount = toLong(parameters.get(SPARK_SQL_STATS_PREFIX + NUM_ROWS)); @@ -84,6 +102,63 @@ public static HiveBasicStatistics getSparkBasicStatistics(Map pa return new HiveBasicStatistics(fileCount, rowCount, inMemoryDataSizeInBytes, onDiskDataSizeInBytes); } + @VisibleForTesting + static HiveColumnStatistics fromMetastoreColumnStatistics(String columnName, HiveType type, Map parameters) + { + TypeInfo typeInfo = type.getTypeInfo(); + if (typeInfo.getCategory() != PRIMITIVE) { + // Spark does not support table statistics for non-primitive types + return HiveColumnStatistics.empty(); + } + String field = COLUMN_STATS_PREFIX + columnName + "."; + OptionalLong maxLength = toLong(parameters.get(field + "maxLen")); + OptionalDouble avgLength = toDouble(parameters.get(field + "avgLen")); + OptionalLong nullsCount = toLong(parameters.get(field + "nullCount")); + OptionalLong distinctValuesWithNullCount = toLong(parameters.get(field + "distinctCount")); + + return switch (((PrimitiveTypeInfo) typeInfo).getPrimitiveCategory()) { + case BOOLEAN -> createBooleanColumnStatistics( + OptionalLong.empty(), + OptionalLong.empty(), + nullsCount); + case BYTE, SHORT, INT, LONG -> createIntegerColumnStatistics( + toLong(parameters.get(field + COLUMN_MIN)), + toLong(parameters.get(field + COLUMN_MAX)), + nullsCount, + distinctValuesWithNullCount); + case TIMESTAMP -> createIntegerColumnStatistics( + OptionalLong.empty(), + OptionalLong.empty(), + nullsCount, + distinctValuesWithNullCount); + case FLOAT, DOUBLE -> createDoubleColumnStatistics( + toDouble(parameters.get(field + COLUMN_MIN)), + toDouble(parameters.get(field + COLUMN_MAX)), + nullsCount, + distinctValuesWithNullCount); + case STRING, VARCHAR, CHAR -> createStringColumnStatistics( + maxLength, + avgLength, + nullsCount, + distinctValuesWithNullCount); + case DATE -> createDateColumnStatistics( + toDate(parameters.get(field + COLUMN_MIN)), + toDate(parameters.get(field + COLUMN_MAX)), + nullsCount, + distinctValuesWithNullCount); + case BINARY -> createBinaryColumnStatistics( + maxLength, + avgLength, + nullsCount); + case DECIMAL -> createDecimalColumnStatistics( + toDecimal(parameters.get(field + COLUMN_MIN)), + toDecimal(parameters.get(field + COLUMN_MAX)), + nullsCount, + distinctValuesWithNullCount); + case TIMESTAMPLOCALTZ, INTERVAL_YEAR_MONTH, INTERVAL_DAY_TIME, VOID, UNKNOWN -> HiveColumnStatistics.empty(); + }; + } + @VisibleForTesting static HiveColumnStatistics fromMetastoreColumnStatistics(FieldSchema fieldSchema, Map columnStatistics, long rowCount) { @@ -97,7 +172,7 @@ static HiveColumnStatistics fromMetastoreColumnStatistics(FieldSchema fieldSchem OptionalLong maxLength = toLong(columnStatistics.get(field + "maxLen")); OptionalDouble avgLength = toDouble(columnStatistics.get(field + "avgLen")); OptionalLong nullsCount = toLong(columnStatistics.get(field + "nullCount")); - OptionalLong distinctValuesCount = toLong(columnStatistics.get(field + "distinctCount")); + OptionalLong distinctValuesWithNullCount = toLong(columnStatistics.get(field + "distinctCount")); return switch (((PrimitiveTypeInfo) typeInfo).getPrimitiveCategory()) { case BOOLEAN -> createBooleanColumnStatistics( @@ -108,60 +183,37 @@ static HiveColumnStatistics fromMetastoreColumnStatistics(FieldSchema fieldSchem toLong(columnStatistics.get(field + COLUMN_MIN)), toLong(columnStatistics.get(field + COLUMN_MAX)), nullsCount, - fromMetastoreDistinctValuesCount(distinctValuesCount, nullsCount, rowCount)); + distinctValuesWithNullCount); case TIMESTAMP -> createIntegerColumnStatistics( OptionalLong.empty(), OptionalLong.empty(), nullsCount, - fromMetastoreDistinctValuesCount(distinctValuesCount, nullsCount, rowCount)); + distinctValuesWithNullCount); case FLOAT, DOUBLE -> createDoubleColumnStatistics( toDouble(columnStatistics.get(field + COLUMN_MIN)), toDouble(columnStatistics.get(field + COLUMN_MAX)), nullsCount, - fromMetastoreDistinctValuesCount(distinctValuesCount, nullsCount, rowCount)); + distinctValuesWithNullCount); case STRING, VARCHAR, CHAR -> createStringColumnStatistics( maxLength, - getTotalSizeInBytes(avgLength, OptionalLong.of(rowCount), nullsCount), + avgLength, nullsCount, - fromMetastoreDistinctValuesCount(distinctValuesCount, nullsCount, rowCount)); + distinctValuesWithNullCount); case DATE -> createDateColumnStatistics( toDate(columnStatistics.get(field + COLUMN_MIN)), toDate(columnStatistics.get(field + COLUMN_MAX)), nullsCount, - fromMetastoreDistinctValuesCount(distinctValuesCount, nullsCount, rowCount)); + distinctValuesWithNullCount); case BINARY -> createBinaryColumnStatistics( maxLength, - getTotalSizeInBytes(avgLength, OptionalLong.of(rowCount), nullsCount), + avgLength, nullsCount); case DECIMAL -> createDecimalColumnStatistics( toDecimal(columnStatistics.get(field + COLUMN_MIN)), toDecimal(columnStatistics.get(field + COLUMN_MAX)), nullsCount, - fromMetastoreDistinctValuesCount(distinctValuesCount, nullsCount, rowCount)); + distinctValuesWithNullCount); case TIMESTAMPLOCALTZ, INTERVAL_YEAR_MONTH, INTERVAL_DAY_TIME, VOID, UNKNOWN -> HiveColumnStatistics.empty(); }; } - - /** - * Hive calculates NDV considering null as a distinct value, but Spark doesn't - */ - private static OptionalLong fromMetastoreDistinctValuesCount(OptionalLong distinctValuesCount, OptionalLong nullsCount, long rowCount) - { - if (distinctValuesCount.isPresent() && nullsCount.isPresent()) { - return OptionalLong.of(fromMetastoreDistinctValuesCount(distinctValuesCount.getAsLong(), nullsCount.getAsLong(), rowCount)); - } - return OptionalLong.empty(); - } - - private static long fromMetastoreDistinctValuesCount(long distinctValuesCount, long nullsCount, long rowCount) - { - long nonNullsCount = rowCount - nullsCount; - // normalize distinctValuesCount in case there is a non-null element - if (nonNullsCount > 0 && distinctValuesCount == 0) { - distinctValuesCount = 1; - } - - // the metastore may store an estimate, so the value stored may be higher than the total number of rows - return Math.min(distinctValuesCount, nonNullsCount); - } } diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/Transport.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/Transport.java index f2f62cf47426..817dd5974106 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/Transport.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/Transport.java @@ -36,14 +36,15 @@ public static TTransport create( HostAndPort address, Optional sslContext, Optional socksProxy, - int timeoutMillis, + int connectTimeoutMillis, + int readTimeoutMillis, HiveMetastoreAuthentication authentication, Optional delegationToken) throws TTransportException { requireNonNull(address, "address is null"); try { - TTransport rawTransport = createRaw(address, sslContext, socksProxy, timeoutMillis); + TTransport rawTransport = createRaw(address, sslContext, socksProxy, connectTimeoutMillis, readTimeoutMillis); TTransport authenticatedTransport = authentication.authenticate(rawTransport, address.getHost(), delegationToken); if (!authenticatedTransport.isOpen()) { authenticatedTransport.open(); @@ -57,7 +58,7 @@ public static TTransport create( private Transport() {} - private static TTransport createRaw(HostAndPort address, Optional sslContext, Optional socksProxy, int timeoutMillis) + private static TTransport createRaw(HostAndPort address, Optional sslContext, Optional socksProxy, int connectTimeoutMillis, int readTimeoutMillis) throws TTransportException { Proxy proxy = socksProxy @@ -66,8 +67,8 @@ private static TTransport createRaw(HostAndPort address, Optional ss Socket socket = new Socket(proxy); try { - socket.connect(new InetSocketAddress(address.getHost(), address.getPort()), timeoutMillis); - socket.setSoTimeout(timeoutMillis); + socket.connect(new InetSocketAddress(address.getHost(), address.getPort()), connectTimeoutMillis); + socket.setSoTimeout(readTimeoutMillis); if (sslContext.isPresent()) { // SSL will connect to the SOCKS address when present diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/orc/OrcFileWriterFactory.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/orc/OrcFileWriterFactory.java index 8ad0a0bd51c3..247f7881824e 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/orc/OrcFileWriterFactory.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/orc/OrcFileWriterFactory.java @@ -44,16 +44,16 @@ import java.io.Closeable; import java.io.IOException; import java.util.List; +import java.util.Map; import java.util.Optional; import java.util.OptionalInt; -import java.util.Properties; import java.util.function.Supplier; import static io.trino.orc.metadata.OrcType.createRootOrcType; import static io.trino.plugin.hive.HiveErrorCode.HIVE_WRITER_OPEN_ERROR; import static io.trino.plugin.hive.HiveErrorCode.HIVE_WRITE_VALIDATION_FAILED; -import static io.trino.plugin.hive.HiveMetadata.PRESTO_QUERY_ID_NAME; -import static io.trino.plugin.hive.HiveMetadata.PRESTO_VERSION_NAME; +import static io.trino.plugin.hive.HiveMetadata.TRINO_QUERY_ID_NAME; +import static io.trino.plugin.hive.HiveMetadata.TRINO_VERSION_NAME; import static io.trino.plugin.hive.HiveSessionProperties.getOrcOptimizedWriterMaxDictionaryMemory; import static io.trino.plugin.hive.HiveSessionProperties.getOrcOptimizedWriterMaxStripeRows; import static io.trino.plugin.hive.HiveSessionProperties.getOrcOptimizedWriterMaxStripeSize; @@ -125,7 +125,7 @@ public Optional createFileWriter( List inputColumnNames, StorageFormat storageFormat, HiveCompressionCodec compressionCodec, - Properties schema, + Map schema, ConnectorSession session, OptionalInt bucketNumber, AcidTransaction transaction, @@ -199,8 +199,8 @@ public Optional createFileWriter( .withMaxStringStatisticsLimit(getOrcStringStatisticsLimit(session)), fileInputColumnIndexes, ImmutableMap.builder() - .put(PRESTO_VERSION_NAME, nodeVersion.toString()) - .put(PRESTO_QUERY_ID_NAME, session.getQueryId()) + .put(TRINO_VERSION_NAME, nodeVersion.toString()) + .put(TRINO_QUERY_ID_NAME, session.getQueryId()) .buildOrThrow(), validationInputFactory, getOrcOptimizedWriterValidateMode(session), diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/orc/OrcPageSource.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/orc/OrcPageSource.java index 1b92a84529a0..6fa77929d1b2 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/orc/OrcPageSource.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/orc/OrcPageSource.java @@ -13,7 +13,6 @@ */ package io.trino.plugin.hive.orc; -import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.common.io.Closer; import io.trino.memory.context.AggregatedMemoryContext; @@ -42,7 +41,6 @@ import java.io.IOException; import java.io.UncheckedIOException; -import java.util.List; import java.util.Optional; import java.util.OptionalLong; @@ -70,7 +68,6 @@ public class OrcPageSource public static final String ORC_CODEC_METRIC_PREFIX = "OrcReaderCompressionFormat_"; private final OrcRecordReader recordReader; - private final List columnAdaptations; private final OrcDataSource orcDataSource; private final Optional deletedRows; @@ -91,7 +88,6 @@ public class OrcPageSource public OrcPageSource( OrcRecordReader recordReader, - List columnAdaptations, OrcDataSource orcDataSource, Optional deletedRows, Optional originalFileRowId, @@ -100,7 +96,6 @@ public OrcPageSource( CompressionKind compressionKind) { this.recordReader = requireNonNull(recordReader, "recordReader is null"); - this.columnAdaptations = ImmutableList.copyOf(requireNonNull(columnAdaptations, "columnAdaptations is null")); this.orcDataSource = requireNonNull(orcDataSource, "orcDataSource is null"); this.deletedRows = requireNonNull(deletedRows, "deletedRows is null"); this.stats = requireNonNull(stats, "stats is null"); @@ -182,19 +177,7 @@ public Page getNextPage() } } - MaskDeletedRowsFunction maskDeletedRowsFunction = deletedRows - .map(deletedRows -> deletedRows.getMaskDeletedRowsFunction(page, startRowId)) - .orElseGet(() -> MaskDeletedRowsFunction.noMaskForPage(page)); - return getColumnAdaptationsPage(page, maskDeletedRowsFunction, recordReader.getFilePosition(), startRowId); - } - - private Page getColumnAdaptationsPage(Page page, MaskDeletedRowsFunction maskDeletedRowsFunction, long filePosition, OptionalLong startRowId) - { - Block[] blocks = new Block[columnAdaptations.size()]; - for (int i = 0; i < columnAdaptations.size(); i++) { - blocks[i] = columnAdaptations.get(i).block(page, maskDeletedRowsFunction, filePosition, startRowId); - } - return new Page(maskDeletedRowsFunction.getPositionCount(), blocks); + return page; } static TrinoException handleException(OrcDataSourceId dataSourceId, Exception exception) @@ -243,7 +226,6 @@ public String toString() { return toStringHelper(this) .add("orcDataSource", orcDataSource.getId()) - .add("columns", columnAdaptations) .toString(); } diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/orc/OrcPageSourceFactory.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/orc/OrcPageSourceFactory.java index 439599f87b51..9b74ff110449 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/orc/OrcPageSourceFactory.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/orc/OrcPageSourceFactory.java @@ -15,7 +15,6 @@ import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; -import com.google.common.collect.Maps; import com.google.inject.Inject; import io.airlift.slice.Slice; import io.trino.filesystem.Location; @@ -42,6 +41,7 @@ import io.trino.plugin.hive.HivePageSourceFactory; import io.trino.plugin.hive.ReaderColumns; import io.trino.plugin.hive.ReaderPageSource; +import io.trino.plugin.hive.Schema; import io.trino.plugin.hive.acid.AcidSchema; import io.trino.plugin.hive.acid.AcidTransaction; import io.trino.plugin.hive.coercions.TypeCoercer; @@ -61,7 +61,6 @@ import java.util.Map; import java.util.Optional; import java.util.OptionalInt; -import java.util.Properties; import java.util.regex.Pattern; import java.util.stream.Collectors; @@ -95,11 +94,8 @@ import static io.trino.plugin.hive.orc.OrcPageSource.ColumnAdaptation.mergedRowColumns; import static io.trino.plugin.hive.orc.OrcPageSource.handleException; import static io.trino.plugin.hive.orc.OrcTypeTranslator.createCoercer; -import static io.trino.plugin.hive.util.AcidTables.isFullAcidTable; import static io.trino.plugin.hive.util.HiveClassNames.ORC_SERDE_CLASS; -import static io.trino.plugin.hive.util.HiveUtil.getDeserializerClassName; import static io.trino.plugin.hive.util.HiveUtil.splitError; -import static io.trino.plugin.hive.util.SerdeConstants.SERIALIZATION_LIB; import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED; import static io.trino.spi.type.BigintType.BIGINT; import static io.trino.spi.type.IntegerType.INTEGER; @@ -160,14 +156,9 @@ public OrcPageSourceFactory( this.fileSystemFactory = requireNonNull(fileSystemFactory, "fileSystemFactory is null"); } - public static Properties stripUnnecessaryProperties(Properties schema) + public static boolean stripUnnecessaryProperties(String serializationLibraryName) { - if (ORC_SERDE_CLASS.equals(getDeserializerClassName(schema)) && !isFullAcidTable(Maps.fromProperties(schema))) { - Properties stripped = new Properties(); - stripped.put(SERIALIZATION_LIB, schema.getProperty(SERIALIZATION_LIB)); - return stripped; - } - return schema; + return ORC_SERDE_CLASS.equals(serializationLibraryName); } @Override @@ -177,7 +168,8 @@ public Optional createPageSource( long start, long length, long estimatedFileSize, - Properties schema, + long fileModifiedTime, + Schema schema, List columns, TupleDomain effectivePredicate, Optional acidInfo, @@ -185,7 +177,7 @@ public Optional createPageSource( boolean originalFile, AcidTransaction transaction) { - if (!ORC_SERDE_CLASS.equals(getDeserializerClassName(schema))) { + if (!ORC_SERDE_CLASS.equals(schema.serializationLibraryName())) { return Optional.empty(); } @@ -207,7 +199,7 @@ public Optional createPageSource( readerColumnHandles, columns, isUseOrcColumnNames(session), - isFullAcidTable(Maps.fromProperties(schema)), + schema.isFullAcidTable(), effectivePredicate, legacyTimeZone, orcReaderOptions @@ -435,7 +427,6 @@ else if (column.getBaseHiveColumnIndex() < fileColumns.size()) { return new OrcPageSource( recordReader, - columnAdaptations, orcDataSource, deletedRows, originalFileRowId, diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/parquet/MemoryParquetDataSource.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/parquet/MemoryParquetDataSource.java new file mode 100644 index 000000000000..b5b848cb24b4 --- /dev/null +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/parquet/MemoryParquetDataSource.java @@ -0,0 +1,151 @@ +/* + * 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.hive.parquet; + +import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ListMultimap; +import io.airlift.slice.Slice; +import io.trino.filesystem.TrinoInput; +import io.trino.filesystem.TrinoInputFile; +import io.trino.memory.context.AggregatedMemoryContext; +import io.trino.memory.context.LocalMemoryContext; +import io.trino.parquet.ChunkReader; +import io.trino.parquet.DiskRange; +import io.trino.parquet.ParquetDataSource; +import io.trino.parquet.ParquetDataSourceId; +import io.trino.parquet.reader.ChunkedInputStream; +import io.trino.plugin.hive.FileFormatDataSourceStats; +import jakarta.annotation.Nullable; + +import java.io.IOException; +import java.util.Collection; +import java.util.List; +import java.util.Map; + +import static com.google.common.collect.ImmutableList.toImmutableList; +import static java.lang.Math.min; +import static java.lang.Math.toIntExact; +import static java.util.Objects.requireNonNull; + +public class MemoryParquetDataSource + implements ParquetDataSource +{ + private final ParquetDataSourceId id; + private final long readTimeNanos; + private final long readBytes; + private final LocalMemoryContext memoryUsage; + @Nullable + private Slice data; + + public MemoryParquetDataSource(TrinoInputFile inputFile, AggregatedMemoryContext memoryContext, FileFormatDataSourceStats stats) + throws IOException + { + try (TrinoInput input = inputFile.newInput()) { + long readStart = System.nanoTime(); + this.data = input.readTail(toIntExact(inputFile.length())); + this.readTimeNanos = System.nanoTime() - readStart; + stats.readDataBytesPerSecond(data.length(), readTimeNanos); + } + this.memoryUsage = memoryContext.newLocalMemoryContext(MemoryParquetDataSource.class.getSimpleName()); + this.memoryUsage.setBytes(data.length()); + this.readBytes = data.length(); + this.id = new ParquetDataSourceId(inputFile.location().toString()); + } + + @Override + public ParquetDataSourceId getId() + { + return id; + } + + @Override + public long getReadBytes() + { + return readBytes; + } + + @Override + public long getReadTimeNanos() + { + return readTimeNanos; + } + + @Override + public long getEstimatedSize() + { + return readBytes; + } + + @Override + public Slice readTail(int length) + { + int readSize = min(data.length(), length); + return readFully(data.length() - readSize, readSize); + } + + @Override + public final Slice readFully(long position, int length) + { + return data.slice(toIntExact(position), length); + } + + @Override + public Map planRead(ListMultimap diskRanges, AggregatedMemoryContext memoryContext) + { + requireNonNull(diskRanges, "diskRanges is null"); + + if (diskRanges.isEmpty()) { + return ImmutableMap.of(); + } + + ImmutableMap.Builder builder = ImmutableMap.builder(); + for (Map.Entry> entry : diskRanges.asMap().entrySet()) { + List chunkReaders = entry.getValue().stream() + .map(diskRange -> new ChunkReader() + { + @Override + public long getDiskOffset() + { + return diskRange.getOffset(); + } + + @Override + public Slice read() + { + return data.slice(toIntExact(diskRange.getOffset()), toIntExact(diskRange.getLength())); + } + + @Override + public void free() {} + }) + .collect(toImmutableList()); + builder.put(entry.getKey(), new ChunkedInputStream(chunkReaders)); + } + return builder.buildOrThrow(); + } + + @Override + public void close() + throws IOException + { + data = null; + memoryUsage.close(); + } + + @Override + public final String toString() + { + return id.toString(); + } +} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/parquet/ParquetFileWriter.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/parquet/ParquetFileWriter.java index bf875ac6223b..f989508a3359 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/parquet/ParquetFileWriter.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/parquet/ParquetFileWriter.java @@ -27,6 +27,7 @@ import io.trino.spi.block.RunLengthEncodedBlock; import io.trino.spi.type.Type; import org.apache.parquet.format.CompressionCodec; +import org.apache.parquet.format.FileMetaData; import org.apache.parquet.schema.MessageType; import org.joda.time.DateTimeZone; @@ -75,7 +76,6 @@ public ParquetFileWriter( int[] fileInputColumnIndexes, CompressionCodec compressionCodec, String trinoVersion, - boolean useBatchColumnReadersForVerification, Optional parquetTimeZone, Optional> validationInputFactory) throws IOException @@ -92,7 +92,6 @@ public ParquetFileWriter( parquetWriterOptions, compressionCodec, trinoVersion, - useBatchColumnReadersForVerification, parquetTimeZone, validationInputFactory.isPresent() ? Optional.of(new ParquetWriteValidationBuilder(fileColumnTypes, fileColumnNames)) @@ -200,4 +199,9 @@ public String toString() .add("writer", parquetWriter) .toString(); } + + public FileMetaData getFileMetadata() + { + return parquetWriter.getFileMetaData(); + } } diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/parquet/ParquetFileWriterFactory.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/parquet/ParquetFileWriterFactory.java index d509dc18ae12..d4a97aec97fd 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/parquet/ParquetFileWriterFactory.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/parquet/ParquetFileWriterFactory.java @@ -43,9 +43,9 @@ import java.io.Closeable; import java.io.IOException; import java.util.List; +import java.util.Map; import java.util.Optional; import java.util.OptionalInt; -import java.util.Properties; import java.util.function.Supplier; import static io.trino.parquet.writer.ParquetSchemaConverter.HIVE_PARQUET_USE_INT96_TIMESTAMP_ENCODING; @@ -53,7 +53,6 @@ import static io.trino.plugin.hive.HiveErrorCode.HIVE_WRITER_OPEN_ERROR; import static io.trino.plugin.hive.HiveErrorCode.HIVE_WRITE_VALIDATION_FAILED; import static io.trino.plugin.hive.HiveSessionProperties.getTimestampPrecision; -import static io.trino.plugin.hive.HiveSessionProperties.isParquetOptimizedReaderEnabled; import static io.trino.plugin.hive.HiveSessionProperties.isParquetOptimizedWriterValidate; import static io.trino.plugin.hive.util.HiveClassNames.MAPRED_PARQUET_OUTPUT_FORMAT_CLASS; import static io.trino.plugin.hive.util.HiveUtil.getColumnNames; @@ -91,7 +90,7 @@ public Optional createFileWriter( List inputColumnNames, StorageFormat storageFormat, HiveCompressionCodec compressionCodec, - Properties schema, + Map schema, ConnectorSession session, OptionalInt bucketNumber, AcidTransaction transaction, @@ -150,9 +149,10 @@ public Optional createFileWriter( schemaConverter.getPrimitiveTypes(), parquetWriterOptions, fileInputColumnIndexes, - compressionCodec.getParquetCompressionCodec(), + compressionCodec.getParquetCompressionCodec() + // Ensured by the caller + .orElseThrow(() -> new IllegalArgumentException("Unsupported compression codec for Parquet: " + compressionCodec)), nodeVersion.toString(), - isParquetOptimizedReaderEnabled(session), Optional.of(parquetTimeZone), validationInputFactory)); } diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/parquet/ParquetPageSource.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/parquet/ParquetPageSource.java index 283eab238afa..7e01334f88ae 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/parquet/ParquetPageSource.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/parquet/ParquetPageSource.java @@ -14,12 +14,15 @@ package io.trino.plugin.hive.parquet; import com.google.common.collect.ImmutableList; +import io.trino.parquet.Column; import io.trino.parquet.ParquetCorruptionException; import io.trino.parquet.ParquetDataSourceId; import io.trino.parquet.reader.ParquetReader; +import io.trino.plugin.hive.coercions.TypeCoercer; import io.trino.spi.Page; import io.trino.spi.TrinoException; import io.trino.spi.block.Block; +import io.trino.spi.block.LazyBlock; import io.trino.spi.block.LongArrayBlock; import io.trino.spi.block.RunLengthEncodedBlock; import io.trino.spi.connector.ConnectorPageSource; @@ -32,6 +35,7 @@ import java.util.Optional; import java.util.OptionalLong; +import static com.google.common.base.MoreObjects.toStringHelper; import static com.google.common.base.Preconditions.checkArgument; import static io.trino.plugin.base.util.Closables.closeAllSuppress; import static io.trino.plugin.hive.HiveErrorCode.HIVE_BAD_DATA; @@ -43,19 +47,18 @@ public class ParquetPageSource implements ConnectorPageSource { private final ParquetReader parquetReader; - private final List columnAdaptations; - private final boolean isColumnAdaptationRequired; private boolean closed; private long completedPositions; - private ParquetPageSource( - ParquetReader parquetReader, - List columnAdaptations) + public ParquetPageSource(ParquetReader parquetReader) { this.parquetReader = requireNonNull(parquetReader, "parquetReader is null"); - this.columnAdaptations = ImmutableList.copyOf(requireNonNull(columnAdaptations, "columnAdaptations is null")); - this.isColumnAdaptationRequired = isColumnAdaptationRequired(columnAdaptations); + } + + public List getColumnFields() + { + return parquetReader.getColumnFields(); } @Override @@ -93,7 +96,7 @@ public Page getNextPage() { Page page; try { - page = getColumnAdaptationsPage(parquetReader.nextPage()); + page = parquetReader.nextPage(); } catch (IOException | RuntimeException e) { closeAllSuppress(e, this); @@ -166,27 +169,16 @@ public Builder addRowIndexColumn() return this; } - public ConnectorPageSource build(ParquetReader parquetReader) + public Builder addCoercedColumn(int sourceChannel, TypeCoercer typeCoercer) { - return new ParquetPageSource(parquetReader, this.columns.build()); + columns.add(new CoercedColumn(new SourceColumn(sourceChannel), typeCoercer)); + return this; } - } - private Page getColumnAdaptationsPage(Page page) - { - if (!isColumnAdaptationRequired) { - return page; - } - if (page == null) { - return null; - } - int batchSize = page.getPositionCount(); - Block[] blocks = new Block[columnAdaptations.size()]; - long startRowId = parquetReader.lastBatchStartRow(); - for (int columnChannel = 0; columnChannel < columnAdaptations.size(); columnChannel++) { - blocks[columnChannel] = columnAdaptations.get(columnChannel).getBlock(page, startRowId); + public ConnectorPageSource build(ParquetReader parquetReader) + { + return new ParquetPageSource(parquetReader); } - return new Page(batchSize, blocks); } static TrinoException handleException(ParquetDataSourceId dataSourceId, Exception exception) @@ -230,9 +222,7 @@ private static class NullColumn private NullColumn(Type type) { - this.nullBlock = type.createBlockBuilder(null, 1, 0) - .appendNull() - .build(); + this.nullBlock = type.createNullBlock(); } @Override @@ -293,6 +283,36 @@ public Block getBlock(Page sourcePage, long startRowId) } } + private static class CoercedColumn + implements ParquetPageSource.ColumnAdaptation + { + private final ParquetPageSource.SourceColumn sourceColumn; + private final TypeCoercer typeCoercer; + + public CoercedColumn(ParquetPageSource.SourceColumn sourceColumn, TypeCoercer typeCoercer) + { + this.sourceColumn = requireNonNull(sourceColumn, "sourceColumn is null"); + this.typeCoercer = requireNonNull(typeCoercer, "typeCoercer is null"); + } + + @Override + public Block getBlock(Page sourcePage, long startRowId) + { + Block block = sourceColumn.getBlock(sourcePage, startRowId); + return new LazyBlock(block.getPositionCount(), () -> typeCoercer.apply(block.getLoadedBlock())); + } + + @Override + public String toString() + { + return toStringHelper(this) + .add("sourceColumn", sourceColumn) + .add("fromType", typeCoercer.getFromType()) + .add("toType", typeCoercer.getToType()) + .toString(); + } + } + private static Block createRowNumberBlock(long baseIndex, int size) { long[] rowIndices = new long[size]; diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/parquet/ParquetPageSourceFactory.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/parquet/ParquetPageSourceFactory.java index 30b356e11aa6..17e83f7567f5 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/parquet/ParquetPageSourceFactory.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/parquet/ParquetPageSourceFactory.java @@ -21,17 +21,20 @@ import io.trino.filesystem.TrinoFileSystem; import io.trino.filesystem.TrinoFileSystemFactory; import io.trino.filesystem.TrinoInputFile; -import io.trino.parquet.BloomFilterStore; +import io.trino.memory.context.AggregatedMemoryContext; +import io.trino.parquet.Column; import io.trino.parquet.Field; import io.trino.parquet.ParquetCorruptionException; import io.trino.parquet.ParquetDataSource; import io.trino.parquet.ParquetDataSourceId; import io.trino.parquet.ParquetReaderOptions; import io.trino.parquet.ParquetWriteValidation; +import io.trino.parquet.metadata.FileMetadata; +import io.trino.parquet.metadata.ParquetMetadata; import io.trino.parquet.predicate.TupleDomainParquetPredicate; import io.trino.parquet.reader.MetadataReader; import io.trino.parquet.reader.ParquetReader; -import io.trino.parquet.reader.TrinoColumnIndexStore; +import io.trino.parquet.reader.RowGroupInfo; import io.trino.plugin.hive.AcidInfo; import io.trino.plugin.hive.FileFormatDataSourceStats; import io.trino.plugin.hive.HiveColumnHandle; @@ -41,65 +44,65 @@ import io.trino.plugin.hive.HiveType; import io.trino.plugin.hive.ReaderColumns; import io.trino.plugin.hive.ReaderPageSource; +import io.trino.plugin.hive.Schema; +import io.trino.plugin.hive.TransformConnectorPageSource; import io.trino.plugin.hive.acid.AcidTransaction; +import io.trino.plugin.hive.coercions.TypeCoercer; +import io.trino.spi.Page; import io.trino.spi.TrinoException; +import io.trino.spi.block.Block; import io.trino.spi.connector.ConnectorPageSource; import io.trino.spi.connector.ConnectorSession; import io.trino.spi.predicate.Domain; import io.trino.spi.predicate.TupleDomain; import org.apache.parquet.column.ColumnDescriptor; -import org.apache.parquet.hadoop.metadata.BlockMetaData; -import org.apache.parquet.hadoop.metadata.ColumnChunkMetaData; -import org.apache.parquet.hadoop.metadata.ColumnPath; -import org.apache.parquet.hadoop.metadata.FileMetaData; -import org.apache.parquet.hadoop.metadata.ParquetMetadata; -import org.apache.parquet.internal.filter2.columnindex.ColumnIndexStore; +import org.apache.parquet.io.ColumnIO; import org.apache.parquet.io.MessageColumnIO; import org.apache.parquet.schema.GroupType; import org.apache.parquet.schema.MessageType; +import org.apache.parquet.schema.PrimitiveType; import org.apache.parquet.schema.Type; import org.joda.time.DateTimeZone; import java.io.IOException; -import java.util.HashSet; +import java.util.ArrayList; +import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.Map.Entry; import java.util.Optional; import java.util.OptionalInt; -import java.util.Properties; +import java.util.OptionalLong; import java.util.Set; +import java.util.function.Function; import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.collect.ImmutableList.toImmutableList; -import static com.google.common.collect.ImmutableSet.toImmutableSet; import static io.trino.memory.context.AggregatedMemoryContext.newSimpleAggregatedMemoryContext; -import static io.trino.parquet.BloomFilterStore.getBloomFilterStore; import static io.trino.parquet.ParquetTypeUtils.constructField; import static io.trino.parquet.ParquetTypeUtils.getColumnIO; import static io.trino.parquet.ParquetTypeUtils.getDescriptors; import static io.trino.parquet.ParquetTypeUtils.getParquetTypeByName; import static io.trino.parquet.ParquetTypeUtils.lookupColumnByName; import static io.trino.parquet.predicate.PredicateUtils.buildPredicate; -import static io.trino.parquet.predicate.PredicateUtils.predicateMatches; +import static io.trino.parquet.predicate.PredicateUtils.getFilteredRowGroups; import static io.trino.plugin.hive.HiveColumnHandle.ColumnType.REGULAR; import static io.trino.plugin.hive.HiveErrorCode.HIVE_BAD_DATA; import static io.trino.plugin.hive.HiveErrorCode.HIVE_CANNOT_OPEN_SPLIT; +import static io.trino.plugin.hive.HivePageSourceProvider.getProjection; import static io.trino.plugin.hive.HivePageSourceProvider.projectBaseColumns; import static io.trino.plugin.hive.HivePageSourceProvider.projectSufficientColumns; import static io.trino.plugin.hive.HiveSessionProperties.getParquetMaxReadBlockRowCount; import static io.trino.plugin.hive.HiveSessionProperties.getParquetMaxReadBlockSize; +import static io.trino.plugin.hive.HiveSessionProperties.getParquetSmallFileThreshold; import static io.trino.plugin.hive.HiveSessionProperties.isParquetIgnoreStatistics; -import static io.trino.plugin.hive.HiveSessionProperties.isParquetOptimizedNestedReaderEnabled; -import static io.trino.plugin.hive.HiveSessionProperties.isParquetOptimizedReaderEnabled; import static io.trino.plugin.hive.HiveSessionProperties.isParquetUseColumnIndex; import static io.trino.plugin.hive.HiveSessionProperties.isUseParquetColumnNames; import static io.trino.plugin.hive.HiveSessionProperties.useParquetBloomFilter; import static io.trino.plugin.hive.parquet.ParquetPageSource.handleException; +import static io.trino.plugin.hive.parquet.ParquetTypeTranslator.createCoercer; import static io.trino.plugin.hive.type.Category.PRIMITIVE; import static io.trino.plugin.hive.util.HiveClassNames.PARQUET_HIVE_SERDE_CLASS; -import static io.trino.plugin.hive.util.HiveUtil.getDeserializerClassName; -import static io.trino.plugin.hive.util.SerdeConstants.SERIALIZATION_LIB; import static io.trino.spi.type.BigintType.BIGINT; import static java.lang.String.format; import static java.util.Objects.requireNonNull; @@ -147,14 +150,9 @@ public ParquetPageSourceFactory( domainCompactionThreshold = hiveConfig.getDomainCompactionThreshold(); } - public static Properties stripUnnecessaryProperties(Properties schema) + public static boolean stripUnnecessaryProperties(String serializationLibraryName) { - if (PARQUET_SERDE_CLASS_NAMES.contains(getDeserializerClassName(schema))) { - Properties stripped = new Properties(); - stripped.put(SERIALIZATION_LIB, schema.getProperty(SERIALIZATION_LIB)); - return stripped; - } - return schema; + return PARQUET_SERDE_CLASS_NAMES.contains(serializationLibraryName); } @Override @@ -164,7 +162,8 @@ public Optional createPageSource( long start, long length, long estimatedFileSize, - Properties schema, + long fileModifiedTime, + Schema schema, List columns, TupleDomain effectivePredicate, Optional acidInfo, @@ -172,7 +171,7 @@ public Optional createPageSource( boolean originalFile, AcidTransaction transaction) { - if (!PARQUET_SERDE_CLASS_NAMES.contains(getDeserializerClassName(schema))) { + if (!PARQUET_SERDE_CLASS_NAMES.contains(schema.serializationLibraryName())) { return Optional.empty(); } @@ -186,19 +185,19 @@ public Optional createPageSource( start, length, columns, - effectivePredicate, + ImmutableList.of(effectivePredicate), isUseParquetColumnNames(session), timeZone, stats, options.withIgnoreStatistics(isParquetIgnoreStatistics(session)) .withMaxReadBlockSize(getParquetMaxReadBlockSize(session)) .withMaxReadBlockRowCount(getParquetMaxReadBlockRowCount(session)) + .withSmallFileThreshold(getParquetSmallFileThreshold(session)) .withUseColumnIndex(isParquetUseColumnIndex(session)) - .withBloomFilter(useParquetBloomFilter(session)) - .withBatchColumnReaders(isParquetOptimizedReaderEnabled(session)) - .withBatchNestedColumnReaders(isParquetOptimizedNestedReaderEnabled(session)), + .withBloomFilter(useParquetBloomFilter(session)), Optional.empty(), - domainCompactionThreshold)); + domainCompactionThreshold, + OptionalLong.of(estimatedFileSize))); } /** @@ -209,23 +208,25 @@ public static ReaderPageSource createPageSource( long start, long length, List columns, - TupleDomain effectivePredicate, + List> disjunctTupleDomains, boolean useColumnNames, DateTimeZone timeZone, FileFormatDataSourceStats stats, ParquetReaderOptions options, Optional parquetWriteValidation, - int domainCompactionThreshold) + int domainCompactionThreshold, + OptionalLong estimatedFileSize) { MessageType fileSchema; MessageType requestedSchema; MessageColumnIO messageColumn; ParquetDataSource dataSource = null; try { - dataSource = new TrinoParquetDataSource(inputFile, options, stats); + AggregatedMemoryContext memoryContext = newSimpleAggregatedMemoryContext(); + dataSource = createDataSource(inputFile, estimatedFileSize, options, memoryContext, stats); ParquetMetadata parquetMetadata = MetadataReader.readFooter(dataSource, parquetWriteValidation); - FileMetaData fileMetaData = parquetMetadata.getFileMetaData(); + FileMetadata fileMetaData = parquetMetadata.getFileMetaData(); fileSchema = fileMetaData.getSchema(); Optional message = getParquetMessageType(columns, useColumnNames, fileSchema); @@ -234,39 +235,36 @@ public static ReaderPageSource createPageSource( messageColumn = getColumnIO(fileSchema, requestedSchema); Map, ColumnDescriptor> descriptorsByPath = getDescriptors(fileSchema, requestedSchema); - TupleDomain parquetTupleDomain = options.isIgnoreStatistics() - ? TupleDomain.all() - : getParquetTupleDomain(descriptorsByPath, effectivePredicate, fileSchema, useColumnNames); - - TupleDomainParquetPredicate parquetPredicate = buildPredicate(requestedSchema, parquetTupleDomain, descriptorsByPath, timeZone); - - long nextStart = 0; - ImmutableList.Builder blocks = ImmutableList.builder(); - ImmutableList.Builder blockStarts = ImmutableList.builder(); - ImmutableList.Builder> columnIndexes = ImmutableList.builder(); - for (BlockMetaData block : parquetMetadata.getBlocks()) { - long firstDataPage = block.getColumns().get(0).getFirstDataPageOffset(); - Optional columnIndex = getColumnIndexStore(dataSource, block, descriptorsByPath, parquetTupleDomain, options); - Optional bloomFilterStore = getBloomFilterStore(dataSource, block, parquetTupleDomain, options); - - if (start <= firstDataPage && firstDataPage < start + length - && predicateMatches( - parquetPredicate, - block, - dataSource, - descriptorsByPath, - parquetTupleDomain, - columnIndex, - bloomFilterStore, - timeZone, - domainCompactionThreshold)) { - blocks.add(block); - blockStarts.add(nextStart); - columnIndexes.add(columnIndex); + List> parquetTupleDomains; + List parquetPredicates; + if (options.isIgnoreStatistics()) { + parquetTupleDomains = ImmutableList.of(TupleDomain.all()); + parquetPredicates = ImmutableList.of(buildPredicate(requestedSchema, TupleDomain.all(), descriptorsByPath, timeZone)); + } + else { + ImmutableList.Builder> parquetTupleDomainsBuilder = ImmutableList.builderWithExpectedSize(disjunctTupleDomains.size()); + ImmutableList.Builder parquetPredicatesBuilder = ImmutableList.builderWithExpectedSize(disjunctTupleDomains.size()); + for (TupleDomain tupleDomain : disjunctTupleDomains) { + TupleDomain parquetTupleDomain = getParquetTupleDomain(descriptorsByPath, tupleDomain, fileSchema, useColumnNames); + parquetTupleDomainsBuilder.add(parquetTupleDomain); + parquetPredicatesBuilder.add(buildPredicate(requestedSchema, parquetTupleDomain, descriptorsByPath, timeZone)); } - nextStart += block.getRowCount(); + parquetTupleDomains = parquetTupleDomainsBuilder.build(); + parquetPredicates = parquetPredicatesBuilder.build(); } + List rowGroups = getFilteredRowGroups( + start, + length, + dataSource, + parquetMetadata, + parquetTupleDomains, + parquetPredicates, + descriptorsByPath, + timeZone, + domainCompactionThreshold, + options); + Optional readerProjections = projectBaseColumns(columns, useColumnNames); List baseColumns = readerProjections.map(projection -> projection.get().stream() @@ -276,18 +274,19 @@ && predicateMatches( ParquetDataSourceId dataSourceId = dataSource.getId(); ParquetDataSource finalDataSource = dataSource; - ParquetReaderProvider parquetReaderProvider = fields -> new ParquetReader( + ParquetReaderProvider parquetReaderProvider = (fields, appendRowNumberColumn) -> new ParquetReader( Optional.ofNullable(fileMetaData.getCreatedBy()), fields, - blocks.build(), - blockStarts.build(), + appendRowNumberColumn, + rowGroups, finalDataSource, timeZone, - newSimpleAggregatedMemoryContext(), + memoryContext, options, exception -> handleException(dataSourceId, exception), - Optional.of(parquetPredicate), - columnIndexes.build(), + // We avoid using disjuncts of parquetPredicate for page pruning in ParquetReader as currently column indexes + // are not present in the Parquet files which are read with disjunct predicates. + parquetPredicates.size() == 1 ? Optional.of(parquetPredicates.get(0)) : Optional.empty(), parquetWriteValidation); ConnectorPageSource parquetPageSource = createParquetPageSource(baseColumns, fileSchema, messageColumn, useColumnNames, parquetReaderProvider); return new ReaderPageSource(parquetPageSource, readerProjections); @@ -298,7 +297,7 @@ && predicateMatches( dataSource.close(); } } - catch (IOException ignored) { + catch (IOException ignore) { } if (e instanceof TrinoException) { throw (TrinoException) e; @@ -311,6 +310,20 @@ && predicateMatches( } } + public static ParquetDataSource createDataSource( + TrinoInputFile inputFile, + OptionalLong estimatedFileSize, + ParquetReaderOptions options, + AggregatedMemoryContext memoryContext, + FileFormatDataSourceStats stats) + throws IOException + { + if (estimatedFileSize.isEmpty() || estimatedFileSize.getAsLong() > options.getSmallFileThreshold().toBytes()) { + return new TrinoParquetDataSource(inputFile, options, stats); + } + return new MemoryParquetDataSource(inputFile, memoryContext, stats); + } + public static Optional getParquetMessageType(List columns, boolean useColumnNames, MessageType fileSchema) { Optional message = projectSufficientColumns(columns) @@ -350,43 +363,6 @@ public static Optional getColumnType(HiveColumnH return Optional.of(new GroupType(baseType.getRepetition(), baseType.getName(), ImmutableList.of(type))); } - public static Optional getColumnIndexStore( - ParquetDataSource dataSource, - BlockMetaData blockMetadata, - Map, ColumnDescriptor> descriptorsByPath, - TupleDomain parquetTupleDomain, - ParquetReaderOptions options) - { - if (!options.isUseColumnIndex() || parquetTupleDomain.isAll() || parquetTupleDomain.isNone()) { - return Optional.empty(); - } - - boolean hasColumnIndex = false; - for (ColumnChunkMetaData column : blockMetadata.getColumns()) { - if (column.getColumnIndexReference() != null && column.getOffsetIndexReference() != null) { - hasColumnIndex = true; - break; - } - } - - if (!hasColumnIndex) { - return Optional.empty(); - } - - Set columnsReadPaths = new HashSet<>(descriptorsByPath.size()); - for (List path : descriptorsByPath.keySet()) { - columnsReadPaths.add(ColumnPath.get(path.toArray(new String[0]))); - } - - Map parquetDomains = parquetTupleDomain.getDomains() - .orElseThrow(() -> new IllegalStateException("Predicate other than none should have domains")); - Set columnsFilteredPaths = parquetDomains.keySet().stream() - .map(column -> ColumnPath.get(column.getPath())) - .collect(toImmutableSet()); - - return Optional.of(new TrinoColumnIndexStore(dataSource, blockMetadata, columnsReadPaths, columnsFilteredPaths)); - } - public static TupleDomain getParquetTupleDomain( Map, ColumnDescriptor> descriptorsByPath, TupleDomain effectivePredicate, @@ -441,44 +417,73 @@ public static TupleDomain getParquetTupleDomain( public interface ParquetReaderProvider { - ParquetReader createParquetReader(List fields) + ParquetReader createParquetReader(List fields, boolean appendRowNumberColumn) throws IOException; } public static ConnectorPageSource createParquetPageSource( - List baseColumns, + List columnHandles, MessageType fileSchema, MessageColumnIO messageColumn, boolean useColumnNames, ParquetReaderProvider parquetReaderProvider) throws IOException { - ParquetPageSource.Builder pageSourceBuilder = ParquetPageSource.builder(); - ImmutableList.Builder parquetColumnFieldsBuilder = ImmutableList.builder(); - int sourceChannel = 0; - for (HiveColumnHandle column : baseColumns) { + List parquetColumnFieldsBuilder = new ArrayList<>(columnHandles.size()); + Map baseColumnIdToOrdinal = new HashMap<>(); + TransformConnectorPageSource.Builder transforms = TransformConnectorPageSource.builder(); + boolean appendRowNumberColumn = false; + for (HiveColumnHandle column : columnHandles) { if (column == PARQUET_ROW_INDEX_COLUMN) { - pageSourceBuilder.addRowIndexColumn(); + appendRowNumberColumn = true; + transforms.transform(new GetRowPositionFromSource()); continue; } - checkArgument(column.getColumnType() == REGULAR, "column type must be REGULAR: %s", column); - Optional parquetType = getBaseColumnParquetType(column, fileSchema, useColumnNames); + + HiveColumnHandle baseColumn = column.getBaseColumn(); + Optional parquetType = getBaseColumnParquetType(baseColumn, fileSchema, useColumnNames); if (parquetType.isEmpty()) { - pageSourceBuilder.addNullColumn(column.getBaseType()); + transforms.constantValue(column.getType().createNullBlock()); continue; } - String columnName = useColumnNames ? column.getBaseColumnName() : fileSchema.getFields().get(column.getBaseHiveColumnIndex()).getName(); - Optional field = constructField(column.getBaseType(), lookupColumnByName(messageColumn, columnName)); - if (field.isEmpty()) { - pageSourceBuilder.addNullColumn(column.getBaseType()); - continue; + String baseColumnName = useColumnNames ? baseColumn.getBaseColumnName() : fileSchema.getFields().get(baseColumn.getBaseHiveColumnIndex()).getName(); + + Optional> coercer = Optional.empty(); + Integer ordinal = baseColumnIdToOrdinal.get(baseColumnName); + if (ordinal == null) { + ColumnIO columnIO = lookupColumnByName(messageColumn, baseColumnName); + if (columnIO != null && columnIO.getType().isPrimitive()) { + PrimitiveType primitiveType = columnIO.getType().asPrimitiveType(); + coercer = createCoercer(primitiveType.getPrimitiveTypeName(), primitiveType.getLogicalTypeAnnotation(), baseColumn.getBaseType()); + } + io.trino.spi.type.Type readType = coercer.map(TypeCoercer::getFromType).orElseGet(baseColumn::getBaseType); + + Optional field = constructField(readType, columnIO); + if (field.isEmpty()) { + transforms.constantValue(column.getType().createNullBlock()); + continue; + } + + ordinal = parquetColumnFieldsBuilder.size(); + parquetColumnFieldsBuilder.add(new Column(baseColumnName, field.get())); + baseColumnIdToOrdinal.put(baseColumnName, ordinal); } - parquetColumnFieldsBuilder.add(field.get()); - pageSourceBuilder.addSourceColumn(sourceChannel); - sourceChannel++; - } - return pageSourceBuilder.build(parquetReaderProvider.createParquetReader(parquetColumnFieldsBuilder.build())); + if (column.isBaseColumn()) { + transforms.column(ordinal, coercer.map(Function.identity())); + } + else { + transforms.dereferenceField( + ImmutableList.builder() + .add(ordinal) + .addAll(getProjection(column, baseColumn)) + .build(), + coercer.map(Function.identity())); + } + } + ParquetReader parquetReader = parquetReaderProvider.createParquetReader(parquetColumnFieldsBuilder, appendRowNumberColumn); + ConnectorPageSource pageSource = new ParquetPageSource(parquetReader); + return transforms.build(pageSource); } private static Optional getBaseColumnParquetType(HiveColumnHandle column, MessageType messageType, boolean useParquetColumnNames) @@ -520,4 +525,14 @@ private static Optional> dereferenceSubFiel return Optional.of(typeBuilder.build()); } + + private record GetRowPositionFromSource() + implements Function + { + @Override + public Block apply(Page page) + { + return page.getBlock(page.getChannelCount() - 1); + } + } } diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/parquet/ParquetReaderConfig.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/parquet/ParquetReaderConfig.java index 187b1790846f..147a7f0fdad3 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/parquet/ParquetReaderConfig.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/parquet/ParquetReaderConfig.java @@ -18,6 +18,7 @@ import io.airlift.configuration.DefunctConfig; import io.airlift.configuration.LegacyConfig; import io.airlift.units.DataSize; +import io.airlift.units.MaxDataSize; import io.airlift.units.MinDataSize; import io.trino.parquet.ParquetReaderOptions; import jakarta.validation.constraints.Max; @@ -30,15 +31,15 @@ }) public class ParquetReaderConfig { + public static final String PARQUET_READER_MAX_SMALL_FILE_THRESHOLD = "15MB"; + private ParquetReaderOptions options = new ParquetReaderOptions(); - @Deprecated public boolean isIgnoreStatistics() { return options.isIgnoreStatistics(); } - @Deprecated @Config("parquet.ignore-statistics") @ConfigDescription("Ignore statistics from Parquet to allow querying files with corrupted or incorrect statistics") public ParquetReaderConfig setIgnoreStatistics(boolean ignoreStatistics) @@ -116,43 +117,32 @@ public boolean isUseColumnIndex() return options.isUseColumnIndex(); } - @Config("parquet.optimized-reader.enabled") - @ConfigDescription("Use optimized Parquet reader") - public ParquetReaderConfig setOptimizedReaderEnabled(boolean optimizedReaderEnabled) - { - options = options.withBatchColumnReaders(optimizedReaderEnabled); - return this; - } - - public boolean isOptimizedReaderEnabled() - { - return options.useBatchColumnReaders(); - } - - @Config("parquet.optimized-nested-reader.enabled") - @ConfigDescription("Use optimized Parquet reader for nested columns") - public ParquetReaderConfig setOptimizedNestedReaderEnabled(boolean optimizedNestedReaderEnabled) + @Config("parquet.use-bloom-filter") + @ConfigDescription("Use Parquet Bloom filters") + public ParquetReaderConfig setUseBloomFilter(boolean useBloomFilter) { - options = options.withBatchNestedColumnReaders(optimizedNestedReaderEnabled); + options = options.withBloomFilter(useBloomFilter); return this; } - public boolean isOptimizedNestedReaderEnabled() + public boolean isUseBloomFilter() { - return options.useBatchNestedColumnReaders(); + return options.useBloomFilter(); } - @Config("parquet.use-bloom-filter") - @ConfigDescription("Use Parquet Bloom filters") - public ParquetReaderConfig setUseBloomFilter(boolean useBloomFilter) + @Config("parquet.small-file-threshold") + @ConfigDescription("Size below which a parquet file will be read entirely") + public ParquetReaderConfig setSmallFileThreshold(DataSize smallFileThreshold) { - options = options.withBloomFilter(useBloomFilter); + options = options.withSmallFileThreshold(smallFileThreshold); return this; } - public boolean isUseBloomFilter() + @NotNull + @MaxDataSize(PARQUET_READER_MAX_SMALL_FILE_THRESHOLD) + public DataSize getSmallFileThreshold() { - return options.useBloomFilter(); + return options.getSmallFileThreshold(); } public ParquetReaderOptions toParquetReaderOptions() diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/parquet/ParquetTypeTranslator.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/parquet/ParquetTypeTranslator.java new file mode 100644 index 000000000000..ac96b70c6eec --- /dev/null +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/parquet/ParquetTypeTranslator.java @@ -0,0 +1,85 @@ +/* + * 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.hive.parquet; + +import io.trino.plugin.hive.coercions.IntegerNumberToDoubleCoercer; +import io.trino.plugin.hive.coercions.IntegerNumberToVarcharCoercer; +import io.trino.plugin.hive.coercions.TypeCoercer; +import io.trino.spi.type.DecimalType; +import io.trino.spi.type.DoubleType; +import io.trino.spi.type.Type; +import io.trino.spi.type.VarcharType; +import org.apache.parquet.schema.LogicalTypeAnnotation; +import org.apache.parquet.schema.LogicalTypeAnnotation.DecimalLogicalTypeAnnotation; + +import java.util.Optional; + +import static io.trino.parquet.reader.ColumnReaderFactory.isIntegerAnnotationAndPrimitive; +import static io.trino.plugin.hive.coercions.DecimalCoercers.createDecimalToVarcharCoercer; +import static io.trino.plugin.hive.coercions.DoubleToVarcharCoercers.createDoubleToVarcharCoercer; +import static io.trino.plugin.hive.coercions.FloatToVarcharCoercers.createFloatToVarcharCoercer; +import static io.trino.plugin.hive.coercions.TimestampCoercer.LongTimestampToVarcharCoercer; +import static io.trino.spi.type.BigintType.BIGINT; +import static io.trino.spi.type.IntegerType.INTEGER; +import static io.trino.spi.type.TimestampType.TIMESTAMP_NANOS; +import static org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName; +import static org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName.DOUBLE; +import static org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName.FLOAT; +import static org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName.INT32; +import static org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName.INT64; +import static org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName.INT96; + +public final class ParquetTypeTranslator +{ + private ParquetTypeTranslator() {} + + public static Optional> createCoercer(PrimitiveTypeName fromParquetType, LogicalTypeAnnotation typeAnnotation, Type toTrinoType) + { + if (toTrinoType instanceof DoubleType) { + if (isIntegerAnnotationAndPrimitive(typeAnnotation, fromParquetType)) { + if (fromParquetType == INT32) { + return Optional.of(new IntegerNumberToDoubleCoercer<>(INTEGER)); + } + if (fromParquetType == INT64) { + return Optional.of(new IntegerNumberToDoubleCoercer<>(BIGINT)); + } + } + } + if (toTrinoType instanceof VarcharType varcharType) { + if (isIntegerAnnotationAndPrimitive(typeAnnotation, fromParquetType)) { + if (fromParquetType == INT32) { + return Optional.of(new IntegerNumberToVarcharCoercer<>(INTEGER, varcharType)); + } + if (fromParquetType == INT64) { + return Optional.of(new IntegerNumberToVarcharCoercer<>(BIGINT, varcharType)); + } + } + if (fromParquetType == FLOAT) { + return Optional.of(createFloatToVarcharCoercer(varcharType, false)); + } + if (fromParquetType == DOUBLE) { + return Optional.of(createDoubleToVarcharCoercer(varcharType, false)); + } + if (typeAnnotation instanceof DecimalLogicalTypeAnnotation decimalAnnotation) { + return Optional.of(createDecimalToVarcharCoercer( + DecimalType.createDecimalType(decimalAnnotation.getPrecision(), decimalAnnotation.getScale()), + varcharType)); + } + if (fromParquetType == INT96) { + return Optional.of(new LongTimestampToVarcharCoercer(TIMESTAMP_NANOS, varcharType)); + } + } + return Optional.empty(); + } +} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/parquet/ParquetUtil.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/parquet/ParquetUtil.java new file mode 100644 index 000000000000..c2bb087ec90f --- /dev/null +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/parquet/ParquetUtil.java @@ -0,0 +1,124 @@ +/* + * 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.hive.parquet; + +import com.google.common.collect.ImmutableMap; +import io.trino.filesystem.Location; +import io.trino.filesystem.memory.MemoryFileSystemFactory; +import io.trino.plugin.hive.FileFormatDataSourceStats; +import io.trino.plugin.hive.HiveColumnHandle; +import io.trino.plugin.hive.HiveConfig; +import io.trino.plugin.hive.HivePageSourceFactory; +import io.trino.plugin.hive.HiveStorageFormat; +import io.trino.plugin.hive.Schema; +import io.trino.spi.connector.ConnectorPageSource; +import io.trino.spi.connector.ConnectorSession; +import io.trino.spi.predicate.TupleDomain; +import io.trino.spi.security.ConnectorIdentity; +import io.trino.spi.type.Type; +import org.joda.time.DateTimeZone; + +import java.io.File; +import java.io.IOException; +import java.io.OutputStream; +import java.nio.file.Files; +import java.util.List; +import java.util.Optional; +import java.util.OptionalInt; +import java.util.stream.IntStream; + +import static com.google.common.base.Preconditions.checkArgument; +import static com.google.common.collect.ImmutableList.toImmutableList; +import static io.trino.plugin.hive.HiveColumnHandle.ColumnType.REGULAR; +import static io.trino.plugin.hive.HiveColumnHandle.createBaseColumn; +import static io.trino.plugin.hive.acid.AcidTransaction.NO_ACID_TRANSACTION; +import static io.trino.plugin.hive.util.HiveTypeTranslator.toHiveType; + +final class ParquetUtil +{ + private ParquetUtil() {} + + public static ConnectorPageSource createPageSource(ConnectorSession session, File parquetFile, List columnNames, List columnTypes) + throws IOException + { + return createPageSource(session, parquetFile, getBaseColumns(columnNames, columnTypes), TupleDomain.all()); + } + + public static ConnectorPageSource createPageSource(ConnectorSession session, File parquetFile, List columnNames, List columnTypes, DateTimeZone timeZone) + throws IOException + { + return createPageSource(session, parquetFile, getBaseColumns(columnNames, columnTypes), TupleDomain.all(), new HiveConfig().setParquetTimeZone(timeZone.toString())); + } + + public static ConnectorPageSource createPageSource(ConnectorSession session, File parquetFile, List columns, TupleDomain domain, DateTimeZone timeZone) + throws IOException + { + return createPageSource(session, parquetFile, columns, domain, new HiveConfig().setParquetTimeZone(timeZone.toString())); + } + + public static ConnectorPageSource createPageSource(ConnectorSession session, File parquetFile, List columns, TupleDomain domain) + throws IOException + { + return createPageSource(session, parquetFile, columns, domain, new HiveConfig()); + } + + private static ConnectorPageSource createPageSource(ConnectorSession session, File parquetFile, List columns, TupleDomain domain, HiveConfig hiveConfig) + throws IOException + { + // copy the test file into the memory filesystem + MemoryFileSystemFactory fileSystemFactory = new MemoryFileSystemFactory(); + Location location = Location.of("memory:///test.file"); + try (OutputStream out = fileSystemFactory.create(ConnectorIdentity.ofUser("test")).newOutputFile(location).create()) { + out.write(Files.readAllBytes(parquetFile.toPath())); + } + + HivePageSourceFactory hivePageSourceFactory = new ParquetPageSourceFactory( + fileSystemFactory, + new FileFormatDataSourceStats(), + new ParquetReaderConfig(), + hiveConfig); + + return hivePageSourceFactory.createPageSource( + session, + location, + 0, + parquetFile.length(), + parquetFile.length(), + parquetFile.lastModified(), + new Schema(HiveStorageFormat.PARQUET.getSerde(), false, ImmutableMap.of()), + columns, + domain, + Optional.empty(), + OptionalInt.empty(), + false, + NO_ACID_TRANSACTION) + .orElseThrow() + .get(); + } + + private static List getBaseColumns(List columnNames, List columnTypes) + { + checkArgument(columnNames.size() == columnTypes.size(), "columnNames and columnTypes should have the same size"); + + return IntStream.range(0, columnNames.size()) + .mapToObj(index -> createBaseColumn( + columnNames.get(index), + index, + toHiveType(columnTypes.get(index)), + columnTypes.get(index), + REGULAR, + Optional.empty())) + .collect(toImmutableList()); + } +} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/parquet/ParquetWriterConfig.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/parquet/ParquetWriterConfig.java index d870a6e097c0..fdd1b1de99ac 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/parquet/ParquetWriterConfig.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/parquet/ParquetWriterConfig.java @@ -23,6 +23,8 @@ import io.trino.parquet.writer.ParquetWriterOptions; import jakarta.validation.constraints.DecimalMax; import jakarta.validation.constraints.DecimalMin; +import jakarta.validation.constraints.Max; +import jakarta.validation.constraints.Min; import org.apache.parquet.hadoop.ParquetWriter; @DefunctConfig({ @@ -35,9 +37,12 @@ public class ParquetWriterConfig public static final String PARQUET_WRITER_MAX_BLOCK_SIZE = "2GB"; public static final String PARQUET_WRITER_MIN_PAGE_SIZE = "8kB"; public static final String PARQUET_WRITER_MAX_PAGE_SIZE = "8MB"; + public static final int PARQUET_WRITER_MIN_PAGE_VALUE_COUNT = 1000; + public static final int PARQUET_WRITER_MAX_PAGE_VALUE_COUNT = 200_000; private DataSize blockSize = DataSize.ofBytes(ParquetWriter.DEFAULT_BLOCK_SIZE); private DataSize pageSize = DataSize.ofBytes(ParquetWriter.DEFAULT_PAGE_SIZE); + private int pageValueCount = ParquetWriterOptions.DEFAULT_MAX_PAGE_VALUE_COUNT; private int batchSize = ParquetWriterOptions.DEFAULT_BATCH_SIZE; private double validationPercentage = 5; @@ -70,6 +75,20 @@ public ParquetWriterConfig setPageSize(DataSize pageSize) return this; } + @Min(PARQUET_WRITER_MIN_PAGE_VALUE_COUNT) + @Max(PARQUET_WRITER_MAX_PAGE_VALUE_COUNT) + public int getPageValueCount() + { + return pageValueCount; + } + + @Config("parquet.writer.page-value-count") + public ParquetWriterConfig setPageValueCount(int pageValueCount) + { + this.pageValueCount = pageValueCount; + return this; + } + @Config("parquet.writer.batch-size") @ConfigDescription("Maximum number of rows passed to the writer in each batch") public ParquetWriterConfig setBatchSize(int batchSize) diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/procedure/CreateEmptyPartitionProcedure.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/procedure/CreateEmptyPartitionProcedure.java index 6a1096222001..891c7b568224 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/procedure/CreateEmptyPartitionProcedure.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/procedure/CreateEmptyPartitionProcedure.java @@ -21,7 +21,6 @@ import io.airlift.slice.Slices; import io.trino.plugin.hive.HiveColumnHandle; import io.trino.plugin.hive.HiveInsertTableHandle; -import io.trino.plugin.hive.HiveMetastoreClosure; import io.trino.plugin.hive.HiveTableHandle; import io.trino.plugin.hive.LocationService; import io.trino.plugin.hive.LocationService.WriteInfo; @@ -29,6 +28,7 @@ import io.trino.plugin.hive.PartitionUpdate.UpdateMode; import io.trino.plugin.hive.TransactionalMetadata; import io.trino.plugin.hive.TransactionalMetadataFactory; +import io.trino.plugin.hive.metastore.HiveMetastore; import io.trino.spi.TrinoException; import io.trino.spi.classloader.ThreadContextClassLoader; import io.trino.spi.connector.ConnectorAccessControl; @@ -124,8 +124,8 @@ private void doCreateEmptyPartition(ConnectorSession session, ConnectorAccessCon throw new TrinoException(INVALID_PROCEDURE_ARGUMENT, "Provided partition column names do not match actual partition column names: " + actualPartitionColumnNames); } - HiveMetastoreClosure metastore = hiveMetadata.getMetastore().unsafeGetRawHiveMetastoreClosure(); - if (metastore.getPartition(schemaName, tableName, partitionValues).isPresent()) { + HiveMetastore metastore = hiveMetadata.getMetastore().unsafeGetRawHiveMetastore(); + if (metastore.getTable(schemaName, tableName).flatMap(table -> metastore.getPartition(table, partitionValues)).isPresent()) { throw new TrinoException(ALREADY_EXISTS, "Partition already exists"); } HiveInsertTableHandle hiveInsertTableHandle = (HiveInsertTableHandle) hiveMetadata.beginInsert(session, tableHandle, ImmutableList.of(), NO_RETRIES); diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/procedure/DropStatsProcedure.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/procedure/DropStatsProcedure.java index 834784802597..ed6474cd9e4e 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/procedure/DropStatsProcedure.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/procedure/DropStatsProcedure.java @@ -14,20 +14,22 @@ package io.trino.plugin.hive.procedure; import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; import com.google.inject.Inject; import com.google.inject.Provider; import io.trino.plugin.hive.HiveColumnHandle; -import io.trino.plugin.hive.HiveMetastoreClosure; import io.trino.plugin.hive.HiveTableHandle; import io.trino.plugin.hive.PartitionStatistics; import io.trino.plugin.hive.TransactionalMetadata; import io.trino.plugin.hive.TransactionalMetadataFactory; +import io.trino.plugin.hive.metastore.HiveMetastore; import io.trino.spi.TrinoException; import io.trino.spi.classloader.ThreadContextClassLoader; import io.trino.spi.connector.ColumnHandle; import io.trino.spi.connector.ConnectorAccessControl; import io.trino.spi.connector.ConnectorSession; import io.trino.spi.connector.SchemaTableName; +import io.trino.spi.connector.TableNotFoundException; import io.trino.spi.predicate.TupleDomain; import io.trino.spi.procedure.Procedure; import io.trino.spi.procedure.Procedure.Argument; @@ -36,10 +38,11 @@ import java.lang.invoke.MethodHandle; import java.util.List; import java.util.Map; +import java.util.OptionalLong; import static com.google.common.collect.ImmutableList.toImmutableList; import static io.trino.plugin.base.util.Procedures.checkProcedureArgument; -import static io.trino.plugin.hive.acid.AcidTransaction.NO_ACID_TRANSACTION; +import static io.trino.plugin.hive.metastore.StatisticsUpdateMode.CLEAR_ALL; import static io.trino.plugin.hive.util.HiveUtil.makePartName; import static io.trino.spi.StandardErrorCode.INVALID_PROCEDURE_ARGUMENT; import static io.trino.spi.type.VarcharType.VARCHAR; @@ -100,6 +103,7 @@ private void doDropStats(ConnectorSession session, ConnectorAccessControl access checkProcedureArgument(table != null, "table_name cannot be null"); TransactionalMetadata hiveMetadata = hiveMetadataFactory.create(session.getIdentity(), true); + SchemaTableName schemaTableName = new SchemaTableName(schema, table); HiveTableHandle handle = (HiveTableHandle) hiveMetadata.getTableHandle(session, new SchemaTableName(schema, table)); if (handle == null) { throw new TrinoException(INVALID_PROCEDURE_ARGUMENT, format("Table '%s' does not exist", new SchemaTableName(schema, table))); @@ -114,7 +118,7 @@ private void doDropStats(ConnectorSession session, ConnectorAccessControl access .map(HiveColumnHandle::getName) .collect(toImmutableList()); - HiveMetastoreClosure metastore = hiveMetadata.getMetastore().unsafeGetRawHiveMetastoreClosure(); + HiveMetastore metastore = hiveMetadata.getMetastore().unsafeGetRawHiveMetastore(); if (partitionValues != null) { // drop stats for specified partitions List> partitionStringValues = partitionValues.stream() @@ -123,29 +127,34 @@ private void doDropStats(ConnectorSession session, ConnectorAccessControl access validatePartitions(partitionStringValues, partitionColumns); partitionStringValues.forEach(values -> metastore.updatePartitionStatistics( - schema, - table, - makePartName(partitionColumns, values), - stats -> PartitionStatistics.empty())); + metastore.getTable(schema, table) + .orElseThrow(() -> new TableNotFoundException(schemaTableName)), + CLEAR_ALL, + ImmutableMap.of( + makePartName(partitionColumns, values), + PartitionStatistics.empty()))); } else { // no partition specified, so drop stats for the entire table if (partitionColumns.isEmpty()) { // for non-partitioned tables, just wipe table stats metastore.updateTableStatistics( - schema, - table, - NO_ACID_TRANSACTION, - stats -> PartitionStatistics.empty()); + schema, + table, + OptionalLong.empty(), + CLEAR_ALL, + PartitionStatistics.empty()); } else { // the table is partitioned; remove stats for every partition - metastore.getPartitionNamesByFilter(handle.getSchemaName(), handle.getTableName(), partitionColumns, TupleDomain.all()) - .ifPresent(partitions -> partitions.forEach(partitionName -> metastore.updatePartitionStatistics( - schema, - table, - partitionName, - stats -> PartitionStatistics.empty()))); + hiveMetadata.getMetastore().getPartitionNamesByFilter(schemaTableName.getSchemaName(), schemaTableName.getTableName(), partitionColumns, TupleDomain.all()) + .ifPresent(partitions -> partitions.forEach(partitionName -> metastore.updatePartitionStatistics( + metastore.getTable(schema, table) + .orElseThrow(() -> new TableNotFoundException(schemaTableName)), + CLEAR_ALL, + ImmutableMap.of( + partitionName, + PartitionStatistics.empty())))); } } diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/procedure/RegisterPartitionProcedure.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/procedure/RegisterPartitionProcedure.java index 308655d520c1..bd294197ae3a 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/procedure/RegisterPartitionProcedure.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/procedure/RegisterPartitionProcedure.java @@ -18,15 +18,16 @@ import com.google.inject.Inject; import com.google.inject.Provider; import io.trino.filesystem.Location; -import io.trino.hdfs.HdfsContext; -import io.trino.hdfs.HdfsEnvironment; +import io.trino.filesystem.TrinoFileSystem; +import io.trino.filesystem.TrinoFileSystemFactory; +import io.trino.plugin.base.util.UncheckedCloseable; import io.trino.plugin.hive.HiveConfig; import io.trino.plugin.hive.PartitionStatistics; +import io.trino.plugin.hive.TransactionalMetadata; import io.trino.plugin.hive.TransactionalMetadataFactory; import io.trino.plugin.hive.metastore.Partition; import io.trino.plugin.hive.metastore.SemiTransactionalHiveMetastore; import io.trino.plugin.hive.metastore.Table; -import io.trino.plugin.hive.util.HiveWriteUtils; import io.trino.spi.TrinoException; import io.trino.spi.classloader.ThreadContextClassLoader; import io.trino.spi.connector.ConnectorAccessControl; @@ -35,14 +36,15 @@ import io.trino.spi.connector.TableNotFoundException; import io.trino.spi.procedure.Procedure; import io.trino.spi.type.ArrayType; -import org.apache.hadoop.fs.Path; +import java.io.IOException; import java.lang.invoke.MethodHandle; import java.util.List; import java.util.Optional; import static io.trino.plugin.base.util.Procedures.checkProcedureArgument; -import static io.trino.plugin.hive.HiveMetadata.PRESTO_QUERY_ID_NAME; +import static io.trino.plugin.hive.HiveErrorCode.HIVE_FILESYSTEM_ERROR; +import static io.trino.plugin.hive.HiveMetadata.TRINO_QUERY_ID_NAME; import static io.trino.plugin.hive.procedure.Procedures.checkIsPartitionedTable; import static io.trino.plugin.hive.procedure.Procedures.checkPartitionColumns; import static io.trino.plugin.hive.util.HiveUtil.makePartName; @@ -70,14 +72,14 @@ public class RegisterPartitionProcedure private final boolean allowRegisterPartition; private final TransactionalMetadataFactory hiveMetadataFactory; - private final HdfsEnvironment hdfsEnvironment; + private final TrinoFileSystemFactory fileSystemFactory; @Inject - public RegisterPartitionProcedure(HiveConfig hiveConfig, TransactionalMetadataFactory hiveMetadataFactory, HdfsEnvironment hdfsEnvironment) + public RegisterPartitionProcedure(HiveConfig hiveConfig, TransactionalMetadataFactory hiveMetadataFactory, TrinoFileSystemFactory fileSystemFactory) { this.allowRegisterPartition = hiveConfig.isAllowRegisterPartition(); this.hiveMetadataFactory = requireNonNull(hiveMetadataFactory, "hiveMetadataFactory is null"); - this.hdfsEnvironment = requireNonNull(hdfsEnvironment, "hdfsEnvironment is null"); + this.fileSystemFactory = requireNonNull(fileSystemFactory, "fileSystemFactory is null"); } @Override @@ -113,59 +115,63 @@ private void doRegisterPartition(ConnectorSession session, ConnectorAccessContro throw new TrinoException(PERMISSION_DENIED, "register_partition procedure is disabled"); } - SemiTransactionalHiveMetastore metastore = hiveMetadataFactory.create(session.getIdentity(), true).getMetastore(); - - HdfsContext hdfsContext = new HdfsContext(session); - SchemaTableName schemaTableName = new SchemaTableName(schemaName, tableName); - - Table table = metastore.getTable(schemaName, tableName) - .orElseThrow(() -> new TableNotFoundException(schemaTableName)); - - accessControl.checkCanInsertIntoTable(null, schemaTableName); - - checkIsPartitionedTable(table); - checkPartitionColumns(table, partitionColumns); - - Optional partition = metastore.unsafeGetRawHiveMetastoreClosure().getPartition(schemaName, tableName, partitionValues); - if (partition.isPresent()) { - String partitionName = makePartName(partitionColumns, partitionValues); - throw new TrinoException(ALREADY_EXISTS, format("Partition [%s] is already registered with location %s", partitionName, partition.get().getStorage().getLocation())); - } - - Path partitionLocation; - - if (location == null) { - partitionLocation = new Path(table.getStorage().getLocation(), makePartName(partitionColumns, partitionValues)); - } - else { - partitionLocation = new Path(location); + TransactionalMetadata hiveMetadata = hiveMetadataFactory.create(session.getIdentity(), true); + hiveMetadata.beginQuery(session); + try (UncheckedCloseable ignore = () -> hiveMetadata.cleanupQuery(session)) { + SemiTransactionalHiveMetastore metastore = hiveMetadata.getMetastore(); + + SchemaTableName schemaTableName = new SchemaTableName(schemaName, tableName); + + Table table = metastore.getTable(schemaName, tableName) + .orElseThrow(() -> new TableNotFoundException(schemaTableName)); + + accessControl.checkCanInsertIntoTable(null, schemaTableName); + + checkIsPartitionedTable(table); + checkPartitionColumns(table, partitionColumns); + + Optional partition = metastore.unsafeGetRawHiveMetastore().getPartition(table, partitionValues); + if (partition.isPresent()) { + String partitionName = makePartName(partitionColumns, partitionValues); + throw new TrinoException(ALREADY_EXISTS, format("Partition [%s] is already registered with location %s", partitionName, partition.get().getStorage().getLocation())); + } + + Location partitionLocation = Optional.ofNullable(location) + .map(Location::of) + .orElseGet(() -> Location.of(table.getStorage().getLocation()).appendPath(makePartName(partitionColumns, partitionValues))); + + TrinoFileSystem fileSystem = fileSystemFactory.create(session); + try { + if (!fileSystem.directoryExists(partitionLocation).orElse(true)) { + throw new TrinoException(INVALID_PROCEDURE_ARGUMENT, "Partition location does not exist: " + partitionLocation); + } + } + catch (IOException e) { + throw new TrinoException(HIVE_FILESYSTEM_ERROR, "Failed checking partition location: " + partitionLocation, e); + } + + metastore.addPartition( + session, + table.getDatabaseName(), + table.getTableName(), + buildPartitionObject(session, table, partitionValues, partitionLocation), + partitionLocation, + Optional.empty(), // no need for failed attempts cleanup + PartitionStatistics.empty(), + false); + + metastore.commit(); } - - if (!HiveWriteUtils.pathExists(hdfsContext, hdfsEnvironment, partitionLocation)) { - throw new TrinoException(INVALID_PROCEDURE_ARGUMENT, "Partition location does not exist: " + partitionLocation); - } - - metastore.addPartition( - session, - table.getDatabaseName(), - table.getTableName(), - buildPartitionObject(session, table, partitionValues, partitionLocation), - Location.of(partitionLocation.toString()), - Optional.empty(), // no need for failed attempts cleanup - PartitionStatistics.empty(), - false); - - metastore.commit(); } - private static Partition buildPartitionObject(ConnectorSession session, Table table, List partitionValues, Path location) + private static Partition buildPartitionObject(ConnectorSession session, Table table, List partitionValues, Location location) { return Partition.builder() .setDatabaseName(table.getDatabaseName()) .setTableName(table.getTableName()) .setColumns(table.getDataColumns()) .setValues(partitionValues) - .setParameters(ImmutableMap.of(PRESTO_QUERY_ID_NAME, session.getQueryId())) + .setParameters(ImmutableMap.of(TRINO_QUERY_ID_NAME, session.getQueryId())) .withStorage(storage -> storage .setStorageFormat(table.getStorage().getStorageFormat()) .setLocation(location.toString()) diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/procedure/SyncPartitionMetadataProcedure.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/procedure/SyncPartitionMetadataProcedure.java index f25b56bd9917..f5fa4f2241e2 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/procedure/SyncPartitionMetadataProcedure.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/procedure/SyncPartitionMetadataProcedure.java @@ -15,14 +15,17 @@ import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableSet; import com.google.common.collect.Lists; -import com.google.common.collect.Sets; import com.google.inject.Inject; import com.google.inject.Provider; import io.trino.filesystem.Location; -import io.trino.hdfs.HdfsContext; -import io.trino.hdfs.HdfsEnvironment; +import io.trino.filesystem.TrinoFileSystem; +import io.trino.filesystem.TrinoFileSystemFactory; +import io.trino.plugin.base.util.UncheckedCloseable; +import io.trino.plugin.hive.HiveConfig; import io.trino.plugin.hive.PartitionStatistics; +import io.trino.plugin.hive.TransactionalMetadata; import io.trino.plugin.hive.TransactionalMetadataFactory; import io.trino.plugin.hive.metastore.Column; import io.trino.plugin.hive.metastore.Partition; @@ -36,27 +39,26 @@ import io.trino.spi.connector.TableNotFoundException; import io.trino.spi.procedure.Procedure; import io.trino.spi.procedure.Procedure.Argument; -import org.apache.hadoop.fs.FileStatus; -import org.apache.hadoop.fs.FileSystem; import org.apache.hadoop.fs.Path; import java.io.IOException; import java.lang.invoke.MethodHandle; -import java.util.HashSet; import java.util.List; import java.util.Optional; import java.util.Set; -import java.util.stream.Stream; -import static com.google.common.collect.ImmutableList.toImmutableList; +import static com.google.common.base.Verify.verify; +import static com.google.common.collect.ImmutableSet.toImmutableSet; +import static com.google.common.collect.Sets.difference; import static io.trino.plugin.base.util.Procedures.checkProcedureArgument; import static io.trino.plugin.hive.HiveErrorCode.HIVE_FILESYSTEM_ERROR; -import static io.trino.plugin.hive.HiveMetadata.PRESTO_QUERY_ID_NAME; +import static io.trino.plugin.hive.HiveMetadata.TRINO_QUERY_ID_NAME; import static io.trino.plugin.hive.HivePartitionManager.extractPartitionValues; import static io.trino.spi.StandardErrorCode.INVALID_PROCEDURE_ARGUMENT; import static io.trino.spi.type.BooleanType.BOOLEAN; import static io.trino.spi.type.VarcharType.VARCHAR; import static java.lang.Boolean.TRUE; +import static java.lang.String.join; import static java.lang.invoke.MethodHandles.lookup; import static java.util.Locale.ENGLISH; import static java.util.Objects.requireNonNull; @@ -83,15 +85,18 @@ public enum SyncMode } private final TransactionalMetadataFactory hiveMetadataFactory; - private final HdfsEnvironment hdfsEnvironment; + private final TrinoFileSystemFactory fileSystemFactory; + private final int maxPartitionBatchSize; @Inject public SyncPartitionMetadataProcedure( TransactionalMetadataFactory hiveMetadataFactory, - HdfsEnvironment hdfsEnvironment) + TrinoFileSystemFactory fileSystemFactory, + HiveConfig config) { this.hiveMetadataFactory = requireNonNull(hiveMetadataFactory, "hiveMetadataFactory is null"); - this.hdfsEnvironment = requireNonNull(hdfsEnvironment, "hdfsEnvironment is null"); + this.fileSystemFactory = requireNonNull(fileSystemFactory, "fileSystemFactory is null"); + maxPartitionBatchSize = config.getMaxPartitionBatchSize(); } @Override @@ -122,94 +127,136 @@ private void doSyncPartitionMetadata(ConnectorSession session, ConnectorAccessCo checkProcedureArgument(mode != null, "mode cannot be null"); SyncMode syncMode = toSyncMode(mode); - HdfsContext hdfsContext = new HdfsContext(session); - SemiTransactionalHiveMetastore metastore = hiveMetadataFactory.create(session.getIdentity(), true).getMetastore(); - SchemaTableName schemaTableName = new SchemaTableName(schemaName, tableName); - - Table table = metastore.getTable(schemaName, tableName) - .orElseThrow(() -> new TableNotFoundException(schemaTableName)); - if (table.getPartitionColumns().isEmpty()) { - throw new TrinoException(INVALID_PROCEDURE_ARGUMENT, "Table is not partitioned: " + schemaTableName); - } - - if (syncMode == SyncMode.ADD || syncMode == SyncMode.FULL) { - accessControl.checkCanInsertIntoTable(null, new SchemaTableName(schemaName, tableName)); - } - if (syncMode == SyncMode.DROP || syncMode == SyncMode.FULL) { - accessControl.checkCanDeleteFromTable(null, new SchemaTableName(schemaName, tableName)); - } + TransactionalMetadata hiveMetadata = hiveMetadataFactory.create(session.getIdentity(), true); + hiveMetadata.beginQuery(session); + try (UncheckedCloseable ignore = () -> hiveMetadata.cleanupQuery(session)) { + SemiTransactionalHiveMetastore metastore = hiveMetadataFactory.create(session.getIdentity(), true).getMetastore(); + SchemaTableName schemaTableName = new SchemaTableName(schemaName, tableName); - Path tableLocation = new Path(table.getStorage().getLocation()); + Table table = metastore.getTable(schemaName, tableName) + .orElseThrow(() -> new TableNotFoundException(schemaTableName)); + if (table.getPartitionColumns().isEmpty()) { + throw new TrinoException(INVALID_PROCEDURE_ARGUMENT, "Table is not partitioned: " + schemaTableName); + } - Set partitionsToAdd; - Set partitionsToDrop; + if (syncMode == SyncMode.ADD || syncMode == SyncMode.FULL) { + accessControl.checkCanInsertIntoTable(null, new SchemaTableName(schemaName, tableName)); + } + if (syncMode == SyncMode.DROP || syncMode == SyncMode.FULL) { + accessControl.checkCanDeleteFromTable(null, new SchemaTableName(schemaName, tableName)); + } - try { - FileSystem fileSystem = hdfsEnvironment.getFileSystem(hdfsContext, tableLocation); - List partitionsNamesInMetastore = metastore.getPartitionNames(schemaName, tableName) + Set partitionNamesInMetastore = metastore.getPartitionNames(schemaName, tableName) + .map(ImmutableSet::copyOf) .orElseThrow(() -> new TableNotFoundException(schemaTableName)); - List partitionsInMetastore = getPartitionsInMetastore(schemaTableName, tableLocation, partitionsNamesInMetastore, metastore); - List partitionsInFileSystem = listDirectory(fileSystem, fileSystem.getFileStatus(tableLocation), table.getPartitionColumns(), table.getPartitionColumns().size(), caseSensitive).stream() - .map(fileStatus -> fileStatus.getPath().toUri()) - .map(uri -> tableLocation.toUri().relativize(uri).getPath()) - .collect(toImmutableList()); + String tableStorageLocation = table.getStorage().getLocation(); + Set canonicalPartitionNamesInMetastore = partitionNamesInMetastore; + if (!caseSensitive) { + canonicalPartitionNamesInMetastore = Lists.partition(ImmutableList.copyOf(partitionNamesInMetastore), maxPartitionBatchSize).stream() + .flatMap(partitionNames -> metastore.getPartitionsByNames(schemaName, tableName, partitionNames).values().stream()) + .flatMap(Optional::stream) // disregard partitions which disappeared in the meantime since listing the partition names + // Disregard the partitions which do not have a canonical Hive location (e.g. `ALTER TABLE ... ADD PARTITION (...) LOCATION '...'`) + .flatMap(partition -> getCanonicalPartitionName(partition, table.getPartitionColumns(), tableStorageLocation).stream()) + .collect(toImmutableSet()); + } + Set partitionsInFileSystem = listPartitions(fileSystemFactory.create(session), Location.of(tableStorageLocation), table.getPartitionColumns(), caseSensitive); // partitions in file system but not in metastore - partitionsToAdd = difference(partitionsInFileSystem, partitionsInMetastore); + Set partitionsToAdd = difference(partitionsInFileSystem, canonicalPartitionNamesInMetastore); + // partitions in metastore but not in file system - partitionsToDrop = difference(partitionsInMetastore, partitionsInFileSystem); + Set partitionsToDrop = difference(canonicalPartitionNamesInMetastore, partitionsInFileSystem); + + syncPartitions(partitionsToAdd, partitionsToDrop, syncMode, metastore, session, table); } - catch (IOException e) { - throw new TrinoException(HIVE_FILESYSTEM_ERROR, e); + } + + private static Optional getCanonicalPartitionName(Partition partition, List partitionColumns, String tableLocation) + { + String partitionStorageLocation = partition.getStorage().getLocation(); + if (!partitionStorageLocation.startsWith(tableLocation)) { + return Optional.empty(); + } + + String partitionName = partitionStorageLocation.substring(tableLocation.length()); + + if (partitionName.startsWith("/")) { + // Remove eventual forward slash from the name of the partition + partitionName = partitionName.replaceFirst("^/+", ""); + } + if (partitionName.endsWith("/")) { + // Remove eventual trailing slash from the name of the partition + partitionName = partitionName.replaceFirst("/+$", ""); + } + + // Ensure that the partition location is corresponding to a canonical Hive partition location + String[] partitionDirectories = partitionName.split("/"); + if (partitionDirectories.length != partitionColumns.size()) { + return Optional.empty(); + } + for (int i = 0; i < partitionDirectories.length; i++) { + String partitionDirectory = partitionDirectories[i]; + Column column = partitionColumns.get(i); + if (!isValidPartitionPath(partitionDirectory, column, false)) { + return Optional.empty(); + } } - syncPartitions(partitionsToAdd, partitionsToDrop, syncMode, metastore, session, table); + return Optional.of(partitionName); } - private List getPartitionsInMetastore(SchemaTableName schemaTableName, Path tableLocation, List partitionsNames, SemiTransactionalHiveMetastore metastore) + private static Set listPartitions(TrinoFileSystem fileSystem, Location directory, List partitionColumns, boolean caseSensitive) { - ImmutableList.Builder partitionsInMetastoreBuilder = ImmutableList.builderWithExpectedSize(partitionsNames.size()); - for (List partitionsNamesBatch : Lists.partition(partitionsNames, BATCH_GET_PARTITIONS_BY_NAMES_MAX_PAGE_SIZE)) { - metastore.getPartitionsByNames(schemaTableName.getSchemaName(), schemaTableName.getTableName(), partitionsNamesBatch).values().stream() - .filter(Optional::isPresent).map(Optional::get) - .map(partition -> new Path(partition.getStorage().getLocation()).toUri()) - .map(uri -> tableLocation.toUri().relativize(uri).getPath()) - .forEach(partitionsInMetastoreBuilder::add); - } - return partitionsInMetastoreBuilder.build(); + return doListPartitions(fileSystem, directory, partitionColumns, partitionColumns.size(), caseSensitive, ImmutableList.of()); } - private static List listDirectory(FileSystem fileSystem, FileStatus current, List partitionColumns, int depth, boolean caseSensitive) + private static Set doListPartitions(TrinoFileSystem fileSystem, Location directory, List partitionColumns, int depth, boolean caseSensitive, List partitions) { if (depth == 0) { - return ImmutableList.of(current); + return ImmutableSet.of(join("/", partitions)); + } + + ImmutableSet.Builder result = ImmutableSet.builder(); + for (Location location : listDirectories(fileSystem, directory)) { + String path = listedDirectoryName(directory, location); + Column column = partitionColumns.get(partitionColumns.size() - depth); + if (!isValidPartitionPath(path, column, caseSensitive)) { + continue; + } + List current = ImmutableList.builder().addAll(partitions).add(path).build(); + result.addAll(doListPartitions(fileSystem, location, partitionColumns, depth - 1, caseSensitive, current)); } + return result.build(); + } + private static Set listDirectories(TrinoFileSystem fileSystem, Location directory) + { try { - return Stream.of(fileSystem.listStatus(current.getPath())) - .filter(fileStatus -> isValidPartitionPath(fileStatus, partitionColumns.get(partitionColumns.size() - depth), caseSensitive)) - .flatMap(directory -> listDirectory(fileSystem, directory, partitionColumns, depth - 1, caseSensitive).stream()) - .collect(toImmutableList()); + return fileSystem.listDirectories(directory); } catch (IOException e) { throw new TrinoException(HIVE_FILESYSTEM_ERROR, e); } } - private static boolean isValidPartitionPath(FileStatus file, Column column, boolean caseSensitive) + private static String listedDirectoryName(Location directory, Location location) { - String path = file.getPath().getName(); - if (!caseSensitive) { - path = path.toLowerCase(ENGLISH); + String prefix = directory.path(); + if (!prefix.isEmpty() && !prefix.endsWith("/")) { + prefix += "/"; } - String prefix = column.getName() + '='; - return file.isDirectory() && path.startsWith(prefix); + String path = location.path(); + verify(path.endsWith("/"), "path does not end with slash: %s", location); + verify(path.startsWith(prefix), "path [%s] is not a child of directory [%s]", location, directory); + return path.substring(prefix.length(), path.length() - 1); } - // calculate relative complement of set b with respect to set a - private static Set difference(List a, List b) + private static boolean isValidPartitionPath(String path, Column column, boolean caseSensitive) { - return Sets.difference(new HashSet<>(a), new HashSet<>(b)); + if (!caseSensitive) { + path = path.toLowerCase(ENGLISH); + } + return path.startsWith(column.getName() + '='); } private static void syncPartitions( @@ -271,7 +318,7 @@ private static Partition buildPartitionObject(ConnectorSession session, Table ta .setTableName(table.getTableName()) .setColumns(table.getDataColumns()) .setValues(extractPartitionValues(partitionName)) - .setParameters(ImmutableMap.of(PRESTO_QUERY_ID_NAME, session.getQueryId())) + .setParameters(ImmutableMap.of(TRINO_QUERY_ID_NAME, session.getQueryId())) .withStorage(storage -> storage .setStorageFormat(table.getStorage().getStorageFormat()) .setLocation(new Path(table.getStorage().getLocation(), partitionName).toString()) diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/procedure/UnregisterPartitionProcedure.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/procedure/UnregisterPartitionProcedure.java index 2ff444f2f30e..31c5461f664f 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/procedure/UnregisterPartitionProcedure.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/procedure/UnregisterPartitionProcedure.java @@ -106,7 +106,7 @@ private void doUnregisterPartition(ConnectorSession session, ConnectorAccessCont String partitionName = makePartName(partitionColumns, partitionValues); - Partition partition = metastore.unsafeGetRawHiveMetastoreClosure().getPartition(schemaName, tableName, partitionValues) + Partition partition = metastore.unsafeGetRawHiveMetastore().getPartition(table, partitionValues) .orElseThrow(() -> new TrinoException(NOT_FOUND, format("Partition '%s' does not exist", partitionName))); metastore.dropPartition( diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/aws/athena/projection/DateProjectionFactory.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/projection/DateProjection.java similarity index 56% rename from plugin/trino-hive/src/main/java/io/trino/plugin/hive/aws/athena/projection/DateProjectionFactory.java rename to plugin/trino-hive/src/main/java/io/trino/plugin/hive/projection/DateProjection.java index 8f6914a2f6e8..18cc546b236d 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/aws/athena/projection/DateProjectionFactory.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/projection/DateProjection.java @@ -11,20 +11,25 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package io.trino.plugin.hive.aws.athena.projection; +package io.trino.plugin.hive.projection; +import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableSet; import io.trino.spi.TrinoException; +import io.trino.spi.predicate.Domain; import io.trino.spi.type.DateType; import io.trino.spi.type.TimestampType; import io.trino.spi.type.Type; import io.trino.spi.type.VarcharType; +import java.text.DateFormat; import java.text.ParseException; import java.text.SimpleDateFormat; import java.time.Instant; import java.time.ZoneId; +import java.time.temporal.ChronoField; import java.time.temporal.ChronoUnit; +import java.util.Date; import java.util.List; import java.util.Locale; import java.util.Map; @@ -35,13 +40,14 @@ import java.util.regex.Pattern; import static com.google.common.collect.ImmutableList.toImmutableList; -import static io.trino.plugin.hive.aws.athena.PartitionProjectionProperties.COLUMN_PROJECTION_FORMAT; -import static io.trino.plugin.hive.aws.athena.PartitionProjectionProperties.COLUMN_PROJECTION_INTERVAL; -import static io.trino.plugin.hive.aws.athena.PartitionProjectionProperties.COLUMN_PROJECTION_INTERVAL_UNIT; -import static io.trino.plugin.hive.aws.athena.PartitionProjectionProperties.COLUMN_PROJECTION_RANGE; -import static io.trino.plugin.hive.aws.athena.PartitionProjectionProperties.getProjectionPropertyRequiredValue; -import static io.trino.plugin.hive.aws.athena.PartitionProjectionProperties.getProjectionPropertyValue; -import static io.trino.plugin.hive.aws.athena.projection.Projection.invalidProjectionException; +import static io.airlift.slice.Slices.utf8Slice; +import static io.trino.plugin.hive.projection.PartitionProjectionProperties.COLUMN_PROJECTION_FORMAT; +import static io.trino.plugin.hive.projection.PartitionProjectionProperties.COLUMN_PROJECTION_INTERVAL; +import static io.trino.plugin.hive.projection.PartitionProjectionProperties.COLUMN_PROJECTION_INTERVAL_UNIT; +import static io.trino.plugin.hive.projection.PartitionProjectionProperties.COLUMN_PROJECTION_RANGE; +import static io.trino.plugin.hive.projection.PartitionProjectionProperties.getProjectionPropertyRequiredValue; +import static io.trino.plugin.hive.projection.PartitionProjectionProperties.getProjectionPropertyValue; +import static io.trino.spi.predicate.Domain.singleValue; import static java.lang.String.format; import static java.time.temporal.ChronoUnit.DAYS; import static java.time.temporal.ChronoUnit.HOURS; @@ -49,31 +55,38 @@ import static java.time.temporal.ChronoUnit.MONTHS; import static java.time.temporal.ChronoUnit.SECONDS; import static java.util.Objects.nonNull; +import static java.util.Objects.requireNonNull; import static java.util.TimeZone.getTimeZone; +import static java.util.concurrent.TimeUnit.MILLISECONDS; -public class DateProjectionFactory - implements ProjectionFactory +final class DateProjection + implements Projection { - public static final ZoneId UTC_TIME_ZONE_ID = ZoneId.of("UTC"); - + private static final ZoneId UTC_TIME_ZONE_ID = ZoneId.of("UTC"); // Limited to only DAYS, HOURS, MINUTES, SECONDS as we are not fully sure how everything above day // is implemented in Athena. So we limit it to a subset of interval units which are explicitly clear how to calculate. - // Rest will be implemented if this will be required as it would require making compatibility tests + // The rest will be implemented if this is required as it would require making compatibility tests // for results received from Athena and verifying if we receive identical with Trino. private static final Set DATE_PROJECTION_INTERVAL_UNITS = ImmutableSet.of(DAYS, HOURS, MINUTES, SECONDS); private static final Pattern DATE_RANGE_BOUND_EXPRESSION_PATTERN = Pattern.compile("^\\s*NOW\\s*(([+-])\\s*([0-9]+)\\s*(DAY|HOUR|MINUTE|SECOND)S?\\s*)?$"); - @Override - public boolean isSupportedColumnType(Type columnType) - { - return columnType instanceof VarcharType - || columnType instanceof DateType - || (columnType instanceof TimestampType && ((TimestampType) columnType).isShort()); - } + private final String columnName; + private final DateFormat dateFormat; + private final Supplier leftBound; + private final Supplier rightBound; + private final int interval; + private final ChronoUnit intervalUnit; - @Override - public Projection create(String columnName, Type columnType, Map columnProperties) + public DateProjection(String columnName, Type columnType, Map columnProperties) { + if (!(columnType instanceof VarcharType) && + !(columnType instanceof DateType) && + !(columnType instanceof TimestampType timestampType && timestampType.isShort())) { + throw new InvalidProjectionException(columnName, columnType); + } + + this.columnName = requireNonNull(columnName, "columnName is null"); + String dateFormatPattern = getProjectionPropertyRequiredValue( columnName, columnProperties, @@ -88,24 +101,26 @@ public Projection create(String columnName, Type columnType, Map .map(String.class::cast) .collect(toImmutableList())); if (range.size() != 2) { - throw invalidRangeProperty(columnName, dateFormatPattern); + throw invalidRangeProperty(columnName, dateFormatPattern, Optional.empty()); } SimpleDateFormat dateFormat = new SimpleDateFormat(dateFormatPattern); dateFormat.setLenient(false); dateFormat.setTimeZone(getTimeZone(UTC_TIME_ZONE_ID)); - Supplier leftBound = parseDateRangerBound(columnName, range.get(0), dateFormat); - Supplier rangeBound = parseDateRangerBound(columnName, range.get(1), dateFormat); - if (!leftBound.get().isBefore(rangeBound.get())) { - throw invalidRangeProperty(columnName, dateFormatPattern); + this.dateFormat = requireNonNull(dateFormat, "dateFormatPattern is null"); + + leftBound = parseDateRangerBound(columnName, range.get(0), dateFormat); + rightBound = parseDateRangerBound(columnName, range.get(1), dateFormat); + if (!leftBound.get().isBefore(rightBound.get())) { + throw invalidRangeProperty(columnName, dateFormatPattern, Optional.empty()); } - int interval = getProjectionPropertyValue(columnProperties, COLUMN_PROJECTION_INTERVAL, Integer.class::cast).orElse(1); - ChronoUnit intervalUnit = getProjectionPropertyValue(columnProperties, COLUMN_PROJECTION_INTERVAL_UNIT, ChronoUnit.class::cast) + interval = getProjectionPropertyValue(columnProperties, COLUMN_PROJECTION_INTERVAL, Integer.class::cast).orElse(1); + intervalUnit = getProjectionPropertyValue(columnProperties, COLUMN_PROJECTION_INTERVAL_UNIT, ChronoUnit.class::cast) .orElseGet(() -> resolveDefaultChronoUnit(columnName, dateFormatPattern)); if (!DATE_PROJECTION_INTERVAL_UNITS.contains(intervalUnit)) { - throw invalidProjectionException( + throw new InvalidProjectionException( columnName, format( "Property: '%s' value '%s' is invalid. Available options: %s", @@ -113,17 +128,72 @@ public Projection create(String columnName, Type columnType, Map intervalUnit, DATE_PROJECTION_INTERVAL_UNITS)); } + } - return new DateProjection(columnName, dateFormat, leftBound, rangeBound, interval, intervalUnit); + @Override + public List getProjectedValues(Optional partitionValueFilter) + { + ImmutableList.Builder builder = ImmutableList.builder(); + + Instant leftBound = adjustBoundToDateFormat(this.leftBound.get()); + Instant rightBound = adjustBoundToDateFormat(this.rightBound.get()); + + Instant currentValue = leftBound; + while (!currentValue.isAfter(rightBound)) { + String currentValueFormatted = formatValue(currentValue); + if (isValueInDomain(partitionValueFilter, currentValue, currentValueFormatted)) { + builder.add(currentValueFormatted); + } + currentValue = currentValue.atZone(UTC_TIME_ZONE_ID) + .plus(interval, intervalUnit) + .toInstant(); + } + + return builder.build(); } - private ChronoUnit resolveDefaultChronoUnit(String columnName, String dateFormatPattern) + private Instant adjustBoundToDateFormat(Instant value) + { + String formatted = formatValue(value.with(ChronoField.MILLI_OF_SECOND, 0)); + try { + return dateFormat.parse(formatted).toInstant(); + } + catch (ParseException e) { + throw new InvalidProjectionException(formatted, e.getMessage()); + } + } + + private String formatValue(Instant current) + { + return dateFormat.format(new Date(current.toEpochMilli())); + } + + private boolean isValueInDomain(Optional valueDomain, Instant value, String formattedValue) + { + if (valueDomain.isEmpty() || valueDomain.get().isAll()) { + return true; + } + Domain domain = valueDomain.get(); + Type type = domain.getType(); + if (type instanceof VarcharType) { + return domain.contains(singleValue(type, utf8Slice(formattedValue))); + } + if (type instanceof DateType) { + return domain.contains(singleValue(type, MILLISECONDS.toDays(value.toEpochMilli()))); + } + if (type instanceof TimestampType && ((TimestampType) type).isShort()) { + return domain.contains(singleValue(type, MILLISECONDS.toMicros(value.toEpochMilli()))); + } + throw new InvalidProjectionException(columnName, type); + } + + private static ChronoUnit resolveDefaultChronoUnit(String columnName, String dateFormatPattern) { String datePatternWithoutText = dateFormatPattern.replaceAll("'.*?'", ""); if (datePatternWithoutText.contains("S") || datePatternWithoutText.contains("s") || datePatternWithoutText.contains("m") || datePatternWithoutText.contains("H")) { // When the provided dates are at single-day or single-month precision. - throw invalidProjectionException( + throw new InvalidProjectionException( columnName, format( "Property: '%s' needs to be set when provided '%s' is less that single-day precision. Interval defaults to 1 day or 1 month, respectively. Otherwise, interval is required", @@ -136,7 +206,7 @@ private ChronoUnit resolveDefaultChronoUnit(String columnName, String dateFormat return MONTHS; } - private Supplier parseDateRangerBound(String columnName, String value, SimpleDateFormat dateFormat) + private static Supplier parseDateRangerBound(String columnName, String value, SimpleDateFormat dateFormat) { Matcher matcher = DATE_RANGE_BOUND_EXPRESSION_PATTERN.matcher(value); if (matcher.matches()) { @@ -167,37 +237,21 @@ private Supplier parseDateRangerBound(String columnName, String value, return () -> dateBound; } - private TrinoException invalidRangeProperty(String columnName, String dateFormatPattern) - { - return invalidRangeProperty(columnName, dateFormatPattern, Optional.empty()); - } - - private TrinoException invalidRangeProperty(String columnName, String dateFormatPattern, Optional errorDetail) + private static TrinoException invalidRangeProperty(String columnName, String dateFormatPattern, Optional errorDetail) { - return invalidProjectionException( + throw new InvalidProjectionException( columnName, format( "Property: '%s' needs to be a list of 2 valid dates formatted as '%s' or '%s' that are sequential%s", COLUMN_PROJECTION_RANGE, dateFormatPattern, DATE_RANGE_BOUND_EXPRESSION_PATTERN.pattern(), - errorDetail.map(error -> ". " + error).orElse(""))); + errorDetail.map(error -> ": " + error).orElse(""))); } - private static class DateExpressionBound + private record DateExpressionBound(int multiplier, ChronoUnit unit, boolean increment) implements Supplier { - private final int multiplier; - private final ChronoUnit unit; - private final boolean increment; - - public DateExpressionBound(int multiplier, ChronoUnit unit, boolean increment) - { - this.multiplier = multiplier; - this.unit = unit; - this.increment = increment; - } - @Override public Instant get() { diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/aws/athena/projection/EnumProjection.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/projection/EnumProjection.java similarity index 58% rename from plugin/trino-hive/src/main/java/io/trino/plugin/hive/aws/athena/projection/EnumProjection.java rename to plugin/trino-hive/src/main/java/io/trino/plugin/hive/projection/EnumProjection.java index 50eb7d1a63fd..45af272920a8 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/aws/athena/projection/EnumProjection.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/projection/EnumProjection.java @@ -11,30 +11,44 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package io.trino.plugin.hive.aws.athena.projection; +package io.trino.plugin.hive.projection; -import com.google.common.collect.ImmutableList; import io.trino.spi.predicate.Domain; import io.trino.spi.type.Type; import io.trino.spi.type.VarcharType; import java.util.List; +import java.util.Map; import java.util.Optional; +import static com.google.common.collect.ImmutableList.toImmutableList; import static io.airlift.slice.Slices.utf8Slice; +import static io.trino.plugin.hive.projection.PartitionProjectionProperties.COLUMN_PROJECTION_VALUES; +import static io.trino.plugin.hive.projection.PartitionProjectionProperties.getProjectionPropertyRequiredValue; import static io.trino.spi.predicate.Domain.singleValue; import static java.util.Objects.requireNonNull; import static java.util.stream.Collectors.toList; -public class EnumProjection - extends Projection +final class EnumProjection + implements Projection { + private final String columnName; private final List values; - public EnumProjection(String columnName, List values) + public EnumProjection(String columnName, Type columnType, Map columnProperties) { - super(columnName); - this.values = ImmutableList.copyOf(requireNonNull(values, "values is null")); + if (!(columnType instanceof VarcharType)) { + throw new InvalidProjectionException(columnName, columnType); + } + + this.columnName = requireNonNull(columnName, "columnName is null"); + this.values = getProjectionPropertyRequiredValue( + columnName, + columnProperties, + COLUMN_PROJECTION_VALUES, + value -> ((List) value).stream() + .map(String::valueOf) + .collect(toImmutableList())); } @Override @@ -54,6 +68,6 @@ private boolean isValueInDomain(Domain valueDomain, String value) if (type instanceof VarcharType) { return valueDomain.contains(singleValue(type, utf8Slice(value))); } - throw unsupportedProjectionColumnTypeException(type); + throw new InvalidProjectionException(columnName, type); } } diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/aws/athena/projection/InjectedProjection.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/projection/InjectedProjection.java similarity index 59% rename from plugin/trino-hive/src/main/java/io/trino/plugin/hive/aws/athena/projection/InjectedProjection.java rename to plugin/trino-hive/src/main/java/io/trino/plugin/hive/projection/InjectedProjection.java index 3c0d3c9ae423..acead53bf047 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/aws/athena/projection/InjectedProjection.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/projection/InjectedProjection.java @@ -11,7 +11,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package io.trino.plugin.hive.aws.athena.projection; +package io.trino.plugin.hive.projection; import com.google.common.collect.ImmutableList; import io.trino.spi.predicate.Domain; @@ -22,26 +22,32 @@ import static io.trino.plugin.hive.metastore.MetastoreUtil.canConvertSqlTypeToStringForParts; import static io.trino.plugin.hive.metastore.MetastoreUtil.sqlScalarToString; +import static java.util.Objects.requireNonNull; -public class InjectedProjection - extends Projection +final class InjectedProjection + implements Projection { - public InjectedProjection(String columnName) + private final String columnName; + + public InjectedProjection(String columnName, Type columnType) { - super(columnName); + if (!canConvertSqlTypeToStringForParts(columnType, true)) { + throw new InvalidProjectionException(columnName, columnType); + } + this.columnName = requireNonNull(columnName, "columnName is null"); } @Override public List getProjectedValues(Optional partitionValueFilter) { Domain domain = partitionValueFilter - .orElseThrow(() -> invalidProjectionException(getColumnName(), "Injected projection requires single predicate for it's column in where clause")); + .orElseThrow(() -> new InvalidProjectionException(columnName, "Injected projection requires single predicate for it's column in where clause")); Type type = domain.getType(); if (!domain.isNullableSingleValue() || !canConvertSqlTypeToStringForParts(type, true)) { - throw invalidProjectionException(getColumnName(), "Injected projection requires single predicate for it's column in where clause. Currently provided can't be converted to single partition."); + throw new InvalidProjectionException(columnName, "Injected projection requires single predicate for it's column in where clause. Currently provided can't be converted to single partition."); } return Optional.ofNullable(sqlScalarToString(type, domain.getNullableSingleValue(), null)) .map(ImmutableList::of) - .orElseThrow(() -> unsupportedProjectionColumnTypeException(type)); + .orElseThrow(() -> new InvalidProjectionException(columnName, type)); } } diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/aws/athena/projection/IntegerProjection.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/projection/IntegerProjection.java similarity index 51% rename from plugin/trino-hive/src/main/java/io/trino/plugin/hive/aws/athena/projection/IntegerProjection.java rename to plugin/trino-hive/src/main/java/io/trino/plugin/hive/projection/IntegerProjection.java index 645b94ca6908..2b07e78bbf15 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/aws/athena/projection/IntegerProjection.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/projection/IntegerProjection.java @@ -11,7 +11,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package io.trino.plugin.hive.aws.athena.projection; +package io.trino.plugin.hive.projection; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableList.Builder; @@ -22,27 +22,52 @@ import io.trino.spi.type.VarcharType; import java.util.List; +import java.util.Map; import java.util.Optional; +import static com.google.common.collect.ImmutableList.toImmutableList; import static io.airlift.slice.Slices.utf8Slice; +import static io.trino.plugin.hive.projection.PartitionProjectionProperties.COLUMN_PROJECTION_DIGITS; +import static io.trino.plugin.hive.projection.PartitionProjectionProperties.COLUMN_PROJECTION_INTERVAL; +import static io.trino.plugin.hive.projection.PartitionProjectionProperties.COLUMN_PROJECTION_RANGE; +import static io.trino.plugin.hive.projection.PartitionProjectionProperties.getProjectionPropertyRequiredValue; +import static io.trino.plugin.hive.projection.PartitionProjectionProperties.getProjectionPropertyValue; import static io.trino.spi.predicate.Domain.singleValue; +import static java.lang.String.format; import static java.util.Objects.requireNonNull; -public class IntegerProjection - extends Projection +final class IntegerProjection + implements Projection { + private final String columnName; private final int leftBound; private final int rightBound; private final int interval; private final Optional digits; - public IntegerProjection(String columnName, int leftBound, int rightBound, int interval, Optional digits) + public IntegerProjection(String columnName, Type columnType, Map columnProperties) { - super(columnName); - this.leftBound = leftBound; - this.rightBound = rightBound; - this.interval = interval; - this.digits = requireNonNull(digits, "digits is null"); + if (!(columnType instanceof VarcharType) && !(columnType instanceof IntegerType) && !(columnType instanceof BigintType)) { + throw new InvalidProjectionException(columnName, columnType); + } + + this.columnName = requireNonNull(columnName, "columnName is null"); + + List range = getProjectionPropertyRequiredValue( + columnName, + columnProperties, + COLUMN_PROJECTION_RANGE, + value -> ((List) value).stream() + .map(element -> Integer.valueOf((String) element)) + .collect(toImmutableList())); + if (range.size() != 2) { + throw new InvalidProjectionException(columnName, format("Property: '%s' needs to be list of 2 integers", COLUMN_PROJECTION_RANGE)); + } + this.leftBound = range.get(0); + this.rightBound = range.get(1); + + this.interval = getProjectionPropertyValue(columnProperties, COLUMN_PROJECTION_INTERVAL, Integer.class::cast).orElse(1); + this.digits = getProjectionPropertyValue(columnProperties, COLUMN_PROJECTION_DIGITS, Integer.class::cast); } @Override @@ -53,7 +78,7 @@ public List getProjectedValues(Optional partitionValueFilter) while (current <= rightBound) { int currentValue = current; String currentValueFormatted = digits - .map(digits -> String.format("%0" + digits + "d", currentValue)) + .map(digits -> format("%0" + digits + "d", currentValue)) .orElseGet(() -> Integer.toString(currentValue)); if (isValueInDomain(partitionValueFilter, current, currentValueFormatted)) { builder.add(currentValueFormatted); @@ -74,8 +99,8 @@ private boolean isValueInDomain(Optional valueDomain, int value, String return domain.contains(singleValue(type, utf8Slice(formattedValue))); } if (type instanceof IntegerType || type instanceof BigintType) { - return domain.contains(singleValue(type, Long.valueOf(value))); + return domain.contains(singleValue(type, (long) value)); } - throw unsupportedProjectionColumnTypeException(type); + throw new InvalidProjectionException(columnName, type); } } diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/projection/InvalidProjectionException.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/projection/InvalidProjectionException.java new file mode 100644 index 000000000000..938152015700 --- /dev/null +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/projection/InvalidProjectionException.java @@ -0,0 +1,44 @@ +/* + * 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.hive.projection; + +import io.trino.spi.TrinoException; +import io.trino.spi.type.Type; + +import static io.trino.spi.StandardErrorCode.INVALID_COLUMN_PROPERTY; +import static java.lang.String.format; + +public class InvalidProjectionException + extends TrinoException +{ + public InvalidProjectionException(String columnName, Type columnType) + { + this(columnName, "Unsupported column type: " + columnType.getDisplayName()); + } + + public InvalidProjectionException(String columnName, String message) + { + this(invalidProjectionMessage(columnName, message)); + } + + public InvalidProjectionException(String message) + { + super(INVALID_COLUMN_PROPERTY, message); + } + + public static String invalidProjectionMessage(String columnName, String message) + { + return format("Column projection for column '%s' failed. %s", columnName, message); + } +} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/aws/athena/PartitionProjection.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/projection/PartitionProjection.java similarity index 85% rename from plugin/trino-hive/src/main/java/io/trino/plugin/hive/aws/athena/PartitionProjection.java rename to plugin/trino-hive/src/main/java/io/trino/plugin/hive/projection/PartitionProjection.java index e20d8ed7f244..b81b4d05943f 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/aws/athena/PartitionProjection.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/projection/PartitionProjection.java @@ -11,11 +11,11 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package io.trino.plugin.hive.aws.athena; +package io.trino.plugin.hive.projection; import com.google.common.base.VerifyException; import com.google.common.collect.ImmutableMap; -import io.trino.plugin.hive.aws.athena.projection.Projection; +import io.trino.plugin.hive.metastore.Column; import io.trino.plugin.hive.metastore.Partition; import io.trino.plugin.hive.metastore.Table; import io.trino.spi.predicate.Domain; @@ -33,7 +33,7 @@ import static com.google.common.collect.ImmutableList.toImmutableList; import static com.google.common.collect.ImmutableSet.toImmutableSet; import static com.google.common.collect.Sets.cartesianProduct; -import static io.trino.plugin.hive.aws.athena.projection.Projection.invalidProjectionMessage; +import static io.trino.plugin.hive.projection.InvalidProjectionException.invalidProjectionMessage; import static io.trino.plugin.hive.util.HiveUtil.escapePathName; import static io.trino.plugin.hive.util.HiveUtil.toPartitionValues; import static java.lang.String.format; @@ -43,23 +43,15 @@ public final class PartitionProjection { private static final Pattern PROJECTION_LOCATION_TEMPLATE_PLACEHOLDER_PATTERN = Pattern.compile("(\\$\\{[^}]+\\})"); - private final boolean enabled; - private final Optional storageLocationTemplate; private final Map columnProjections; - public PartitionProjection(boolean enabled, Optional storageLocationTemplate, Map columnProjections) + public PartitionProjection(Optional storageLocationTemplate, Map columnProjections) { - this.enabled = enabled; this.storageLocationTemplate = requireNonNull(storageLocationTemplate, "storageLocationTemplate is null"); this.columnProjections = ImmutableMap.copyOf(requireNonNull(columnProjections, "columnProjections is null")); } - public boolean isEnabled() - { - return enabled; - } - public Optional> getProjectedPartitionNamesByFilter(List columnNames, TupleDomain partitionKeysFilter) { if (partitionKeysFilter.isNone()) { @@ -108,15 +100,23 @@ private Partition buildPartitionObject(Table table, String partitionName) .map(template -> expandStorageLocationTemplate( template, table.getPartitionColumns().stream() - .map(column -> column.getName()).collect(Collectors.toList()), + .map(Column::getName).collect(Collectors.toList()), partitionValues)) - .orElseGet(() -> format("%s/%s/", table.getStorage().getLocation(), partitionName))) + .orElseGet(() -> getPartitionLocation(table.getStorage().getLocation(), partitionName))) .setBucketProperty(table.getStorage().getBucketProperty()) .setSerdeParameters(table.getStorage().getSerdeParameters())) .build(); } - private String expandStorageLocationTemplate(String template, List partitionColumns, List partitionValues) + private static String getPartitionLocation(String tableLocation, String partitionName) + { + if (tableLocation.endsWith("/")) { + return format("%s%s/", tableLocation, partitionName); + } + return format("%s/%s/", tableLocation, partitionName); + } + + private static String expandStorageLocationTemplate(String template, List partitionColumns, List partitionValues) { Matcher matcher = PROJECTION_LOCATION_TEMPLATE_PLACEHOLDER_PATTERN.matcher(template); StringBuilder location = new StringBuilder(); diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/projection/PartitionProjectionProperties.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/projection/PartitionProjectionProperties.java new file mode 100644 index 000000000000..bc15420154ec --- /dev/null +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/projection/PartitionProjectionProperties.java @@ -0,0 +1,337 @@ +/* + * 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.hive.projection; + +import com.google.common.base.Joiner; +import com.google.common.base.Splitter; +import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableSet; +import io.trino.plugin.hive.metastore.Column; +import io.trino.plugin.hive.metastore.Table; +import io.trino.spi.connector.ColumnMetadata; +import io.trino.spi.connector.ConnectorTableMetadata; +import io.trino.spi.type.Type; +import io.trino.spi.type.TypeManager; + +import java.time.temporal.ChronoUnit; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.function.Function; +import java.util.stream.Collectors; + +import static com.google.common.collect.ImmutableList.toImmutableList; +import static com.google.common.collect.ImmutableMap.toImmutableMap; +import static io.trino.plugin.hive.HiveTableProperties.getPartitionedBy; +import static io.trino.plugin.hive.HiveTimestampPrecision.DEFAULT_PRECISION; +import static io.trino.plugin.hive.util.HiveTypeUtil.getType; +import static java.lang.Boolean.parseBoolean; +import static java.lang.String.format; +import static java.util.Locale.ROOT; + +public final class PartitionProjectionProperties +{ + private static final String COLUMN_PROJECTION_TYPE_SUFFIX = "type"; + private static final String COLUMN_PROJECTION_VALUES_SUFFIX = "values"; + private static final String COLUMN_PROJECTION_RANGE_SUFFIX = "range"; + private static final String COLUMN_PROJECTION_INTERVAL_SUFFIX = "interval"; + private static final String COLUMN_PROJECTION_DIGITS_SUFFIX = "digits"; + private static final String COLUMN_PROJECTION_FORMAT_SUFFIX = "format"; + private static final String METASTORE_PROPERTY_PROJECTION_INTERVAL_UNIT_SUFFIX = "interval.unit"; + private static final String METASTORE_PROPERTY_PROJECTION_ENABLED = "projection.enabled"; + private static final String METASTORE_PROPERTY_PROJECTION_LOCATION_TEMPLATE = "storage.location.template"; + private static final String METASTORE_PROPERTY_PROJECTION_IGNORE = "trino.partition_projection.ignore"; + private static final String PROPERTY_KEY_PREFIX = "partition_projection_"; + + public static final String COLUMN_PROJECTION_FORMAT = PROPERTY_KEY_PREFIX + COLUMN_PROJECTION_FORMAT_SUFFIX; + public static final String COLUMN_PROJECTION_DIGITS = PROPERTY_KEY_PREFIX + COLUMN_PROJECTION_DIGITS_SUFFIX; + public static final String COLUMN_PROJECTION_INTERVAL_UNIT = PROPERTY_KEY_PREFIX + "interval_unit"; + public static final String COLUMN_PROJECTION_INTERVAL = PROPERTY_KEY_PREFIX + COLUMN_PROJECTION_INTERVAL_SUFFIX; + public static final String COLUMN_PROJECTION_RANGE = PROPERTY_KEY_PREFIX + COLUMN_PROJECTION_RANGE_SUFFIX; + public static final String COLUMN_PROJECTION_VALUES = PROPERTY_KEY_PREFIX + COLUMN_PROJECTION_VALUES_SUFFIX; + public static final String COLUMN_PROJECTION_TYPE = PROPERTY_KEY_PREFIX + COLUMN_PROJECTION_TYPE_SUFFIX; + public static final String PARTITION_PROJECTION_IGNORE = PROPERTY_KEY_PREFIX + "ignore"; + public static final String PARTITION_PROJECTION_LOCATION_TEMPLATE = PROPERTY_KEY_PREFIX + "location_template"; + public static final String PARTITION_PROJECTION_ENABLED = PROPERTY_KEY_PREFIX + "enabled"; + + private PartitionProjectionProperties() {} + + public static Map getPartitionProjectionTrinoTableProperties(Table table) + { + Map metastoreTableProperties = table.getParameters(); + ImmutableMap.Builder trinoTablePropertiesBuilder = ImmutableMap.builder(); + + String ignore = metastoreTableProperties.get(METASTORE_PROPERTY_PROJECTION_IGNORE); + if (ignore != null) { + trinoTablePropertiesBuilder.put(PARTITION_PROJECTION_IGNORE, Boolean.valueOf(ignore)); + } + + String enabled = metastoreTableProperties.get(METASTORE_PROPERTY_PROJECTION_ENABLED); + if (enabled != null) { + trinoTablePropertiesBuilder.put(PARTITION_PROJECTION_ENABLED, Boolean.valueOf(enabled)); + } + + String locationTemplate = metastoreTableProperties.get(METASTORE_PROPERTY_PROJECTION_LOCATION_TEMPLATE); + if (locationTemplate != null) { + trinoTablePropertiesBuilder.put(PARTITION_PROJECTION_LOCATION_TEMPLATE, locationTemplate); + } + + return trinoTablePropertiesBuilder.buildOrThrow(); + } + + public static Map getPartitionProjectionTrinoColumnProperties(Table table, String columnName) + { + Map metastoreTableProperties = table.getParameters(); + return rewriteColumnProjectionProperties(metastoreTableProperties, columnName); + } + + public static Map getPartitionProjectionHiveTableProperties(ConnectorTableMetadata tableMetadata) + { + ImmutableMap.Builder metastoreTablePropertiesBuilder = ImmutableMap.builder(); + // Handle Table Properties + Map trinoTableProperties = tableMetadata.getProperties(); + + Object ignore = trinoTableProperties.get(PARTITION_PROJECTION_IGNORE); + if (ignore != null) { + metastoreTablePropertiesBuilder.put(METASTORE_PROPERTY_PROJECTION_IGNORE, ignore.toString().toLowerCase(ROOT)); + } + + Object enabled = trinoTableProperties.get(PARTITION_PROJECTION_ENABLED); + if (enabled != null) { + metastoreTablePropertiesBuilder.put(METASTORE_PROPERTY_PROJECTION_ENABLED, enabled.toString().toLowerCase(ROOT)); + } + + Object locationTemplate = trinoTableProperties.get(PARTITION_PROJECTION_LOCATION_TEMPLATE); + if (locationTemplate != null) { + metastoreTablePropertiesBuilder.put(METASTORE_PROPERTY_PROJECTION_LOCATION_TEMPLATE, locationTemplate.toString()); + } + + // Handle Column Properties + tableMetadata.getColumns().stream() + .filter(columnMetadata -> !columnMetadata.getProperties().isEmpty()) + .forEach(columnMetadata -> { + Map columnProperties = columnMetadata.getProperties(); + String columnName = columnMetadata.getName(); + + if (columnProperties.get(COLUMN_PROJECTION_TYPE) instanceof ProjectionType projectionType) { + metastoreTablePropertiesBuilder.put(getMetastoreProjectionPropertyKey(columnName, COLUMN_PROJECTION_TYPE_SUFFIX), projectionType.name().toLowerCase(ROOT)); + } + + if (columnProperties.get(COLUMN_PROJECTION_VALUES) instanceof List values) { + metastoreTablePropertiesBuilder.put(getMetastoreProjectionPropertyKey(columnName, COLUMN_PROJECTION_VALUES_SUFFIX), Joiner.on(",").join(values)); + } + + if (columnProperties.get(COLUMN_PROJECTION_RANGE) instanceof List range) { + metastoreTablePropertiesBuilder.put(getMetastoreProjectionPropertyKey(columnName, COLUMN_PROJECTION_RANGE_SUFFIX), Joiner.on(",").join(range)); + } + + if (columnProperties.get(COLUMN_PROJECTION_INTERVAL) instanceof Integer interval) { + metastoreTablePropertiesBuilder.put(getMetastoreProjectionPropertyKey(columnName, COLUMN_PROJECTION_INTERVAL_SUFFIX), interval.toString()); + } + + if (columnProperties.get(COLUMN_PROJECTION_INTERVAL_UNIT) instanceof ChronoUnit intervalUnit) { + metastoreTablePropertiesBuilder.put(getMetastoreProjectionPropertyKey(columnName, METASTORE_PROPERTY_PROJECTION_INTERVAL_UNIT_SUFFIX), intervalUnit.name().toLowerCase(ROOT)); + } + + if (columnProperties.get(COLUMN_PROJECTION_DIGITS) instanceof Integer digits) { + metastoreTablePropertiesBuilder.put(getMetastoreProjectionPropertyKey(columnName, COLUMN_PROJECTION_DIGITS_SUFFIX), digits.toString()); + } + + if (columnProperties.get(COLUMN_PROJECTION_FORMAT) instanceof String format) { + metastoreTablePropertiesBuilder.put(getMetastoreProjectionPropertyKey(columnName, COLUMN_PROJECTION_FORMAT_SUFFIX), format); + } + }); + + // We initialize partition projection to validate properties. + Map metastoreTableProperties = metastoreTablePropertiesBuilder.buildOrThrow(); + Set partitionColumnNames = ImmutableSet.copyOf(getPartitionedBy(tableMetadata.getProperties())); + createPartitionProjection( + tableMetadata.getColumns() + .stream() + .map(ColumnMetadata::getName) + .filter(name -> !partitionColumnNames.contains(name)) + .collect(toImmutableList()), + tableMetadata.getColumns().stream() + .filter(columnMetadata -> partitionColumnNames.contains(columnMetadata.getName())) + .collect(toImmutableMap(ColumnMetadata::getName, ColumnMetadata::getType)), + metastoreTableProperties); + + return metastoreTableProperties; + } + + public static boolean arePartitionProjectionPropertiesSet(ConnectorTableMetadata tableMetadata) + { + if (tableMetadata.getProperties().keySet().stream() + .anyMatch(propertyKey -> propertyKey.startsWith(PROPERTY_KEY_PREFIX))) { + return true; + } + return tableMetadata.getColumns().stream() + .map(columnMetadata -> columnMetadata.getProperties().keySet()) + .flatMap(Set::stream) + .anyMatch(propertyKey -> propertyKey.startsWith(PROPERTY_KEY_PREFIX)); + } + + public static Optional getPartitionProjectionFromTable(Table table, TypeManager typeManager) + { + Map tableProperties = table.getParameters(); + if (parseBoolean(tableProperties.get(METASTORE_PROPERTY_PROJECTION_IGNORE)) || + !parseBoolean(tableProperties.get(METASTORE_PROPERTY_PROJECTION_ENABLED))) { + return Optional.empty(); + } + + Set partitionColumnNames = table.getPartitionColumns().stream().map(Column::getName).collect(Collectors.toSet()); + return createPartitionProjection( + table.getDataColumns().stream() + .map(Column::getName) + .filter(partitionColumnNames::contains) + .collect(toImmutableList()), + table.getPartitionColumns().stream() + .collect(toImmutableMap(Column::getName, column -> getType(column.getType(), typeManager, DEFAULT_PRECISION))), + tableProperties); + } + + private static Optional createPartitionProjection(List dataColumns, Map partitionColumns, Map tableProperties) + { + // This method is used during table creation to validate the properties. The validation is performed even if the projection is disabled. + boolean enabled = parseBoolean(tableProperties.get(METASTORE_PROPERTY_PROJECTION_ENABLED)); + + if (!tableProperties.containsKey(METASTORE_PROPERTY_PROJECTION_ENABLED) && + partitionColumns.keySet().stream().anyMatch(columnName -> !rewriteColumnProjectionProperties(tableProperties, columnName).isEmpty())) { + throw new InvalidProjectionException("Columns partition projection properties cannot be set when '%s' is not set".formatted(PARTITION_PROJECTION_ENABLED)); + } + + if (enabled && partitionColumns.isEmpty()) { + throw new InvalidProjectionException("Partition projection cannot be enabled on a table that is not partitioned"); + } + + for (String columnName : dataColumns) { + if (!rewriteColumnProjectionProperties(tableProperties, columnName).isEmpty()) { + throw new InvalidProjectionException("Partition projection cannot be defined for non-partition column: '" + columnName + "'"); + } + } + + Map columnProjections = new HashMap<>(); + partitionColumns.forEach((columnName, type) -> { + Map columnProperties = rewriteColumnProjectionProperties(tableProperties, columnName); + if (enabled) { + columnProjections.put(columnName, parseColumnProjection(columnName, type, columnProperties)); + } + }); + + Optional storageLocationTemplate = Optional.ofNullable(tableProperties.get(METASTORE_PROPERTY_PROJECTION_LOCATION_TEMPLATE)); + storageLocationTemplate.ifPresent(locationTemplate -> { + for (String columnName : partitionColumns.keySet()) { + if (!locationTemplate.contains("${" + columnName + "}")) { + throw new InvalidProjectionException(format("Partition projection location template: %s is missing partition column: '%s' placeholder", locationTemplate, columnName)); + } + } + }); + + if (!enabled) { + return Optional.empty(); + } + return Optional.of(new PartitionProjection(storageLocationTemplate, columnProjections)); + } + + private static Map rewriteColumnProjectionProperties(Map metastoreTableProperties, String columnName) + { + ImmutableMap.Builder trinoTablePropertiesBuilder = ImmutableMap.builder(); + + String type = metastoreTableProperties.get(getMetastoreProjectionPropertyKey(columnName, COLUMN_PROJECTION_TYPE_SUFFIX)); + if (type != null) { + trinoTablePropertiesBuilder.put(COLUMN_PROJECTION_TYPE, ProjectionType.valueOf(type.toUpperCase(ROOT))); + } + + String values = metastoreTableProperties.get(getMetastoreProjectionPropertyKey(columnName, COLUMN_PROJECTION_VALUES_SUFFIX)); + if (values != null) { + trinoTablePropertiesBuilder.put(COLUMN_PROJECTION_VALUES, splitCommaSeparatedString(values)); + } + + String range = metastoreTableProperties.get(getMetastoreProjectionPropertyKey(columnName, COLUMN_PROJECTION_RANGE_SUFFIX)); + if (range != null) { + trinoTablePropertiesBuilder.put(COLUMN_PROJECTION_RANGE, splitCommaSeparatedString(range)); + } + + String interval = metastoreTableProperties.get(getMetastoreProjectionPropertyKey(columnName, COLUMN_PROJECTION_INTERVAL_SUFFIX)); + if (interval != null) { + trinoTablePropertiesBuilder.put(COLUMN_PROJECTION_INTERVAL, Integer.valueOf(interval)); + } + + String intervalUnit = metastoreTableProperties.get(getMetastoreProjectionPropertyKey(columnName, METASTORE_PROPERTY_PROJECTION_INTERVAL_UNIT_SUFFIX)); + if (intervalUnit != null) { + trinoTablePropertiesBuilder.put(COLUMN_PROJECTION_INTERVAL_UNIT, ChronoUnit.valueOf(intervalUnit.toUpperCase(ROOT))); + } + + String digits = metastoreTableProperties.get(getMetastoreProjectionPropertyKey(columnName, COLUMN_PROJECTION_DIGITS_SUFFIX)); + if (digits != null) { + trinoTablePropertiesBuilder.put(COLUMN_PROJECTION_DIGITS, Integer.valueOf(digits)); + } + + String format = metastoreTableProperties.get(getMetastoreProjectionPropertyKey(columnName, COLUMN_PROJECTION_FORMAT_SUFFIX)); + if (format != null) { + trinoTablePropertiesBuilder.put(COLUMN_PROJECTION_FORMAT, format); + } + + return trinoTablePropertiesBuilder.buildOrThrow(); + } + + private static Projection parseColumnProjection(String columnName, Type columnType, Map columnProperties) + { + ProjectionType projectionType = (ProjectionType) columnProperties.get(COLUMN_PROJECTION_TYPE); + if (projectionType == null) { + throw new InvalidProjectionException(columnName, "Projection type property missing"); + } + return switch (projectionType) { + case ENUM -> new EnumProjection(columnName, columnType, columnProperties); + case INTEGER -> new IntegerProjection(columnName, columnType, columnProperties); + case DATE -> new DateProjection(columnName, columnType, columnProperties); + case INJECTED -> new InjectedProjection(columnName, columnType); + }; + } + + private static List splitCommaSeparatedString(String value) + { + return Splitter.on(',') + .trimResults() + .omitEmptyStrings() + .splitToList(value); + } + + private static String getMetastoreProjectionPropertyKey(String columnName, String propertyKeySuffix) + { + return "projection" + "." + columnName + "." + propertyKeySuffix; + } + + static T getProjectionPropertyRequiredValue( + String columnName, + Map columnProjectionProperties, + String propertyKey, + Function decoder) + { + return getProjectionPropertyValue(columnProjectionProperties, propertyKey, decoder) + .orElseThrow(() -> new InvalidProjectionException(columnName, format("Missing required property: '%s'", propertyKey))); + } + + static Optional getProjectionPropertyValue( + Map columnProjectionProperties, + String propertyKey, + Function decoder) + { + return Optional.ofNullable( + columnProjectionProperties.get(propertyKey)) + .map(decoder); + } +} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/projection/Projection.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/projection/Projection.java new file mode 100644 index 000000000000..0677769dfaba --- /dev/null +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/projection/Projection.java @@ -0,0 +1,25 @@ +/* + * 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.hive.projection; + +import io.trino.spi.predicate.Domain; + +import java.util.List; +import java.util.Optional; + +sealed interface Projection + permits DateProjection, EnumProjection, InjectedProjection, IntegerProjection +{ + List getProjectedValues(Optional partitionValueFilter); +} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/aws/athena/projection/ProjectionType.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/projection/ProjectionType.java similarity index 85% rename from plugin/trino-hive/src/main/java/io/trino/plugin/hive/aws/athena/projection/ProjectionType.java rename to plugin/trino-hive/src/main/java/io/trino/plugin/hive/projection/ProjectionType.java index f5eca8363cc0..7521053f87d2 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/aws/athena/projection/ProjectionType.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/projection/ProjectionType.java @@ -11,12 +11,9 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package io.trino.plugin.hive.aws.athena.projection; +package io.trino.plugin.hive.projection; public enum ProjectionType { - ENUM, - INTEGER, - DATE, - INJECTED; + ENUM, INTEGER, DATE, INJECTED } diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/rcfile/RcFilePageSourceFactory.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/rcfile/RcFilePageSourceFactory.java index 2be057f69577..c3d5f7955651 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/rcfile/RcFilePageSourceFactory.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/rcfile/RcFilePageSourceFactory.java @@ -14,7 +14,6 @@ package io.trino.plugin.hive.rcfile; import com.google.common.collect.ImmutableMap; -import com.google.common.collect.Maps; import com.google.inject.Inject; import io.airlift.slice.Slices; import io.airlift.units.DataSize; @@ -31,14 +30,13 @@ import io.trino.hive.formats.encodings.text.TextEncodingOptions; import io.trino.hive.formats.rcfile.RcFileReader; import io.trino.plugin.hive.AcidInfo; -import io.trino.plugin.hive.FileFormatDataSourceStats; import io.trino.plugin.hive.HiveColumnHandle; import io.trino.plugin.hive.HiveConfig; import io.trino.plugin.hive.HivePageSourceFactory; import io.trino.plugin.hive.ReaderColumns; import io.trino.plugin.hive.ReaderPageSource; +import io.trino.plugin.hive.Schema; import io.trino.plugin.hive.acid.AcidTransaction; -import io.trino.plugin.hive.fs.MonitoredInputFile; import io.trino.spi.TrinoException; import io.trino.spi.connector.ConnectorPageSource; import io.trino.spi.connector.ConnectorSession; @@ -51,7 +49,6 @@ import java.util.List; import java.util.Optional; import java.util.OptionalInt; -import java.util.Properties; import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.collect.ImmutableList.toImmutableList; @@ -61,9 +58,7 @@ import static io.trino.plugin.hive.ReaderPageSource.noProjectionAdaptation; import static io.trino.plugin.hive.util.HiveClassNames.COLUMNAR_SERDE_CLASS; import static io.trino.plugin.hive.util.HiveClassNames.LAZY_BINARY_COLUMNAR_SERDE_CLASS; -import static io.trino.plugin.hive.util.HiveUtil.getDeserializerClassName; import static io.trino.plugin.hive.util.HiveUtil.splitError; -import static io.trino.plugin.hive.util.SerdeConstants.SERIALIZATION_LIB; import static java.lang.Math.min; import static java.util.Objects.requireNonNull; @@ -73,25 +68,18 @@ public class RcFilePageSourceFactory private static final DataSize BUFFER_SIZE = DataSize.of(8, Unit.MEGABYTE); private final TrinoFileSystemFactory fileSystemFactory; - private final FileFormatDataSourceStats stats; private final DateTimeZone timeZone; @Inject - public RcFilePageSourceFactory(TrinoFileSystemFactory fileSystemFactory, FileFormatDataSourceStats stats, HiveConfig hiveConfig) + public RcFilePageSourceFactory(TrinoFileSystemFactory fileSystemFactory, HiveConfig hiveConfig) { this.fileSystemFactory = requireNonNull(fileSystemFactory, "fileSystemFactory is null"); - this.stats = requireNonNull(stats, "stats is null"); this.timeZone = hiveConfig.getRcfileDateTimeZone(); } - public static Properties stripUnnecessaryProperties(Properties schema) + public static boolean stripUnnecessaryProperties(String serializationLibraryName) { - if (LAZY_BINARY_COLUMNAR_SERDE_CLASS.equals(getDeserializerClassName(schema))) { - Properties stripped = new Properties(); - stripped.put(SERIALIZATION_LIB, schema.getProperty(SERIALIZATION_LIB)); - return stripped; - } - return schema; + return LAZY_BINARY_COLUMNAR_SERDE_CLASS.equals(serializationLibraryName); } @Override @@ -101,7 +89,8 @@ public Optional createPageSource( long start, long length, long estimatedFileSize, - Properties schema, + long fileModifiedTime, + Schema schema, List columns, TupleDomain effectivePredicate, Optional acidInfo, @@ -110,12 +99,12 @@ public Optional createPageSource( AcidTransaction transaction) { ColumnEncodingFactory columnEncodingFactory; - String deserializerClassName = getDeserializerClassName(schema); + String deserializerClassName = schema.serializationLibraryName(); if (deserializerClassName.equals(LAZY_BINARY_COLUMNAR_SERDE_CLASS)) { columnEncodingFactory = new BinaryColumnEncodingFactory(timeZone); } else if (deserializerClassName.equals(COLUMNAR_SERDE_CLASS)) { - columnEncodingFactory = new TextColumnEncodingFactory(TextEncodingOptions.fromSchema(Maps.fromProperties(schema))); + columnEncodingFactory = new TextColumnEncodingFactory(TextEncodingOptions.fromSchema(schema.serdeProperties())); } else { return Optional.empty(); @@ -133,7 +122,7 @@ else if (deserializerClassName.equals(COLUMNAR_SERDE_CLASS)) { } TrinoFileSystem trinoFileSystem = fileSystemFactory.create(session.getIdentity()); - TrinoInputFile inputFile = new MonitoredInputFile(stats, trinoFileSystem.newInputFile(path)); + TrinoInputFile inputFile = trinoFileSystem.newInputFile(path); try { length = min(inputFile.length() - start, length); if (!inputFile.exists()) { diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/s3select/IonSqlQueryBuilder.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/s3select/IonSqlQueryBuilder.java deleted file mode 100644 index 53a1ca72288b..000000000000 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/s3select/IonSqlQueryBuilder.java +++ /dev/null @@ -1,274 +0,0 @@ -/* - * 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.hive.s3select; - -import com.google.common.base.Joiner; -import com.google.common.collect.ImmutableList; -import com.google.common.primitives.Shorts; -import com.google.common.primitives.SignedBytes; -import io.airlift.slice.Slice; -import io.trino.plugin.hive.HiveColumnHandle; -import io.trino.spi.predicate.Domain; -import io.trino.spi.predicate.Range; -import io.trino.spi.predicate.TupleDomain; -import io.trino.spi.type.Type; -import io.trino.spi.type.TypeManager; -import io.trino.spi.type.VarcharType; -import org.joda.time.format.DateTimeFormatter; - -import java.util.ArrayList; -import java.util.List; -import java.util.Optional; - -import static com.google.common.base.Preconditions.checkArgument; -import static com.google.common.base.Preconditions.checkState; -import static com.google.common.collect.Iterables.getOnlyElement; -import static io.trino.plugin.hive.s3select.S3SelectDataType.CSV; -import static io.trino.spi.type.BigintType.BIGINT; -import static io.trino.spi.type.BooleanType.BOOLEAN; -import static io.trino.spi.type.DateType.DATE; -import static io.trino.spi.type.IntegerType.INTEGER; -import static io.trino.spi.type.SmallintType.SMALLINT; -import static io.trino.spi.type.TinyintType.TINYINT; -import static java.lang.Math.toIntExact; -import static java.lang.String.format; -import static java.util.Objects.requireNonNull; -import static java.util.concurrent.TimeUnit.DAYS; -import static java.util.stream.Collectors.joining; -import static org.joda.time.chrono.ISOChronology.getInstanceUTC; -import static org.joda.time.format.ISODateTimeFormat.date; - -/** - * S3 Select uses Ion SQL++ query language. This class is used to construct a valid Ion SQL++ query - * to be evaluated with S3 Select on an S3 object. - */ -public class IonSqlQueryBuilder -{ - private static final DateTimeFormatter FORMATTER = date().withChronology(getInstanceUTC()); - private static final String DATA_SOURCE = "S3Object s"; - private final TypeManager typeManager; - private final S3SelectDataType s3SelectDataType; - private final String nullPredicate; - private final String notNullPredicate; - - public IonSqlQueryBuilder(TypeManager typeManager, S3SelectDataType s3SelectDataType, Optional optionalNullCharacterEncoding) - { - if (optionalNullCharacterEncoding.isPresent()) { - checkArgument(s3SelectDataType == CSV, "Null character encoding should only be provided for CSV data"); - } - - this.typeManager = requireNonNull(typeManager, "typeManager is null"); - this.s3SelectDataType = requireNonNull(s3SelectDataType, "s3SelectDataType is null"); - - String nullCharacterEncoding = optionalNullCharacterEncoding.orElse(""); - this.nullPredicate = switch (s3SelectDataType) { - case JSON -> "IS NULL"; - case CSV -> "= '%s'".formatted(nullCharacterEncoding); - }; - this.notNullPredicate = switch (s3SelectDataType) { - case JSON -> "IS NOT NULL"; - case CSV -> "!= '%s'".formatted(nullCharacterEncoding); - }; - } - - public String buildSql(List columns, TupleDomain tupleDomain) - { - columns.forEach(column -> checkArgument(column.isBaseColumn(), "%s is not a base column", column)); - tupleDomain.getDomains().ifPresent(domains -> { - domains.keySet().forEach(column -> checkArgument(column.isBaseColumn(), "%s is not a base column", column)); - }); - - // SELECT clause - StringBuilder sql = new StringBuilder("SELECT "); - - if (columns.isEmpty()) { - sql.append("' '"); - } - else { - String columnNames = columns.stream() - .map(this::getFullyQualifiedColumnName) - .collect(joining(", ")); - sql.append(columnNames); - } - - // FROM clause - sql.append(" FROM "); - sql.append(DATA_SOURCE); - - // WHERE clause - List clauses = toConjuncts(columns, tupleDomain); - if (!clauses.isEmpty()) { - sql.append(" WHERE ") - .append(Joiner.on(" AND ").join(clauses)); - } - - return sql.toString(); - } - - private String getFullyQualifiedColumnName(HiveColumnHandle column) - { - return switch (s3SelectDataType) { - case JSON -> "s.%s".formatted(column.getBaseColumnName()); - case CSV -> "s._%d".formatted(column.getBaseHiveColumnIndex() + 1); - }; - } - - private List toConjuncts(List columns, TupleDomain tupleDomain) - { - ImmutableList.Builder builder = ImmutableList.builder(); - for (HiveColumnHandle column : columns) { - Type type = column.getHiveType().getType(typeManager); - if (tupleDomain.getDomains().isPresent() && isSupported(type)) { - Domain domain = tupleDomain.getDomains().get().get(column); - if (domain != null) { - builder.add(toPredicate(domain, type, column)); - } - } - } - return builder.build(); - } - - private static boolean isSupported(Type type) - { - Type validType = requireNonNull(type, "type is null"); - return validType.equals(BIGINT) || - validType.equals(TINYINT) || - validType.equals(SMALLINT) || - validType.equals(INTEGER) || - validType.equals(BOOLEAN) || - validType.equals(DATE) || - validType instanceof VarcharType; - } - - private String toPredicate(Domain domain, Type type, HiveColumnHandle column) - { - checkArgument(domain.getType().isOrderable(), "Domain type must be orderable"); - - if (domain.getValues().isNone()) { - if (domain.isNullAllowed()) { - return getFullyQualifiedColumnName(column) + " " + nullPredicate; - } - return "FALSE"; - } - - if (domain.getValues().isAll()) { - if (domain.isNullAllowed()) { - return "TRUE"; - } - return getFullyQualifiedColumnName(column) + " " + notNullPredicate; - } - - List disjuncts = new ArrayList<>(); - List singleValues = new ArrayList<>(); - for (Range range : domain.getValues().getRanges().getOrderedRanges()) { - checkState(!range.isAll()); - if (range.isSingleValue()) { - singleValues.add(range.getSingleValue()); - continue; - } - List rangeConjuncts = new ArrayList<>(); - if (!range.isLowUnbounded()) { - rangeConjuncts.add(toPredicate(range.isLowInclusive() ? ">=" : ">", range.getLowBoundedValue(), type, column)); - } - if (!range.isHighUnbounded()) { - rangeConjuncts.add(toPredicate(range.isHighInclusive() ? "<=" : "<", range.getHighBoundedValue(), type, column)); - } - // If rangeConjuncts is null, then the range was ALL, which should already have been checked for - checkState(!rangeConjuncts.isEmpty()); - if (rangeConjuncts.size() == 1) { - disjuncts.add("%s %s AND %s".formatted(getFullyQualifiedColumnName(column), notNullPredicate, getOnlyElement(rangeConjuncts))); - } - else { - disjuncts.add("(%s %s AND %s)".formatted(getFullyQualifiedColumnName(column), notNullPredicate, Joiner.on(" AND ").join(rangeConjuncts))); - } - } - - // Add back all of the possible single values either as an equality or an IN predicate - if (singleValues.size() == 1) { - disjuncts.add("%s %s AND %s".formatted(getFullyQualifiedColumnName(column), notNullPredicate, toPredicate("=", getOnlyElement(singleValues), type, column))); - } - else if (singleValues.size() > 1) { - List values = new ArrayList<>(); - for (Object value : singleValues) { - checkType(type); - values.add(valueToQuery(type, value)); - } - disjuncts.add("%s %s AND %s IN (%s)".formatted( - getFullyQualifiedColumnName(column), - notNullPredicate, - createColumn(type, column), - Joiner.on(",").join(values))); - } - - // Add nullability disjuncts - checkState(!disjuncts.isEmpty()); - if (domain.isNullAllowed()) { - disjuncts.add(getFullyQualifiedColumnName(column) + " " + nullPredicate); - } - - return "(" + Joiner.on(" OR ").join(disjuncts) + ")"; - } - - private String toPredicate(String operator, Object value, Type type, HiveColumnHandle column) - { - checkType(type); - - return format("%s %s %s", createColumn(type, column), operator, valueToQuery(type, value)); - } - - private static void checkType(Type type) - { - checkArgument(isSupported(type), "Type not supported: %s", type); - } - - private static String valueToQuery(Type type, Object value) - { - if (type.equals(BIGINT)) { - return String.valueOf((long) value); - } - if (type.equals(INTEGER)) { - return String.valueOf(toIntExact((long) value)); - } - if (type.equals(SMALLINT)) { - return String.valueOf(Shorts.checkedCast((long) value)); - } - if (type.equals(TINYINT)) { - return String.valueOf(SignedBytes.checkedCast((long) value)); - } - if (type.equals(BOOLEAN)) { - return String.valueOf((boolean) value); - } - if (type.equals(DATE)) { - // CAST('2007-04-05T14:30Z' AS TIMESTAMP) - return "'" + FORMATTER.print(DAYS.toMillis((long) value)) + "'"; - } - if (type.equals(VarcharType.VARCHAR)) { - return "'" + ((Slice) value).toStringUtf8().replace("'", "''") + "'"; - } - return "'" + ((Slice) value).toStringUtf8() + "'"; - } - - private String createColumn(Type type, HiveColumnHandle columnHandle) - { - String column = getFullyQualifiedColumnName(columnHandle); - - if (type.equals(BIGINT) || type.equals(INTEGER) || type.equals(SMALLINT) || type.equals(TINYINT)) { - return "CAST(" + column + " AS INT)"; - } - if (type.equals(BOOLEAN)) { - return "CAST(" + column + " AS BOOL)"; - } - return column; - } -} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/s3select/S3SelectLineRecordReader.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/s3select/S3SelectLineRecordReader.java deleted file mode 100644 index ec74e14a7759..000000000000 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/s3select/S3SelectLineRecordReader.java +++ /dev/null @@ -1,308 +0,0 @@ -/* - * 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.hive.s3select; - -import com.amazonaws.AbortedException; -import com.amazonaws.services.s3.model.AmazonS3Exception; -import com.amazonaws.services.s3.model.CompressionType; -import com.amazonaws.services.s3.model.InputSerialization; -import com.amazonaws.services.s3.model.OutputSerialization; -import com.amazonaws.services.s3.model.ScanRange; -import com.amazonaws.services.s3.model.SelectObjectContentRequest; -import com.google.common.annotations.VisibleForTesting; -import com.google.common.io.Closer; -import com.google.errorprone.annotations.ThreadSafe; -import io.airlift.units.Duration; -import io.trino.hdfs.s3.HiveS3Config; -import io.trino.hdfs.s3.TrinoS3FileSystem; -import io.trino.spi.TrinoException; -import org.apache.hadoop.conf.Configuration; -import org.apache.hadoop.fs.Path; -import org.apache.hadoop.io.LongWritable; -import org.apache.hadoop.io.Text; -import org.apache.hadoop.io.compress.BZip2Codec; -import org.apache.hadoop.io.compress.CompressionCodec; -import org.apache.hadoop.io.compress.CompressionCodecFactory; -import org.apache.hadoop.io.compress.GzipCodec; -import org.apache.hadoop.mapred.RecordReader; -import org.apache.hadoop.util.LineReader; - -import java.io.IOException; -import java.io.InputStream; -import java.io.InterruptedIOException; -import java.net.URI; -import java.nio.charset.StandardCharsets; -import java.util.Properties; - -import static com.amazonaws.services.s3.model.ExpressionType.SQL; -import static com.google.common.base.Throwables.throwIfInstanceOf; -import static com.google.common.base.Throwables.throwIfUnchecked; -import static io.trino.hdfs.s3.TrinoS3FileSystem.S3_MAX_BACKOFF_TIME; -import static io.trino.hdfs.s3.TrinoS3FileSystem.S3_MAX_CLIENT_RETRIES; -import static io.trino.hdfs.s3.TrinoS3FileSystem.S3_MAX_RETRY_TIME; -import static io.trino.plugin.hive.util.RetryDriver.retry; -import static io.trino.plugin.hive.util.SerdeConstants.LINE_DELIM; -import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED; -import static java.lang.String.format; -import static java.net.HttpURLConnection.HTTP_BAD_REQUEST; -import static java.net.HttpURLConnection.HTTP_FORBIDDEN; -import static java.net.HttpURLConnection.HTTP_NOT_FOUND; -import static java.util.Objects.requireNonNull; -import static java.util.concurrent.TimeUnit.SECONDS; - -@ThreadSafe -public abstract class S3SelectLineRecordReader - implements RecordReader -{ - private InputStream selectObjectContent; - private long processedRecords; - private long recordsFromS3; - private long position; - private LineReader reader; - private boolean isFirstLine; - private static final Duration BACKOFF_MIN_SLEEP = new Duration(1, SECONDS); - private final TrinoS3SelectClient selectClient; - private final long start; - private final long end; - private final int maxAttempts; - private final Duration maxBackoffTime; - private final Duration maxRetryTime; - private final Closer closer = Closer.create(); - private final SelectObjectContentRequest selectObjectContentRequest; - private final CompressionCodecFactory compressionCodecFactory; - private final String lineDelimiter; - private final Properties schema; - private final CompressionType compressionType; - - public S3SelectLineRecordReader( - Configuration configuration, - Path path, - long start, - long length, - Properties schema, - String ionSqlQuery, - TrinoS3ClientFactory s3ClientFactory) - { - requireNonNull(configuration, "configuration is null"); - requireNonNull(schema, "schema is null"); - requireNonNull(path, "path is null"); - requireNonNull(ionSqlQuery, "ionSqlQuery is null"); - requireNonNull(s3ClientFactory, "s3ClientFactory is null"); - this.lineDelimiter = (schema).getProperty(LINE_DELIM, "\n"); - this.processedRecords = 0; - this.recordsFromS3 = 0; - this.start = start; - this.position = this.start; - this.end = this.start + length; - this.isFirstLine = true; - - this.compressionCodecFactory = new CompressionCodecFactory(configuration); - this.compressionType = getCompressionType(path); - this.schema = schema; - this.selectObjectContentRequest = buildSelectObjectRequest(ionSqlQuery, path); - - HiveS3Config defaults = new HiveS3Config(); - this.maxAttempts = configuration.getInt(S3_MAX_CLIENT_RETRIES, defaults.getS3MaxClientRetries()) + 1; - this.maxBackoffTime = Duration.valueOf(configuration.get(S3_MAX_BACKOFF_TIME, defaults.getS3MaxBackoffTime().toString())); - this.maxRetryTime = Duration.valueOf(configuration.get(S3_MAX_RETRY_TIME, defaults.getS3MaxRetryTime().toString())); - - this.selectClient = new TrinoS3SelectClient(configuration, s3ClientFactory); - closer.register(selectClient); - } - - protected abstract InputSerialization buildInputSerialization(); - - protected abstract OutputSerialization buildOutputSerialization(); - - protected abstract boolean shouldEnableScanRange(); - - protected Properties getSchema() - { - return schema; - } - - protected CompressionType getCompressionType() - { - return compressionType; - } - - public SelectObjectContentRequest buildSelectObjectRequest(String query, Path path) - { - SelectObjectContentRequest selectObjectRequest = new SelectObjectContentRequest(); - URI uri = path.toUri(); - selectObjectRequest.setBucketName(TrinoS3FileSystem.extractBucketName(uri)); - selectObjectRequest.setKey(TrinoS3FileSystem.keyFromPath(path)); - selectObjectRequest.setExpression(query); - selectObjectRequest.setExpressionType(SQL); - - InputSerialization selectObjectInputSerialization = buildInputSerialization(); - selectObjectRequest.setInputSerialization(selectObjectInputSerialization); - - OutputSerialization selectObjectOutputSerialization = buildOutputSerialization(); - selectObjectRequest.setOutputSerialization(selectObjectOutputSerialization); - - if (shouldEnableScanRange()) { - ScanRange scanRange = new ScanRange(); - scanRange.setStart(getStart()); - scanRange.setEnd(getEnd()); - selectObjectRequest.setScanRange(scanRange); - } - - return selectObjectRequest; - } - - protected CompressionType getCompressionType(Path path) - { - CompressionCodec codec = compressionCodecFactory.getCodec(path); - if (codec == null) { - return CompressionType.NONE; - } - if (codec instanceof GzipCodec) { - return CompressionType.GZIP; - } - if (codec instanceof BZip2Codec) { - return CompressionType.BZIP2; - } - throw new TrinoException(NOT_SUPPORTED, "Compression extension not supported for S3 Select: " + path); - } - - private int readLine(Text value) - throws IOException - { - try { - return retry() - .maxAttempts(maxAttempts) - .exponentialBackoff(BACKOFF_MIN_SLEEP, maxBackoffTime, maxRetryTime, 2.0) - .stopOn(InterruptedException.class, UnrecoverableS3OperationException.class, AbortedException.class) - .run("readRecordsContentStream", () -> { - if (isFirstLine) { - recordsFromS3 = 0; - selectObjectContent = selectClient.getRecordsContent(selectObjectContentRequest); - closer.register(selectObjectContent); - reader = new LineReader(selectObjectContent, lineDelimiter.getBytes(StandardCharsets.UTF_8)); - closer.register(reader); - isFirstLine = false; - } - try { - return reader.readLine(value); - } - catch (RuntimeException e) { - isFirstLine = true; - recordsFromS3 = 0; - if (e instanceof AmazonS3Exception) { - switch (((AmazonS3Exception) e).getStatusCode()) { - case HTTP_FORBIDDEN: - case HTTP_NOT_FOUND: - case HTTP_BAD_REQUEST: - throw new UnrecoverableS3OperationException(selectClient.getBucketName(), selectClient.getKeyName(), e); - } - } - throw e; - } - }); - } - catch (InterruptedException | AbortedException e) { - Thread.currentThread().interrupt(); - throw new InterruptedIOException(); - } - catch (Exception e) { - throwIfInstanceOf(e, IOException.class); - throwIfUnchecked(e); - throw new RuntimeException(e); - } - } - - @Override - public synchronized boolean next(LongWritable key, Text value) - throws IOException - { - while (true) { - int bytes = readLine(value); - if (bytes <= 0) { - if (!selectClient.isRequestComplete()) { - throw new IOException("S3 Select request was incomplete as End Event was not received"); - } - return false; - } - recordsFromS3++; - if (recordsFromS3 > processedRecords) { - position += bytes; - processedRecords++; - key.set(processedRecords); - return true; - } - } - } - - @Override - public LongWritable createKey() - { - return new LongWritable(); - } - - @Override - public Text createValue() - { - return new Text(); - } - - @Override - public long getPos() - { - return position; - } - - @Override - public void close() - throws IOException - { - closer.close(); - } - - @Override - public float getProgress() - { - return ((float) (position - start)) / (end - start); - } - - /** - * This exception is for stopping retries for S3 Select calls that shouldn't be retried. - * For example, "Caused by: com.amazonaws.services.s3.model.AmazonS3Exception: Forbidden (Service: Amazon S3; Status Code: 403 ..." - */ - @VisibleForTesting - static class UnrecoverableS3OperationException - extends RuntimeException - { - public UnrecoverableS3OperationException(String bucket, String key, Throwable cause) - { - // append bucket and key to the message - super(format("%s (Bucket: %s, Key: %s)", cause, bucket, key)); - } - } - - protected long getStart() - { - return start; - } - - protected long getEnd() - { - return end; - } - - protected String getLineDelimiter() - { - return lineDelimiter; - } -} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/s3select/S3SelectLineRecordReaderProvider.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/s3select/S3SelectLineRecordReaderProvider.java deleted file mode 100644 index 49221c398280..000000000000 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/s3select/S3SelectLineRecordReaderProvider.java +++ /dev/null @@ -1,51 +0,0 @@ -/* - * 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.hive.s3select; - -import io.trino.plugin.hive.s3select.csv.S3SelectCsvRecordReader; -import io.trino.plugin.hive.s3select.json.S3SelectJsonRecordReader; -import org.apache.hadoop.conf.Configuration; -import org.apache.hadoop.fs.Path; - -import java.util.Optional; -import java.util.Properties; - -/** - * Returns an S3SelectLineRecordReader based on the serDe class. It supports CSV and JSON formats, and - * will not push down any other formats. - */ -public class S3SelectLineRecordReaderProvider -{ - private S3SelectLineRecordReaderProvider() {} - - public static Optional get(Configuration configuration, - Path path, - long start, - long length, - Properties schema, - String ionSqlQuery, - TrinoS3ClientFactory s3ClientFactory, - S3SelectDataType dataType) - { - switch (dataType) { - case CSV: - return Optional.of(new S3SelectCsvRecordReader(configuration, path, start, length, schema, ionSqlQuery, s3ClientFactory)); - case JSON: - return Optional.of(new S3SelectJsonRecordReader(configuration, path, start, length, schema, ionSqlQuery, s3ClientFactory)); - default: - // return empty if data type is not returned by the serDeMapper or unrecognizable by the LineRecordReader - return Optional.empty(); - } - } -} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/s3select/S3SelectPushdown.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/s3select/S3SelectPushdown.java deleted file mode 100644 index b98e3f7c3bd8..000000000000 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/s3select/S3SelectPushdown.java +++ /dev/null @@ -1,165 +0,0 @@ -/* - * 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.hive.s3select; - -import com.google.common.collect.ImmutableSet; -import io.trino.hive.formats.compression.CompressionKind; -import io.trino.plugin.hive.metastore.Column; -import io.trino.plugin.hive.metastore.Partition; -import io.trino.plugin.hive.metastore.Table; -import io.trino.spi.connector.ConnectorSession; - -import java.util.List; -import java.util.Objects; -import java.util.Optional; -import java.util.Properties; -import java.util.Set; - -import static io.trino.plugin.hive.HiveMetadata.SKIP_FOOTER_COUNT_KEY; -import static io.trino.plugin.hive.HiveMetadata.SKIP_HEADER_COUNT_KEY; -import static io.trino.plugin.hive.HiveSessionProperties.isS3SelectPushdownEnabled; -import static io.trino.plugin.hive.metastore.MetastoreUtil.getHiveSchema; -import static io.trino.plugin.hive.util.HiveClassNames.TEXT_INPUT_FORMAT_CLASS; -import static io.trino.plugin.hive.util.HiveUtil.getDeserializerClassName; -import static io.trino.plugin.hive.util.HiveUtil.getInputFormatName; -import static java.util.Objects.requireNonNull; - -/** - * S3SelectPushdown uses Amazon S3 Select to push down queries to Amazon S3. This allows Presto to retrieve only a - * subset of data rather than retrieving the full S3 object thus improving Presto query performance. - */ -public final class S3SelectPushdown -{ - private static final Set SUPPORTED_S3_PREFIXES = ImmutableSet.of("s3://", "s3a://", "s3n://"); - - /* - * Double and Real Types lose precision. Thus, they are not pushed down to S3. - * Correctness problems have also been observed with Decimal columns. - * - * When S3 select support was added, Trino did not properly implement TIMESTAMP semantic. This was fixed in 2020, and TIMESTAMPS may be supportable now - * (https://github.com/trinodb/trino/issues/10962). Pushing down timestamps to s3select maybe still be problematic due to ION SQL comparing timestamps - * using precision. This means timestamps with different precisions are not equal even actually they present the same instant of time. - */ - private static final Set SUPPORTED_COLUMN_TYPES = ImmutableSet.of( - "boolean", - "int", - "tinyint", - "smallint", - "bigint", - "string", - "date"); - - private S3SelectPushdown() {} - - private static boolean isSerDeSupported(Properties schema) - { - String serdeName = getDeserializerClassName(schema); - return S3SelectSerDeDataTypeMapper.doesSerDeExist(serdeName); - } - - private static boolean isInputFormatSupported(Properties schema) - { - if (isTextInputFormat(schema)) { - if (!Objects.equals(schema.getProperty(SKIP_HEADER_COUNT_KEY, "0"), "0")) { - // S3 Select supports skipping one line of headers, but it was returning incorrect results for trino-hive-hadoop2/conf/files/test_table_with_header.csv.gz - // TODO https://github.com/trinodb/trino/issues/2349 - return false; - } - - // S3 Select does not support skipping footers - return Objects.equals(schema.getProperty(SKIP_FOOTER_COUNT_KEY, "0"), "0"); - } - - return false; - } - - public static boolean isCompressionCodecSupported(Properties schema, String path) - { - if (isTextInputFormat(schema)) { - // S3 Select supports the following formats: uncompressed, GZIP and BZIP2. - return CompressionKind.forFile(path) - .map(kind -> kind == CompressionKind.GZIP || kind == CompressionKind.BZIP2) - .orElse(true); - } - - return false; - } - - public static boolean isSplittable(boolean s3SelectPushdownEnabled, Properties schema, String path) - { - if (!s3SelectPushdownEnabled) { - return true; - } - - // S3 Select supports splitting uncompressed files - if (isTextInputFormat(schema) && CompressionKind.forFile(path).isEmpty()) { - return isSerDeSupported(schema); - } - - return false; - } - - private static boolean isTextInputFormat(Properties schema) - { - return TEXT_INPUT_FORMAT_CLASS.equals(getInputFormatName(schema).orElse(null)); - } - - private static boolean areColumnTypesSupported(List columns) - { - requireNonNull(columns, "columns is null"); - - if (columns.isEmpty()) { - return false; - } - - for (Column column : columns) { - if (!SUPPORTED_COLUMN_TYPES.contains(column.getType().getHiveTypeName().toString())) { - return false; - } - } - - return true; - } - - private static boolean isS3Storage(String path) - { - return SUPPORTED_S3_PREFIXES.stream().anyMatch(path::startsWith); - } - - public static boolean shouldEnablePushdownForTable(ConnectorSession session, Table table, String path, Optional optionalPartition) - { - if (!isS3SelectPushdownEnabled(session)) { - return false; - } - - if (path == null) { - return false; - } - - // Hive table partitions could be on different storages, - // as a result, we have to check each individual optionalPartition - Properties schema = optionalPartition - .map(partition -> getHiveSchema(partition, table)) - .orElseGet(() -> getHiveSchema(table)); - return shouldEnablePushdownForTable(table, path, schema); - } - - private static boolean shouldEnablePushdownForTable(Table table, String path, Properties schema) - { - return isS3Storage(path) && - isSerDeSupported(schema) && - isInputFormatSupported(schema) && - areColumnTypesSupported(table.getDataColumns()); - } -} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/s3select/S3SelectRecordCursor.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/s3select/S3SelectRecordCursor.java deleted file mode 100644 index ac8b03646ec2..000000000000 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/s3select/S3SelectRecordCursor.java +++ /dev/null @@ -1,81 +0,0 @@ -/* - * 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.hive.s3select; - -import com.google.common.annotations.VisibleForTesting; -import io.trino.plugin.hive.GenericHiveRecordCursor; -import io.trino.plugin.hive.HiveColumnHandle; -import org.apache.hadoop.conf.Configuration; -import org.apache.hadoop.fs.Path; -import org.apache.hadoop.io.Writable; -import org.apache.hadoop.mapred.RecordReader; - -import java.util.List; -import java.util.Properties; -import java.util.stream.Collectors; - -import static io.trino.plugin.hive.util.SerdeConstants.LIST_COLUMNS; -import static io.trino.plugin.hive.util.SerdeConstants.LIST_COLUMN_TYPES; -import static java.util.Objects.requireNonNull; - -class S3SelectRecordCursor - extends GenericHiveRecordCursor -{ - public S3SelectRecordCursor( - Configuration configuration, - Path path, - RecordReader recordReader, - long totalBytes, - Properties splitSchema, - List columns) - { - super(configuration, path, recordReader, totalBytes, updateSplitSchema(splitSchema, columns), columns); - } - - // since s3select only returns the required column, not the whole columns - // we need to update the split schema to include only the required columns - // otherwise, Serde could not deserialize output from s3select to row data correctly - @VisibleForTesting - static Properties updateSplitSchema(Properties splitSchema, List columns) - { - requireNonNull(splitSchema, "splitSchema is null"); - requireNonNull(columns, "columns is null"); - // clone split properties for update so as not to affect the original one - Properties updatedSchema = new Properties(); - updatedSchema.putAll(splitSchema); - updatedSchema.setProperty(LIST_COLUMNS, buildColumns(columns)); - updatedSchema.setProperty(LIST_COLUMN_TYPES, buildColumnTypes(columns)); - return updatedSchema; - } - - private static String buildColumns(List columns) - { - if (columns == null || columns.isEmpty()) { - return ""; - } - return columns.stream() - .map(HiveColumnHandle::getName) - .collect(Collectors.joining(",")); - } - - private static String buildColumnTypes(List columns) - { - if (columns == null || columns.isEmpty()) { - return ""; - } - return columns.stream() - .map(column -> column.getHiveType().getTypeInfo().getTypeName()) - .collect(Collectors.joining(",")); - } -} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/s3select/S3SelectRecordCursorProvider.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/s3select/S3SelectRecordCursorProvider.java deleted file mode 100644 index da318adc4df1..000000000000 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/s3select/S3SelectRecordCursorProvider.java +++ /dev/null @@ -1,198 +0,0 @@ -/* - * 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.hive.s3select; - -import com.google.common.collect.ImmutableList; -import com.google.common.collect.ImmutableSet; -import com.google.inject.Inject; -import io.trino.filesystem.Location; -import io.trino.hdfs.HdfsEnvironment; -import io.trino.plugin.hive.HiveColumnHandle; -import io.trino.plugin.hive.HiveConfig; -import io.trino.plugin.hive.HiveRecordCursorProvider; -import io.trino.plugin.hive.ReaderColumns; -import io.trino.plugin.hive.s3select.csv.S3SelectCsvRecordReader; -import io.trino.plugin.hive.type.TypeInfo; -import io.trino.spi.TrinoException; -import io.trino.spi.connector.ConnectorSession; -import io.trino.spi.connector.RecordCursor; -import io.trino.spi.predicate.TupleDomain; -import io.trino.spi.type.TypeManager; -import org.apache.hadoop.conf.Configuration; -import org.apache.hadoop.fs.Path; - -import java.io.IOException; -import java.util.Arrays; -import java.util.List; -import java.util.Optional; -import java.util.Properties; -import java.util.Set; -import java.util.function.Function; - -import static com.google.common.collect.ImmutableList.toImmutableList; -import static com.google.common.collect.ImmutableSet.toImmutableSet; -import static io.trino.plugin.hive.HiveErrorCode.HIVE_FILESYSTEM_ERROR; -import static io.trino.plugin.hive.HivePageSourceProvider.projectBaseColumns; -import static io.trino.plugin.hive.s3select.S3SelectDataType.CSV; -import static io.trino.plugin.hive.type.TypeInfoUtils.getTypeInfosFromTypeString; -import static io.trino.plugin.hive.util.HiveUtil.getDeserializerClassName; -import static io.trino.plugin.hive.util.SerdeConstants.COLUMN_NAME_DELIMITER; -import static io.trino.plugin.hive.util.SerdeConstants.LIST_COLUMNS; -import static io.trino.plugin.hive.util.SerdeConstants.LIST_COLUMN_TYPES; -import static java.util.Objects.requireNonNull; - -public class S3SelectRecordCursorProvider - implements HiveRecordCursorProvider -{ - private final HdfsEnvironment hdfsEnvironment; - private final TrinoS3ClientFactory s3ClientFactory; - private final boolean experimentalPushdownEnabled; - - @Inject - public S3SelectRecordCursorProvider(HdfsEnvironment hdfsEnvironment, TrinoS3ClientFactory s3ClientFactory, HiveConfig hiveConfig) - { - this.hdfsEnvironment = requireNonNull(hdfsEnvironment, "hdfsEnvironment is null"); - this.s3ClientFactory = requireNonNull(s3ClientFactory, "s3ClientFactory is null"); - this.experimentalPushdownEnabled = hiveConfig.isS3SelectExperimentalPushdownEnabled(); - } - - @Override - public Optional createRecordCursor( - Configuration configuration, - ConnectorSession session, - Location location, - long start, - long length, - long fileSize, - Properties schema, - List columns, - TupleDomain effectivePredicate, - TypeManager typeManager, - boolean s3SelectPushdownEnabled) - { - if (!s3SelectPushdownEnabled) { - return Optional.empty(); - } - - Path path = new Path(location.toString()); - try { - this.hdfsEnvironment.getFileSystem(session.getIdentity(), path, configuration); - } - catch (IOException e) { - throw new TrinoException(HIVE_FILESYSTEM_ERROR, "Failed getting FileSystem: " + path, e); - } - - Optional projectedReaderColumns = projectBaseColumns(columns); - // Ignore predicates on partial columns for now. - effectivePredicate = effectivePredicate.filter((column, domain) -> column.isBaseColumn()); - - List readerColumns = projectedReaderColumns - .map(readColumns -> readColumns.get().stream().map(HiveColumnHandle.class::cast).collect(toImmutableList())) - .orElseGet(() -> ImmutableList.copyOf(columns)); - // Query is not going to filter any data, no need to use S3 Select - if (!hasFilters(schema, effectivePredicate, readerColumns)) { - return Optional.empty(); - } - - String serdeName = getDeserializerClassName(schema); - Optional s3SelectDataTypeOptional = S3SelectSerDeDataTypeMapper.getDataType(serdeName); - - if (s3SelectDataTypeOptional.isPresent()) { - S3SelectDataType s3SelectDataType = s3SelectDataTypeOptional.get(); - if (s3SelectDataType == CSV && !experimentalPushdownEnabled) { - return Optional.empty(); - } - - Optional nullCharacterEncoding = Optional.empty(); - if (s3SelectDataType == CSV) { - nullCharacterEncoding = S3SelectCsvRecordReader.nullCharacterEncoding(schema); - } - IonSqlQueryBuilder queryBuilder = new IonSqlQueryBuilder(typeManager, s3SelectDataType, nullCharacterEncoding); - String ionSqlQuery = queryBuilder.buildSql(readerColumns, effectivePredicate); - Optional recordReader = S3SelectLineRecordReaderProvider.get(configuration, path, start, length, schema, - ionSqlQuery, s3ClientFactory, s3SelectDataType); - - if (recordReader.isEmpty()) { - // S3 Select data type is not mapped to an S3SelectLineRecordReader - return Optional.empty(); - } - - RecordCursor cursor = new S3SelectRecordCursor<>(configuration, path, recordReader.get(), length, schema, readerColumns); - return Optional.of(new ReaderRecordCursorWithProjections(cursor, projectedReaderColumns)); - } - // unsupported serdes - return Optional.empty(); - } - - private static boolean hasFilters( - Properties schema, - TupleDomain effectivePredicate, - List readerColumns) - { - //There are no effective predicates and readercolumns and columntypes are identical to schema - //means getting all data out of S3. We can use S3 GetObject instead of S3 SelectObjectContent in these cases. - if (effectivePredicate.isAll()) { - return !isEquivalentSchema(readerColumns, schema); - } - return true; - } - - private static boolean isEquivalentSchema(List readerColumns, Properties schema) - { - Set projectedColumnNames = getColumnProperty(readerColumns, HiveColumnHandle::getName); - Set projectedColumnTypes = getColumnProperty(readerColumns, column -> column.getHiveType().getTypeInfo().getTypeName()); - return isEquivalentColumns(projectedColumnNames, schema) && isEquivalentColumnTypes(projectedColumnTypes, schema); - } - - private static boolean isEquivalentColumns(Set projectedColumnNames, Properties schema) - { - Set columnNames; - String columnNameProperty = schema.getProperty(LIST_COLUMNS); - if (columnNameProperty.length() == 0) { - columnNames = ImmutableSet.of(); - } - else { - String columnNameDelimiter = (String) schema.getOrDefault(COLUMN_NAME_DELIMITER, ","); - columnNames = Arrays.stream(columnNameProperty.split(columnNameDelimiter)) - .collect(toImmutableSet()); - } - return projectedColumnNames.equals(columnNames); - } - - private static boolean isEquivalentColumnTypes(Set projectedColumnTypes, Properties schema) - { - String columnTypeProperty = schema.getProperty(LIST_COLUMN_TYPES); - Set columnTypes; - if (columnTypeProperty.length() == 0) { - columnTypes = ImmutableSet.of(); - } - else { - columnTypes = getTypeInfosFromTypeString(columnTypeProperty) - .stream() - .map(TypeInfo::getTypeName) - .collect(toImmutableSet()); - } - return projectedColumnTypes.equals(columnTypes); - } - - private static Set getColumnProperty(List readerColumns, Function mapper) - { - if (readerColumns.isEmpty()) { - return ImmutableSet.of(); - } - return readerColumns.stream() - .map(mapper) - .collect(toImmutableSet()); - } -} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/s3select/S3SelectSerDeDataTypeMapper.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/s3select/S3SelectSerDeDataTypeMapper.java deleted file mode 100644 index 4695eb1a7e3b..000000000000 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/s3select/S3SelectSerDeDataTypeMapper.java +++ /dev/null @@ -1,40 +0,0 @@ -/* - * 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.hive.s3select; - -import java.util.Map; -import java.util.Optional; - -import static io.trino.plugin.hive.util.HiveClassNames.JSON_SERDE_CLASS; -import static io.trino.plugin.hive.util.HiveClassNames.LAZY_SIMPLE_SERDE_CLASS; - -public class S3SelectSerDeDataTypeMapper -{ - // Contains mapping of SerDe class name -> data type. Multiple SerDe classes can be mapped to the same data type. - private static final Map serDeToDataTypeMapping = Map.of( - LAZY_SIMPLE_SERDE_CLASS, S3SelectDataType.CSV, - JSON_SERDE_CLASS, S3SelectDataType.JSON); - - private S3SelectSerDeDataTypeMapper() {} - - public static Optional getDataType(String serdeName) - { - return Optional.ofNullable(serDeToDataTypeMapping.get(serdeName)); - } - - public static boolean doesSerDeExist(String serdeName) - { - return serDeToDataTypeMapping.containsKey(serdeName); - } -} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/s3select/TrinoS3ClientFactory.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/s3select/TrinoS3ClientFactory.java deleted file mode 100644 index 9a016c8e41ec..000000000000 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/s3select/TrinoS3ClientFactory.java +++ /dev/null @@ -1,243 +0,0 @@ -/* - * 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.hive.s3select; - -import com.amazonaws.ClientConfiguration; -import com.amazonaws.Protocol; -import com.amazonaws.SdkClientException; -import com.amazonaws.auth.AWSCredentials; -import com.amazonaws.auth.AWSCredentialsProvider; -import com.amazonaws.auth.AWSStaticCredentialsProvider; -import com.amazonaws.auth.BasicAWSCredentials; -import com.amazonaws.auth.BasicSessionCredentials; -import com.amazonaws.auth.DefaultAWSCredentialsProviderChain; -import com.amazonaws.auth.STSAssumeRoleSessionCredentialsProvider; -import com.amazonaws.regions.DefaultAwsRegionProviderChain; -import com.amazonaws.services.s3.AmazonS3; -import com.amazonaws.services.s3.AmazonS3Builder; -import com.amazonaws.services.s3.AmazonS3Client; -import com.amazonaws.services.securitytoken.AWSSecurityTokenServiceClientBuilder; -import com.google.errorprone.annotations.concurrent.GuardedBy; -import com.google.inject.Inject; -import io.airlift.log.Logger; -import io.airlift.units.Duration; -import io.trino.hdfs.s3.HiveS3Config; -import io.trino.hdfs.s3.TrinoS3FileSystem; -import io.trino.plugin.hive.HiveConfig; -import org.apache.hadoop.conf.Configuration; - -import java.net.URI; -import java.util.Optional; - -import static com.amazonaws.client.builder.AwsClientBuilder.EndpointConfiguration; -import static com.amazonaws.regions.Regions.US_EAST_1; -import static com.google.common.base.Strings.isNullOrEmpty; -import static com.google.common.base.Verify.verify; -import static io.trino.hdfs.s3.AwsCurrentRegionHolder.getCurrentRegionFromEC2Metadata; -import static io.trino.hdfs.s3.TrinoS3FileSystem.S3_ACCESS_KEY; -import static io.trino.hdfs.s3.TrinoS3FileSystem.S3_CONNECT_TIMEOUT; -import static io.trino.hdfs.s3.TrinoS3FileSystem.S3_CONNECT_TTL; -import static io.trino.hdfs.s3.TrinoS3FileSystem.S3_CREDENTIALS_PROVIDER; -import static io.trino.hdfs.s3.TrinoS3FileSystem.S3_ENDPOINT; -import static io.trino.hdfs.s3.TrinoS3FileSystem.S3_EXTERNAL_ID; -import static io.trino.hdfs.s3.TrinoS3FileSystem.S3_IAM_ROLE; -import static io.trino.hdfs.s3.TrinoS3FileSystem.S3_MAX_ERROR_RETRIES; -import static io.trino.hdfs.s3.TrinoS3FileSystem.S3_PIN_CLIENT_TO_CURRENT_REGION; -import static io.trino.hdfs.s3.TrinoS3FileSystem.S3_ROLE_SESSION_NAME; -import static io.trino.hdfs.s3.TrinoS3FileSystem.S3_SECRET_KEY; -import static io.trino.hdfs.s3.TrinoS3FileSystem.S3_SESSION_TOKEN; -import static io.trino.hdfs.s3.TrinoS3FileSystem.S3_SOCKET_TIMEOUT; -import static io.trino.hdfs.s3.TrinoS3FileSystem.S3_SSL_ENABLED; -import static io.trino.hdfs.s3.TrinoS3FileSystem.S3_STS_ENDPOINT; -import static io.trino.hdfs.s3.TrinoS3FileSystem.S3_STS_REGION; -import static io.trino.hdfs.s3.TrinoS3FileSystem.S3_USER_AGENT_PREFIX; -import static java.lang.Math.toIntExact; -import static java.lang.String.format; - -/** - * This factory provides AmazonS3 client required for executing S3SelectPushdown requests. - * Normal S3 GET requests use AmazonS3 clients initialized in {@link TrinoS3FileSystem} or EMRFS. - * The ideal state will be to merge this logic with the two file systems and get rid of this - * factory class. - * Please do not use the client provided by this factory for any other use cases. - */ -public class TrinoS3ClientFactory -{ - private static final Logger log = Logger.get(TrinoS3ClientFactory.class); - private static final String S3_SELECT_PUSHDOWN_MAX_CONNECTIONS = "hive.s3select-pushdown.max-connections"; - - private final boolean enabled; - private final int defaultMaxConnections; - - @GuardedBy("this") - private AmazonS3 s3Client; - - @Inject - public TrinoS3ClientFactory(HiveConfig config) - { - this.enabled = config.isS3SelectPushdownEnabled(); - this.defaultMaxConnections = config.getS3SelectPushdownMaxConnections(); - } - - synchronized AmazonS3 getS3Client(Configuration config) - { - if (s3Client == null) { - s3Client = createS3Client(config); - } - return s3Client; - } - - private AmazonS3 createS3Client(Configuration config) - { - HiveS3Config defaults = new HiveS3Config(); - String userAgentPrefix = config.get(S3_USER_AGENT_PREFIX, defaults.getS3UserAgentPrefix()); - int maxErrorRetries = config.getInt(S3_MAX_ERROR_RETRIES, defaults.getS3MaxErrorRetries()); - boolean sslEnabled = config.getBoolean(S3_SSL_ENABLED, defaults.isS3SslEnabled()); - Duration connectTimeout = Duration.valueOf(config.get(S3_CONNECT_TIMEOUT, defaults.getS3ConnectTimeout().toString())); - Duration socketTimeout = Duration.valueOf(config.get(S3_SOCKET_TIMEOUT, defaults.getS3SocketTimeout().toString())); - int maxConnections = config.getInt(S3_SELECT_PUSHDOWN_MAX_CONNECTIONS, defaultMaxConnections); - - ClientConfiguration clientConfiguration = new ClientConfiguration() - .withMaxErrorRetry(maxErrorRetries) - .withProtocol(sslEnabled ? Protocol.HTTPS : Protocol.HTTP) - .withConnectionTimeout(toIntExact(connectTimeout.toMillis())) - .withSocketTimeout(toIntExact(socketTimeout.toMillis())) - .withMaxConnections(maxConnections) - .withUserAgentPrefix(userAgentPrefix) - .withUserAgentSuffix(enabled ? "Trino-select" : "Trino"); - - String connectTtlValue = config.get(S3_CONNECT_TTL); - if (!isNullOrEmpty(connectTtlValue)) { - clientConfiguration.setConnectionTTL(Duration.valueOf(connectTtlValue).toMillis()); - } - - AWSCredentialsProvider awsCredentialsProvider = getAwsCredentialsProvider(config); - AmazonS3Builder, ? extends AmazonS3> clientBuilder = AmazonS3Client.builder() - .withCredentials(awsCredentialsProvider) - .withClientConfiguration(clientConfiguration) - .withMetricsCollector(TrinoS3FileSystem.getFileSystemStats().newRequestMetricCollector()) - .enablePathStyleAccess(); - - boolean regionOrEndpointSet = false; - - String endpoint = config.get(S3_ENDPOINT); - boolean pinS3ClientToCurrentRegion = config.getBoolean(S3_PIN_CLIENT_TO_CURRENT_REGION, defaults.isPinS3ClientToCurrentRegion()); - verify(!pinS3ClientToCurrentRegion || endpoint == null, - "Invalid configuration: either endpoint can be set or S3 client can be pinned to the current region"); - - // use local region when running inside of EC2 - if (pinS3ClientToCurrentRegion) { - clientBuilder.setRegion(getCurrentRegionFromEC2Metadata().getName()); - regionOrEndpointSet = true; - } - - if (!isNullOrEmpty(endpoint)) { - clientBuilder.withEndpointConfiguration(new EndpointConfiguration(endpoint, null)); - regionOrEndpointSet = true; - } - - if (!regionOrEndpointSet) { - clientBuilder.withRegion(US_EAST_1); - clientBuilder.setForceGlobalBucketAccessEnabled(true); - } - - return clientBuilder.build(); - } - - private static AWSCredentialsProvider getAwsCredentialsProvider(Configuration conf) - { - Optional credentials = getAwsCredentials(conf); - if (credentials.isPresent()) { - return new AWSStaticCredentialsProvider(credentials.get()); - } - - String providerClass = conf.get(S3_CREDENTIALS_PROVIDER); - if (!isNullOrEmpty(providerClass)) { - return getCustomAWSCredentialsProvider(conf, providerClass); - } - - AWSCredentialsProvider provider = getAwsCredentials(conf) - .map(value -> (AWSCredentialsProvider) new AWSStaticCredentialsProvider(value)) - .orElseGet(DefaultAWSCredentialsProviderChain::getInstance); - - String iamRole = conf.get(S3_IAM_ROLE); - if (iamRole != null) { - String stsEndpointOverride = conf.get(S3_STS_ENDPOINT); - String stsRegionOverride = conf.get(S3_STS_REGION); - String s3RoleSessionName = conf.get(S3_ROLE_SESSION_NAME); - String externalId = conf.get(S3_EXTERNAL_ID); - - AWSSecurityTokenServiceClientBuilder stsClientBuilder = AWSSecurityTokenServiceClientBuilder.standard() - .withCredentials(provider); - - String region; - if (!isNullOrEmpty(stsRegionOverride)) { - region = stsRegionOverride; - } - else { - DefaultAwsRegionProviderChain regionProviderChain = new DefaultAwsRegionProviderChain(); - try { - region = regionProviderChain.getRegion(); - } - catch (SdkClientException ex) { - log.warn("Falling back to default AWS region %s", US_EAST_1); - region = US_EAST_1.getName(); - } - } - - if (!isNullOrEmpty(stsEndpointOverride)) { - stsClientBuilder.withEndpointConfiguration(new EndpointConfiguration(stsEndpointOverride, region)); - } - else { - stsClientBuilder.withRegion(region); - } - - provider = new STSAssumeRoleSessionCredentialsProvider.Builder(iamRole, s3RoleSessionName) - .withExternalId(externalId) - .withStsClient(stsClientBuilder.build()) - .build(); - } - - return provider; - } - - private static AWSCredentialsProvider getCustomAWSCredentialsProvider(Configuration conf, String providerClass) - { - try { - return conf.getClassByName(providerClass) - .asSubclass(AWSCredentialsProvider.class) - .getConstructor(URI.class, Configuration.class) - .newInstance(null, conf); - } - catch (ReflectiveOperationException e) { - throw new RuntimeException(format("Error creating an instance of %s", providerClass), e); - } - } - - private static Optional getAwsCredentials(Configuration conf) - { - String accessKey = conf.get(S3_ACCESS_KEY); - String secretKey = conf.get(S3_SECRET_KEY); - - if (isNullOrEmpty(accessKey) || isNullOrEmpty(secretKey)) { - return Optional.empty(); - } - String sessionToken = conf.get(S3_SESSION_TOKEN); - if (!isNullOrEmpty(sessionToken)) { - return Optional.of(new BasicSessionCredentials(accessKey, secretKey, sessionToken)); - } - - return Optional.of(new BasicAWSCredentials(accessKey, secretKey)); - } -} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/s3select/TrinoS3SelectClient.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/s3select/TrinoS3SelectClient.java deleted file mode 100644 index e42777ac40f8..000000000000 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/s3select/TrinoS3SelectClient.java +++ /dev/null @@ -1,85 +0,0 @@ -/* - * 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.hive.s3select; - -import com.amazonaws.services.s3.AmazonS3; -import com.amazonaws.services.s3.model.SelectObjectContentEventVisitor; -import com.amazonaws.services.s3.model.SelectObjectContentRequest; -import com.amazonaws.services.s3.model.SelectObjectContentResult; -import org.apache.hadoop.conf.Configuration; - -import java.io.Closeable; -import java.io.IOException; -import java.io.InputStream; - -import static com.amazonaws.services.s3.model.SelectObjectContentEvent.EndEvent; -import static java.util.Objects.requireNonNull; - -class TrinoS3SelectClient - implements Closeable -{ - private final AmazonS3 s3Client; - private boolean requestComplete; - private SelectObjectContentRequest selectObjectRequest; - private SelectObjectContentResult selectObjectContentResult; - - public TrinoS3SelectClient(Configuration configuration, TrinoS3ClientFactory s3ClientFactory) - { - requireNonNull(configuration, "configuration is null"); - requireNonNull(s3ClientFactory, "s3ClientFactory is null"); - this.s3Client = s3ClientFactory.getS3Client(configuration); - } - - public InputStream getRecordsContent(SelectObjectContentRequest selectObjectRequest) - { - this.selectObjectRequest = requireNonNull(selectObjectRequest, "selectObjectRequest is null"); - this.selectObjectContentResult = s3Client.selectObjectContent(selectObjectRequest); - return selectObjectContentResult.getPayload() - .getRecordsInputStream( - new SelectObjectContentEventVisitor() - { - @Override - public void visit(EndEvent endEvent) - { - requestComplete = true; - } - }); - } - - @Override - public void close() - throws IOException - { - selectObjectContentResult.close(); - } - - public String getKeyName() - { - return selectObjectRequest.getKey(); - } - - public String getBucketName() - { - return selectObjectRequest.getBucketName(); - } - - /** - * The End Event indicates all matching records have been transmitted. - * If the End Event is not received, the results may be incomplete. - */ - public boolean isRequestComplete() - { - return requestComplete; - } -} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/s3select/csv/S3SelectCsvRecordReader.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/s3select/csv/S3SelectCsvRecordReader.java deleted file mode 100644 index 89dc6d3c9102..000000000000 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/s3select/csv/S3SelectCsvRecordReader.java +++ /dev/null @@ -1,114 +0,0 @@ -/* - * 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.hive.s3select.csv; - -import com.amazonaws.services.s3.model.CSVInput; -import com.amazonaws.services.s3.model.CSVOutput; -import com.amazonaws.services.s3.model.CompressionType; -import com.amazonaws.services.s3.model.InputSerialization; -import com.amazonaws.services.s3.model.OutputSerialization; -import io.trino.plugin.hive.s3select.S3SelectLineRecordReader; -import io.trino.plugin.hive.s3select.TrinoS3ClientFactory; -import io.trino.plugin.hive.util.SerdeConstants; -import org.apache.hadoop.conf.Configuration; -import org.apache.hadoop.fs.Path; - -import java.util.Optional; -import java.util.Properties; - -import static io.trino.plugin.hive.util.SerdeConstants.ESCAPE_CHAR; -import static io.trino.plugin.hive.util.SerdeConstants.FIELD_DELIM; -import static io.trino.plugin.hive.util.SerdeConstants.QUOTE_CHAR; - -public class S3SelectCsvRecordReader - extends S3SelectLineRecordReader -{ - /* - * Sentinel unicode comment character (http://www.unicode.org/faq/private_use.html#nonchar_codes). - * It is expected that \uFDD0 sentinel comment character is not the first character in any row of user's CSV S3 object. - * The rows starting with \uFDD0 will be skipped by S3Select and will not be a part of the result set or aggregations. - * To process CSV objects that may contain \uFDD0 as first row character please disable S3SelectPushdown. - * TODO: Remove this proxy logic when S3Select API supports disabling of row level comments. - */ - - private static final String COMMENTS_CHAR_STR = "\uFDD0"; - private static final String DEFAULT_FIELD_DELIMITER = ","; - - public S3SelectCsvRecordReader( - Configuration configuration, - Path path, - long start, - long length, - Properties schema, - String ionSqlQuery, - TrinoS3ClientFactory s3ClientFactory) - { - super(configuration, path, start, length, schema, ionSqlQuery, s3ClientFactory); - } - - @Override - public InputSerialization buildInputSerialization() - { - Properties schema = getSchema(); - String fieldDelimiter = schema.getProperty(FIELD_DELIM, DEFAULT_FIELD_DELIMITER); - String quoteChar = schema.getProperty(QUOTE_CHAR, null); - String escapeChar = schema.getProperty(ESCAPE_CHAR, null); - - CSVInput selectObjectCSVInputSerialization = new CSVInput(); - selectObjectCSVInputSerialization.setRecordDelimiter(getLineDelimiter()); - selectObjectCSVInputSerialization.setFieldDelimiter(fieldDelimiter); - selectObjectCSVInputSerialization.setComments(COMMENTS_CHAR_STR); - selectObjectCSVInputSerialization.setQuoteCharacter(quoteChar); - selectObjectCSVInputSerialization.setQuoteEscapeCharacter(escapeChar); - - InputSerialization selectObjectInputSerialization = new InputSerialization(); - selectObjectInputSerialization.setCompressionType(getCompressionType()); - selectObjectInputSerialization.setCsv(selectObjectCSVInputSerialization); - - return selectObjectInputSerialization; - } - - @Override - public OutputSerialization buildOutputSerialization() - { - Properties schema = getSchema(); - String fieldDelimiter = schema.getProperty(FIELD_DELIM, DEFAULT_FIELD_DELIMITER); - String quoteChar = schema.getProperty(QUOTE_CHAR, null); - String escapeChar = schema.getProperty(ESCAPE_CHAR, null); - - OutputSerialization selectObjectOutputSerialization = new OutputSerialization(); - CSVOutput selectObjectCSVOutputSerialization = new CSVOutput(); - selectObjectCSVOutputSerialization.setRecordDelimiter(getLineDelimiter()); - selectObjectCSVOutputSerialization.setFieldDelimiter(fieldDelimiter); - selectObjectCSVOutputSerialization.setQuoteCharacter(quoteChar); - selectObjectCSVOutputSerialization.setQuoteEscapeCharacter(escapeChar); - selectObjectOutputSerialization.setCsv(selectObjectCSVOutputSerialization); - - return selectObjectOutputSerialization; - } - - @Override - public boolean shouldEnableScanRange() - { - // Works for CSV if AllowQuotedRecordDelimiter is disabled. - boolean isQuotedRecordDelimiterAllowed = Boolean.TRUE.equals( - buildInputSerialization().getCsv().getAllowQuotedRecordDelimiter()); - return CompressionType.NONE.equals(getCompressionType()) && !isQuotedRecordDelimiterAllowed; - } - - public static Optional nullCharacterEncoding(Properties schema) - { - return Optional.ofNullable(schema.getProperty(SerdeConstants.SERIALIZATION_NULL_FORMAT)); - } -} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/s3select/json/S3SelectJsonRecordReader.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/s3select/json/S3SelectJsonRecordReader.java deleted file mode 100644 index fa7d7be84654..000000000000 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/s3select/json/S3SelectJsonRecordReader.java +++ /dev/null @@ -1,72 +0,0 @@ -/* - * 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.hive.s3select.json; - -import com.amazonaws.services.s3.model.CompressionType; -import com.amazonaws.services.s3.model.InputSerialization; -import com.amazonaws.services.s3.model.JSONInput; -import com.amazonaws.services.s3.model.JSONOutput; -import com.amazonaws.services.s3.model.JSONType; -import com.amazonaws.services.s3.model.OutputSerialization; -import io.trino.plugin.hive.s3select.S3SelectLineRecordReader; -import io.trino.plugin.hive.s3select.TrinoS3ClientFactory; -import org.apache.hadoop.conf.Configuration; -import org.apache.hadoop.fs.Path; - -import java.util.Properties; - -public class S3SelectJsonRecordReader - extends S3SelectLineRecordReader -{ - public S3SelectJsonRecordReader(Configuration configuration, - Path path, - long start, - long length, - Properties schema, - String ionSqlQuery, - TrinoS3ClientFactory s3ClientFactory) - { - super(configuration, path, start, length, schema, ionSqlQuery, s3ClientFactory); - } - - @Override - public InputSerialization buildInputSerialization() - { - // JSONType.LINES is the only JSON format supported by the Hive JsonSerDe. - JSONInput selectObjectJSONInputSerialization = new JSONInput(); - selectObjectJSONInputSerialization.setType(JSONType.LINES); - - InputSerialization selectObjectInputSerialization = new InputSerialization(); - selectObjectInputSerialization.setCompressionType(getCompressionType()); - selectObjectInputSerialization.setJson(selectObjectJSONInputSerialization); - - return selectObjectInputSerialization; - } - - @Override - public OutputSerialization buildOutputSerialization() - { - OutputSerialization selectObjectOutputSerialization = new OutputSerialization(); - JSONOutput selectObjectJSONOutputSerialization = new JSONOutput(); - selectObjectOutputSerialization.setJson(selectObjectJSONOutputSerialization); - - return selectObjectOutputSerialization; - } - - @Override - public boolean shouldEnableScanRange() - { - return CompressionType.NONE.equals(getCompressionType()); - } -} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/security/HiveSecurityModule.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/security/HiveSecurityModule.java index 5a5a3f51b138..f9a38824b559 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/security/HiveSecurityModule.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/security/HiveSecurityModule.java @@ -41,20 +41,23 @@ protected void setup(Binder binder) LEGACY, combine( new LegacySecurityModule(), - new StaticAccessControlMetadataModule())); + new StaticAccessControlMetadataModule(), + usingSystemSecurity(false))); bindSecurityModule( FILE, combine( new FileBasedAccessControlModule(), - new StaticAccessControlMetadataModule())); + new StaticAccessControlMetadataModule(), + usingSystemSecurity(false))); bindSecurityModule( READ_ONLY, combine( new ReadOnlySecurityModule(), - new StaticAccessControlMetadataModule())); - bindSecurityModule(SQL_STANDARD, new SqlStandardSecurityModule()); - bindSecurityModule(ALLOW_ALL, new AllowAllSecurityModule()); - bindSecurityModule(SYSTEM, new SystemSecurityModule()); + new StaticAccessControlMetadataModule(), + usingSystemSecurity(false))); + bindSecurityModule(SQL_STANDARD, combine(new SqlStandardSecurityModule(), usingSystemSecurity(false))); + bindSecurityModule(ALLOW_ALL, combine(new AllowAllSecurityModule(), usingSystemSecurity(false))); + bindSecurityModule(SYSTEM, combine(new SystemSecurityModule(), usingSystemSecurity(true))); } private void bindSecurityModule(String name, Module module) @@ -65,6 +68,11 @@ private void bindSecurityModule(String name, Module module) module)); } + private static Module usingSystemSecurity(boolean system) + { + return binder -> binder.bind(boolean.class).annotatedWith(UsingSystemSecurity.class).toInstance(system); + } + private static class StaticAccessControlMetadataModule implements Module { diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/security/SqlStandardAccessControlMetadataMetastore.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/security/SqlStandardAccessControlMetadataMetastore.java index 9d5d60230d5f..4df8f99988e6 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/security/SqlStandardAccessControlMetadataMetastore.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/security/SqlStandardAccessControlMetadataMetastore.java @@ -39,8 +39,6 @@ public interface SqlStandardAccessControlMetadataMetastore void revokeRoles(Set roles, Set grantees, boolean adminOption, HivePrincipal grantor); - Set listGrantedPrincipals(String role); - Optional getDatabaseOwner(String databaseName); void revokeTablePrivileges(String databaseName, String tableName, HivePrincipal grantee, HivePrincipal grantor, Set privileges, boolean grantOption); diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/ForGlueColumnStatisticsWrite.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/security/UsingSystemSecurity.java similarity index 91% rename from plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/ForGlueColumnStatisticsWrite.java rename to plugin/trino-hive/src/main/java/io/trino/plugin/hive/security/UsingSystemSecurity.java index f6f7f2731ac5..7d7cbeaa8131 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/ForGlueColumnStatisticsWrite.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/security/UsingSystemSecurity.java @@ -11,7 +11,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package io.trino.plugin.hive.metastore.glue; +package io.trino.plugin.hive.security; import com.google.inject.BindingAnnotation; @@ -26,4 +26,4 @@ @Retention(RUNTIME) @Target({FIELD, PARAMETER, METHOD}) @BindingAnnotation -public @interface ForGlueColumnStatisticsWrite {} +public @interface UsingSystemSecurity {} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/statistics/MetastoreHiveStatisticsProvider.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/statistics/MetastoreHiveStatisticsProvider.java index ded2ad1a766c..e4c90840bb0d 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/statistics/MetastoreHiveStatisticsProvider.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/statistics/MetastoreHiveStatisticsProvider.java @@ -86,6 +86,8 @@ import static io.trino.spi.type.TinyintType.TINYINT; import static java.lang.Double.isFinite; import static java.lang.Double.isNaN; +import static java.lang.Math.max; +import static java.lang.Math.min; import static java.lang.String.format; import static java.util.Collections.unmodifiableList; import static java.util.Objects.requireNonNull; @@ -238,8 +240,8 @@ private static void validateColumnStatistics(SchemaTableName table, String parti { columnStatistics.getMaxValueSizeInBytes().ifPresent(maxValueSizeInBytes -> checkStatistics(maxValueSizeInBytes >= 0, table, partition, column, "maxValueSizeInBytes must be greater than or equal to zero: %s", maxValueSizeInBytes)); - columnStatistics.getTotalSizeInBytes().ifPresent(totalSizeInBytes -> - checkStatistics(totalSizeInBytes >= 0, table, partition, column, "totalSizeInBytes must be greater than or equal to zero: %s", totalSizeInBytes)); + columnStatistics.getAverageColumnLength().ifPresent(averageColumnLength -> + checkStatistics(averageColumnLength >= 0, table, partition, column, "averageColumnLength must be greater than or equal to zero: %s", averageColumnLength)); columnStatistics.getNullsCount().ifPresent(nullsCount -> { checkStatistics(nullsCount >= 0, table, partition, column, "nullsCount must be greater than or equal to zero: %s", nullsCount); if (rowCount.isPresent()) { @@ -253,29 +255,8 @@ private static void validateColumnStatistics(SchemaTableName table, String parti rowCount.getAsLong()); } }); - columnStatistics.getDistinctValuesCount().ifPresent(distinctValuesCount -> { + columnStatistics.getDistinctValuesWithNullCount().ifPresent(distinctValuesCount -> { checkStatistics(distinctValuesCount >= 0, table, partition, column, "distinctValuesCount must be greater than or equal to zero: %s", distinctValuesCount); - if (rowCount.isPresent()) { - checkStatistics( - distinctValuesCount <= rowCount.getAsLong(), - table, - partition, - column, - "distinctValuesCount must be less than or equal to rowCount. distinctValuesCount: %s. rowCount: %s.", - distinctValuesCount, - rowCount.getAsLong()); - } - if (rowCount.isPresent() && columnStatistics.getNullsCount().isPresent()) { - long nonNullsCount = rowCount.getAsLong() - columnStatistics.getNullsCount().getAsLong(); - checkStatistics( - distinctValuesCount <= nonNullsCount, - table, - partition, - column, - "distinctValuesCount must be less than or equal to nonNullsCount. distinctValuesCount: %s. nonNullsCount: %s.", - distinctValuesCount, - nonNullsCount); - } }); columnStatistics.getIntegerStatistics().ifPresent(integerStatistics -> { @@ -715,7 +696,7 @@ static ColumnStatistics createDataColumnStatistics(String column, Type type, dou } return ColumnStatistics.builder() - .setDistinctValuesCount(calculateDistinctValuesCount(columnStatistics)) + .setDistinctValuesCount(calculateDistinctValuesCount(column, partitionStatistics)) .setNullsFraction(calculateNullsFraction(column, partitionStatistics)) .setDataSize(calculateDataSize(column, partitionStatistics, rowsCount)) .setRange(calculateRange(type, columnStatistics)) @@ -723,10 +704,10 @@ static ColumnStatistics createDataColumnStatistics(String column, Type type, dou } @VisibleForTesting - static Estimate calculateDistinctValuesCount(List columnStatistics) + static Estimate calculateDistinctValuesCount(String column, Collection partitionStatistics) { - return columnStatistics.stream() - .map(MetastoreHiveStatisticsProvider::getDistinctValuesCount) + return partitionStatistics.stream() + .map(statistics -> getDistinctValuesCount(column, statistics)) .filter(OptionalLong::isPresent) .map(OptionalLong::getAsLong) .peek(distinctValuesCount -> verify(distinctValuesCount >= 0, "distinctValuesCount must be greater than or equal to zero")) @@ -735,8 +716,13 @@ static Estimate calculateDistinctValuesCount(List columnSt .orElse(Estimate.unknown()); } - private static OptionalLong getDistinctValuesCount(HiveColumnStatistics statistics) + static OptionalLong getDistinctValuesCount(String column, PartitionStatistics partitionStatistics) { + HiveColumnStatistics statistics = partitionStatistics.getColumnStatistics().get(column); + if (statistics == null) { + return OptionalLong.empty(); + } + if (statistics.getBooleanStatistics().isPresent() && statistics.getBooleanStatistics().get().getFalseCount().isPresent() && statistics.getBooleanStatistics().get().getTrueCount().isPresent()) { @@ -744,10 +730,27 @@ private static OptionalLong getDistinctValuesCount(HiveColumnStatistics statisti long trueCount = statistics.getBooleanStatistics().get().getTrueCount().getAsLong(); return OptionalLong.of((falseCount > 0 ? 1 : 0) + (trueCount > 0 ? 1 : 0)); } - if (statistics.getDistinctValuesCount().isPresent()) { - return statistics.getDistinctValuesCount(); + + if (statistics.getDistinctValuesWithNullCount().isEmpty()) { + return OptionalLong.empty(); + } + + long distinctValuesCount = statistics.getDistinctValuesWithNullCount().getAsLong(); + + // Hive includes nulls in the distinct values count, but Trino does not + long nullsCount = statistics.getNullsCount().orElse(0); + if (distinctValuesCount > 0 && nullsCount > 0) { + distinctValuesCount--; + } + + // if there is non-null row the distinct values count should be at least 1 + if (distinctValuesCount == 0 && nullsCount < partitionStatistics.getBasicStatistics().getRowCount().orElse(0)) { + distinctValuesCount = 1; } - return OptionalLong.empty(); + + // Hive can produce distinct values that are much larger than the actual number of rows in the partition + distinctValuesCount = min(distinctValuesCount, partitionStatistics.getBasicStatistics().getRowCount().orElse(Long.MAX_VALUE) - nullsCount); + return OptionalLong.of(distinctValuesCount); } @VisibleForTesting @@ -808,7 +811,7 @@ static Estimate calculateDataSize(String column, Collection if (columnStatistics == null) { return false; } - return columnStatistics.getTotalSizeInBytes().isPresent(); + return columnStatistics.getAverageColumnLength().isPresent(); }) .collect(toImmutableList()); @@ -817,16 +820,21 @@ static Estimate calculateDataSize(String column, Collection } long knownRowCount = 0; - long knownDataSize = 0; + double knownDataSize = 0; for (PartitionStatistics statistics : statisticsWithKnownRowCountAndDataSize) { long rowCount = statistics.getBasicStatistics().getRowCount().orElseThrow(() -> new VerifyException("rowCount is not present")); verify(rowCount >= 0, "rowCount must be greater than or equal to zero"); HiveColumnStatistics columnStatistics = statistics.getColumnStatistics().get(column); verifyNotNull(columnStatistics, "columnStatistics is null"); - long dataSize = columnStatistics.getTotalSizeInBytes().orElseThrow(() -> new VerifyException("totalSizeInBytes is not present")); - verify(dataSize >= 0, "dataSize must be greater than or equal to zero"); + + long nullCount = columnStatistics.getNullsCount().orElse(0); + verify(nullCount >= 0, "nullCount must be greater than or equal to zero"); + long nonNullRowCount = max(rowCount - nullCount, 0); + + double averageColumnLength = columnStatistics.getAverageColumnLength().orElseThrow(() -> new VerifyException("averageColumnLength is not present")); + verify(averageColumnLength >= 0, "averageColumnLength must be greater than or equal to zero"); knownRowCount += rowCount; - knownDataSize += dataSize; + knownDataSize += averageColumnLength * nonNullRowCount; } if (totalRowCount == 0) { @@ -837,7 +845,7 @@ static Estimate calculateDataSize(String column, Collection return Estimate.unknown(); } - double averageValueDataSizeInBytes = ((double) knownDataSize) / knownRowCount; + double averageValueDataSizeInBytes = knownDataSize / knownRowCount; return Estimate.of(averageValueDataSizeInBytes * totalRowCount); } diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/util/AcidTables.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/util/AcidTables.java index a7f92eb734e6..0a7151677947 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/util/AcidTables.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/util/AcidTables.java @@ -33,6 +33,7 @@ import java.util.List; import java.util.Map; import java.util.Optional; +import java.util.Properties; import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.collect.Multimaps.asMap; @@ -56,17 +57,33 @@ public static boolean isInsertOnlyTable(Map parameters) return "insert_only".equalsIgnoreCase(parameters.get(TABLE_TRANSACTIONAL_PROPERTIES)); } + public static boolean isInsertOnlyTable(Properties parameters) + { + return "insert_only".equalsIgnoreCase(parameters.getProperty(TABLE_TRANSACTIONAL_PROPERTIES)); + } + public static boolean isTransactionalTable(Map parameters) { return "true".equalsIgnoreCase(parameters.get(TABLE_IS_TRANSACTIONAL)) || "true".equalsIgnoreCase(parameters.get(TABLE_IS_TRANSACTIONAL.toUpperCase(ENGLISH))); } + public static boolean isTransactionalTable(Properties parameters) + { + return "true".equalsIgnoreCase(parameters.getProperty(TABLE_IS_TRANSACTIONAL)) || + "true".equalsIgnoreCase(parameters.getProperty(TABLE_IS_TRANSACTIONAL.toUpperCase(ENGLISH))); + } + public static boolean isFullAcidTable(Map parameters) { return isTransactionalTable(parameters) && !isInsertOnlyTable(parameters); } + public static boolean isFullAcidTable(Properties parameters) + { + return isTransactionalTable(parameters) && !isInsertOnlyTable(parameters); + } + public static Location bucketFileName(Location subdir, int bucket) { return subdir.appendPath("bucket_%05d".formatted(bucket)); diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/util/CompressionConfigUtil.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/util/CompressionConfigUtil.java index eaea72e21682..60dc3fc712ac 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/util/CompressionConfigUtil.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/util/CompressionConfigUtil.java @@ -48,7 +48,7 @@ public static void configureCompression(Configuration config, HiveCompressionCod } // For Parquet - config.set(ParquetOutputFormat.COMPRESSION, compressionCodec.getParquetCompressionCodec().name()); + compressionCodec.getParquetCompressionCodec().ifPresent(codec -> config.set(ParquetOutputFormat.COMPRESSION, codec.name())); // For Avro compressionCodec.getAvroCompressionKind().ifPresent(kind -> config.set(AvroJob.OUTPUT_CODEC, kind.toString())); diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/util/HiveBucketing.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/util/HiveBucketing.java index 7497161f660d..468f6fa4a22b 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/util/HiveBucketing.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/util/HiveBucketing.java @@ -18,14 +18,13 @@ import com.google.common.annotations.VisibleForTesting; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableSet; -import io.trino.plugin.hive.HiveBucketHandle; import io.trino.plugin.hive.HiveBucketProperty; import io.trino.plugin.hive.HiveColumnHandle; import io.trino.plugin.hive.HiveTableHandle; +import io.trino.plugin.hive.HiveTablePartitioning; import io.trino.plugin.hive.HiveTimestampPrecision; import io.trino.plugin.hive.HiveType; import io.trino.plugin.hive.metastore.Column; -import io.trino.plugin.hive.metastore.SortingColumn; import io.trino.plugin.hive.metastore.Table; import io.trino.plugin.hive.type.ListTypeInfo; import io.trino.plugin.hive.type.MapTypeInfo; @@ -58,6 +57,7 @@ import static io.trino.plugin.hive.HiveColumnHandle.BUCKET_COLUMN_NAME; import static io.trino.plugin.hive.HiveErrorCode.HIVE_INVALID_METADATA; import static io.trino.plugin.hive.HiveSessionProperties.getTimestampPrecision; +import static io.trino.plugin.hive.HiveSessionProperties.isParallelPartitionedBucketedWrites; import static io.trino.plugin.hive.util.HiveBucketing.BucketingVersion.BUCKETING_V1; import static io.trino.plugin.hive.util.HiveBucketing.BucketingVersion.BUCKETING_V2; import static io.trino.plugin.hive.util.HiveUtil.SPARK_TABLE_PROVIDER_KEY; @@ -172,7 +172,17 @@ static int getBucketNumber(int hashCode, int bucketCount) return (hashCode & Integer.MAX_VALUE) % bucketCount; } - public static Optional getHiveBucketHandle(ConnectorSession session, Table table, TypeManager typeManager) + public static Optional getHiveTablePartitioningForRead(ConnectorSession session, Table table, TypeManager typeManager) + { + return getHiveTablePartitioning(false, session, table, typeManager); + } + + public static Optional getHiveTablePartitioningForWrite(ConnectorSession session, Table table, TypeManager typeManager) + { + return getHiveTablePartitioning(true, session, table, typeManager); + } + + private static Optional getHiveTablePartitioning(boolean forWrite, ConnectorSession session, Table table, TypeManager typeManager) { if (table.getParameters().containsKey(SPARK_TABLE_PROVIDER_KEY)) { return Optional.empty(); @@ -202,19 +212,23 @@ public static Optional getHiveBucketHandle(ConnectorSession se bucketColumns.add(bucketColumnHandle); } - BucketingVersion bucketingVersion = hiveBucketProperty.get().getBucketingVersion(); - int bucketCount = hiveBucketProperty.get().getBucketCount(); - List sortedBy = hiveBucketProperty.get().getSortedBy(); - return Optional.of(new HiveBucketHandle(bucketColumns.build(), bucketingVersion, bucketCount, bucketCount, sortedBy)); + return Optional.of(new HiveTablePartitioning( + forWrite, + getBucketingVersion(table.getParameters()), + hiveBucketProperty.get().getBucketCount(), + bucketColumns.build(), + forWrite && !table.getPartitionColumns().isEmpty() && isParallelPartitionedBucketedWrites(session), + hiveBucketProperty.get().getSortedBy(), + forWrite)); } public static Optional getHiveBucketFilter(HiveTableHandle hiveTable, TupleDomain effectivePredicate) { - if (hiveTable.getBucketHandle().isEmpty()) { + if (hiveTable.getTablePartitioning().isEmpty()) { return Optional.empty(); } - HiveBucketProperty hiveBucketProperty = hiveTable.getBucketHandle().get().toTableBucketProperty(); + HiveBucketProperty hiveBucketProperty = hiveTable.getTablePartitioning().get().toTableBucketProperty(); List dataColumns = hiveTable.getDataColumns().stream() .collect(toImmutableList()); @@ -222,7 +236,8 @@ public static Optional getHiveBucketFilter(HiveTableHandle hiv if (bindings.isEmpty()) { return Optional.empty(); } - Optional> buckets = getHiveBuckets(hiveBucketProperty, dataColumns, bindings.get()); + BucketingVersion bucketingVersion = hiveTable.getTablePartitioning().get().partitioningHandle().getBucketingVersion(); + Optional> buckets = getHiveBuckets(bucketingVersion, hiveBucketProperty, dataColumns, bindings.get()); if (buckets.isPresent()) { return Optional.of(new HiveBucketFilter(buckets.get())); } @@ -246,7 +261,7 @@ public static Optional getHiveBucketFilter(HiveTableHandle hiv return Optional.of(new HiveBucketFilter(builder.build())); } - private static Optional> getHiveBuckets(HiveBucketProperty hiveBucketProperty, List dataColumns, Map> bindings) + private static Optional> getHiveBuckets(BucketingVersion bucketingVersion, HiveBucketProperty hiveBucketProperty, List dataColumns, Map> bindings) { if (bindings.isEmpty()) { return Optional.empty(); @@ -291,7 +306,7 @@ private static Optional> getHiveBuckets(HiveBucketProperty hiveBuck .collect(toImmutableList()); return getHiveBuckets( - hiveBucketProperty.getBucketingVersion(), + bucketingVersion, hiveBucketProperty.getBucketCount(), typeInfos, orderedBindings); @@ -313,12 +328,12 @@ public static BucketingVersion getBucketingVersion(Map tableProp public static boolean isSupportedBucketing(Table table) { - return isSupportedBucketing(table.getStorage().getBucketProperty().orElseThrow(), table.getDataColumns(), table.getTableName()); + return isSupportedBucketing(table.getStorage().getBucketProperty().orElseThrow().getBucketedBy(), table.getDataColumns(), table.getTableName()); } - public static boolean isSupportedBucketing(HiveBucketProperty bucketProperty, List dataColumns, String tableName) + public static boolean isSupportedBucketing(List bucketedBy, List dataColumns, String tableName) { - return bucketProperty.getBucketedBy().stream() + return bucketedBy.stream() .map(columnName -> dataColumns.stream().filter(column -> column.getName().equals(columnName)).findFirst() .orElseThrow(() -> new IllegalArgumentException(format("Cannot find column '%s' in %s", columnName, tableName)))) .map(Column::getType) diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/util/HiveReaderUtil.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/util/HiveReaderUtil.java deleted file mode 100644 index 80d82ead8531..000000000000 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/util/HiveReaderUtil.java +++ /dev/null @@ -1,287 +0,0 @@ -/* - * 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.hive.util; - -import com.google.common.base.Joiner; -import com.google.common.base.Splitter; -import io.airlift.compress.lzo.LzoCodec; -import io.airlift.compress.lzo.LzopCodec; -import io.trino.hadoop.TextLineLengthLimitExceededException; -import io.trino.plugin.hive.HiveColumnHandle; -import io.trino.plugin.hive.HiveStorageFormat; -import io.trino.plugin.hive.avro.TrinoAvroSerDe; -import io.trino.spi.TrinoException; -import org.apache.hadoop.conf.Configuration; -import org.apache.hadoop.fs.Path; -import org.apache.hadoop.hive.serde2.AbstractSerDe; -import org.apache.hadoop.hive.serde2.Deserializer; -import org.apache.hadoop.hive.serde2.SerDeException; -import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspector; -import org.apache.hadoop.hive.serde2.objectinspector.StructObjectInspector; -import org.apache.hadoop.io.Writable; -import org.apache.hadoop.io.WritableComparable; -import org.apache.hadoop.mapred.FileSplit; -import org.apache.hadoop.mapred.InputFormat; -import org.apache.hadoop.mapred.JobConf; -import org.apache.hadoop.mapred.RecordReader; -import org.apache.hadoop.mapred.Reporter; -import org.apache.hadoop.util.ReflectionUtils; - -import java.io.IOException; -import java.util.List; -import java.util.Properties; - -import static com.google.common.base.MoreObjects.firstNonNull; -import static com.google.common.base.Preconditions.checkArgument; -import static com.google.common.collect.ImmutableList.toImmutableList; -import static com.google.common.collect.Lists.newArrayList; -import static io.trino.hdfs.ConfigurationUtils.copy; -import static io.trino.hdfs.ConfigurationUtils.toJobConf; -import static io.trino.hive.thrift.metastore.hive_metastoreConstants.FILE_INPUT_FORMAT; -import static io.trino.plugin.hive.HiveColumnHandle.ColumnType.REGULAR; -import static io.trino.plugin.hive.HiveErrorCode.HIVE_BAD_DATA; -import static io.trino.plugin.hive.HiveErrorCode.HIVE_CANNOT_OPEN_SPLIT; -import static io.trino.plugin.hive.HiveErrorCode.HIVE_INVALID_METADATA; -import static io.trino.plugin.hive.HiveErrorCode.HIVE_SERDE_NOT_FOUND; -import static io.trino.plugin.hive.HiveErrorCode.HIVE_UNSUPPORTED_FORMAT; -import static io.trino.plugin.hive.HiveStorageFormat.TEXTFILE; -import static io.trino.plugin.hive.util.HiveClassNames.AVRO_SERDE_CLASS; -import static io.trino.plugin.hive.util.HiveClassNames.LAZY_SIMPLE_SERDE_CLASS; -import static io.trino.plugin.hive.util.HiveClassNames.MAPRED_PARQUET_INPUT_FORMAT_CLASS; -import static io.trino.plugin.hive.util.HiveClassNames.SYMLINK_TEXT_INPUT_FORMAT_CLASS; -import static io.trino.plugin.hive.util.SerdeConstants.COLLECTION_DELIM; -import static org.apache.hadoop.hive.serde2.ColumnProjectionUtils.READ_ALL_COLUMNS; -import static org.apache.hadoop.hive.serde2.ColumnProjectionUtils.READ_COLUMN_IDS_CONF_STR; - -public final class HiveReaderUtil -{ - private HiveReaderUtil() {} - - public static RecordReader createRecordReader(Configuration configuration, Path path, long start, long length, Properties schema, List columns) - { - // determine which hive columns we will read - List readColumns = columns.stream() - .filter(column -> column.getColumnType() == REGULAR) - .collect(toImmutableList()); - - // Projected columns are not supported here - readColumns.forEach(readColumn -> checkArgument(readColumn.isBaseColumn(), "column %s is not a base column", readColumn.getName())); - - List readHiveColumnIndexes = readColumns.stream() - .map(HiveColumnHandle::getBaseHiveColumnIndex) - .collect(toImmutableList()); - - // Tell hive the columns we would like to read, this lets hive optimize reading column oriented files - configuration = copy(configuration); - setReadColumns(configuration, readHiveColumnIndexes); - - InputFormat inputFormat = getInputFormat(configuration, schema); - JobConf jobConf = toJobConf(configuration); - FileSplit fileSplit = new FileSplit(path, start, length, (String[]) null); - - // propagate serialization configuration to getRecordReader - schema.stringPropertyNames().stream() - .filter(name -> name.startsWith("serialization.")) - .forEach(name -> jobConf.set(name, schema.getProperty(name))); - - configureCompressionCodecs(jobConf); - - try { - @SuppressWarnings("unchecked") - RecordReader, ? extends Writable> recordReader = (RecordReader, ? extends Writable>) - inputFormat.getRecordReader(fileSplit, jobConf, Reporter.NULL); - - int headerCount = HiveUtil.getHeaderCount(schema); - // Only skip header rows when the split is at the beginning of the file - if (start == 0 && headerCount > 0) { - skipHeader(recordReader, headerCount); - } - - int footerCount = HiveUtil.getFooterCount(schema); - if (footerCount > 0) { - recordReader = new FooterAwareRecordReader<>(recordReader, footerCount, jobConf); - } - - return recordReader; - } - catch (IOException e) { - if (e instanceof TextLineLengthLimitExceededException) { - throw new TrinoException(HIVE_BAD_DATA, "Line too long in text file: " + path, e); - } - - throw new TrinoException(HIVE_CANNOT_OPEN_SPLIT, String.format("Error opening Hive split %s (offset=%s, length=%s) using %s: %s", - path, - start, - length, - HiveUtil.getInputFormatName(schema).orElse(null), - firstNonNull(e.getMessage(), e.getClass().getName())), - e); - } - } - - private static void skipHeader(RecordReader reader, int headerCount) - throws IOException - { - K key = reader.createKey(); - V value = reader.createValue(); - - while (headerCount > 0) { - if (!reader.next(key, value)) { - return; - } - headerCount--; - } - } - - private static void setReadColumns(Configuration configuration, List readHiveColumnIndexes) - { - configuration.set(READ_COLUMN_IDS_CONF_STR, Joiner.on(',').join(readHiveColumnIndexes)); - configuration.setBoolean(READ_ALL_COLUMNS, false); - } - - private static void configureCompressionCodecs(JobConf jobConf) - { - // add Airlift LZO and LZOP to head of codecs list to not override existing entries - List codecs = newArrayList(Splitter.on(",").trimResults().omitEmptyStrings().split(jobConf.get("io.compression.codecs", ""))); - if (!codecs.contains(LzoCodec.class.getName())) { - codecs.add(0, LzoCodec.class.getName()); - } - if (!codecs.contains(LzopCodec.class.getName())) { - codecs.add(0, LzopCodec.class.getName()); - } - jobConf.set("io.compression.codecs", String.join(",", codecs)); - } - - public static InputFormat getInputFormat(Configuration configuration, Properties schema) - { - String inputFormatName = HiveUtil.getInputFormatName(schema).orElseThrow(() -> - new TrinoException(HIVE_INVALID_METADATA, "Table or partition is missing Hive input format property: " + FILE_INPUT_FORMAT)); - try { - JobConf jobConf = toJobConf(configuration); - configureCompressionCodecs(jobConf); - - Class> inputFormatClass = getInputFormatClass(jobConf, inputFormatName); - if (inputFormatClass.getName().equals(SYMLINK_TEXT_INPUT_FORMAT_CLASS)) { - String serde = HiveUtil.getDeserializerClassName(schema); - // LazySimpleSerDe is used by TEXTFILE and SEQUENCEFILE. Default to TEXTFILE - // per Hive spec (https://hive.apache.org/javadocs/r2.1.1/api/org/apache/hadoop/hive/ql/io/SymlinkTextInputFormat.html) - if (serde.equals(TEXTFILE.getSerde())) { - inputFormatClass = getInputFormatClass(jobConf, TEXTFILE.getInputFormat()); - return ReflectionUtils.newInstance(inputFormatClass, jobConf); - } - for (HiveStorageFormat format : HiveStorageFormat.values()) { - if (serde.equals(format.getSerde())) { - inputFormatClass = getInputFormatClass(jobConf, format.getInputFormat()); - return ReflectionUtils.newInstance(inputFormatClass, jobConf); - } - } - throw new TrinoException(HIVE_UNSUPPORTED_FORMAT, "Unknown SerDe for SymlinkTextInputFormat: " + serde); - } - - return ReflectionUtils.newInstance(inputFormatClass, jobConf); - } - catch (ClassNotFoundException | RuntimeException e) { - throw new TrinoException(HIVE_UNSUPPORTED_FORMAT, "Unable to create input format " + inputFormatName, e); - } - } - - @SuppressWarnings("unchecked") - private static Class> getInputFormatClass(JobConf conf, String inputFormatName) - throws ClassNotFoundException - { - // legacy names for Parquet - if ("parquet.hive.DeprecatedParquetInputFormat".equals(inputFormatName) || - "parquet.hive.MapredParquetInputFormat".equals(inputFormatName)) { - inputFormatName = MAPRED_PARQUET_INPUT_FORMAT_CLASS; - } - - Class clazz = conf.getClassByName(inputFormatName); - return (Class>) clazz.asSubclass(InputFormat.class); - } - - public static StructObjectInspector getTableObjectInspector(Deserializer deserializer) - { - try { - ObjectInspector inspector = deserializer.getObjectInspector(); - checkArgument(inspector.getCategory() == ObjectInspector.Category.STRUCT, "expected STRUCT: %s", inspector.getCategory()); - return (StructObjectInspector) inspector; - } - catch (SerDeException e) { - throw new RuntimeException(e); - } - } - - public static Deserializer getDeserializer(Configuration configuration, Properties schema) - { - String name = HiveUtil.getDeserializerClassName(schema); - - // for collection delimiter, Hive 1.x, 2.x uses "colelction.delim" but Hive 3.x uses "collection.delim" - // see also https://issues.apache.org/jira/browse/HIVE-16922 - if (name.equals(LAZY_SIMPLE_SERDE_CLASS)) { - if (schema.containsKey("colelction.delim") && !schema.containsKey(COLLECTION_DELIM)) { - schema.setProperty(COLLECTION_DELIM, schema.getProperty("colelction.delim")); - } - } - - Deserializer deserializer = createDeserializer(getDeserializerClass(name)); - initializeDeserializer(configuration, deserializer, schema); - return deserializer; - } - - private static Class getDeserializerClass(String name) - { - if (AVRO_SERDE_CLASS.equals(name)) { - return TrinoAvroSerDe.class; - } - - try { - return Class.forName(name).asSubclass(Deserializer.class); - } - catch (ClassNotFoundException e) { - throw new TrinoException(HIVE_SERDE_NOT_FOUND, "deserializer does not exist: " + name); - } - catch (ClassCastException e) { - throw new RuntimeException("invalid deserializer class: " + name); - } - } - - private static Deserializer createDeserializer(Class clazz) - { - try { - return clazz.getConstructor().newInstance(); - } - catch (ReflectiveOperationException e) { - throw new RuntimeException("error creating deserializer: " + clazz.getName(), e); - } - } - - private static void initializeDeserializer(Configuration configuration, Deserializer deserializer, Properties schema) - { - try { - configuration = copy(configuration); // Some SerDes (e.g. Avro) modify passed configuration - deserializer.initialize(configuration, schema); - validate(deserializer); - } - catch (SerDeException | RuntimeException e) { - throw new RuntimeException("error initializing deserializer: " + deserializer.getClass().getName(), e); - } - } - - private static void validate(Deserializer deserializer) - { - if (deserializer instanceof AbstractSerDe && !((AbstractSerDe) deserializer).getConfigurationErrors().isEmpty()) { - throw new RuntimeException("There are configuration errors: " + ((AbstractSerDe) deserializer).getConfigurationErrors()); - } - } -} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/util/HiveTypeTranslator.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/util/HiveTypeTranslator.java index 574b26a76be3..3b0c3f440a8e 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/util/HiveTypeTranslator.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/util/HiveTypeTranslator.java @@ -17,6 +17,7 @@ import com.google.common.collect.Streams; import io.trino.plugin.hive.HiveErrorCode; import io.trino.plugin.hive.HiveTimestampPrecision; +import io.trino.plugin.hive.HiveType; import io.trino.plugin.hive.type.CharTypeInfo; import io.trino.plugin.hive.type.DecimalTypeInfo; import io.trino.plugin.hive.type.ListTypeInfo; @@ -91,6 +92,11 @@ public final class HiveTypeTranslator { private HiveTypeTranslator() {} + public static HiveType toHiveType(Type type) + { + return HiveType.fromTypeInfo(toTypeInfo(type)); + } + public static TypeInfo toTypeInfo(Type type) { requireNonNull(type, "type is null"); diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/util/HiveTypeUtil.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/util/HiveTypeUtil.java new file mode 100644 index 000000000000..1c52cac6c422 --- /dev/null +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/util/HiveTypeUtil.java @@ -0,0 +1,182 @@ +/* + * 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.hive.util; + +import com.google.common.collect.ImmutableList; +import io.trino.plugin.hive.HiveTimestampPrecision; +import io.trino.plugin.hive.HiveType; +import io.trino.plugin.hive.metastore.StorageFormat; +import io.trino.plugin.hive.type.ListTypeInfo; +import io.trino.plugin.hive.type.MapTypeInfo; +import io.trino.plugin.hive.type.PrimitiveCategory; +import io.trino.plugin.hive.type.PrimitiveTypeInfo; +import io.trino.plugin.hive.type.StructTypeInfo; +import io.trino.plugin.hive.type.TypeInfo; +import io.trino.plugin.hive.type.UnionTypeInfo; +import io.trino.spi.type.Type; +import io.trino.spi.type.TypeManager; +import io.trino.spi.type.TypeSignature; + +import java.util.List; +import java.util.Optional; + +import static com.google.common.base.Preconditions.checkArgument; +import static com.google.common.base.Strings.lenientFormat; +import static io.trino.hive.formats.UnionToRowCoercionUtils.UNION_FIELD_FIELD_PREFIX; +import static io.trino.hive.formats.UnionToRowCoercionUtils.UNION_FIELD_TAG_NAME; +import static io.trino.hive.formats.UnionToRowCoercionUtils.UNION_FIELD_TAG_TYPE; +import static io.trino.plugin.hive.HiveStorageFormat.AVRO; +import static io.trino.plugin.hive.HiveStorageFormat.ORC; +import static io.trino.plugin.hive.HiveTimestampPrecision.DEFAULT_PRECISION; +import static io.trino.plugin.hive.util.HiveTypeTranslator.toHiveType; +import static io.trino.plugin.hive.util.HiveTypeTranslator.toTypeSignature; + +public final class HiveTypeUtil +{ + private HiveTypeUtil() {} + + /** + * @deprecated Prefer {@link #getTypeSignature(HiveType, HiveTimestampPrecision)}. + */ + @Deprecated + public static TypeSignature getTypeSignature(HiveType type) + { + return getTypeSignature(type, DEFAULT_PRECISION); + } + + public static TypeSignature getTypeSignature(HiveType type, HiveTimestampPrecision timestampPrecision) + { + return toTypeSignature(type.getTypeInfo(), timestampPrecision); + } + + public static Type getType(HiveType type, TypeManager typeManager, HiveTimestampPrecision timestampPrecision) + { + return typeManager.getType(getTypeSignature(type, timestampPrecision)); + } + + public static boolean typeSupported(TypeInfo typeInfo, StorageFormat storageFormat) + { + return switch (typeInfo.getCategory()) { + case PRIMITIVE -> typeSupported(((PrimitiveTypeInfo) typeInfo).getPrimitiveCategory()); + case MAP -> typeSupported(((MapTypeInfo) typeInfo).getMapKeyTypeInfo(), storageFormat) && + typeSupported(((MapTypeInfo) typeInfo).getMapValueTypeInfo(), storageFormat); + case LIST -> typeSupported(((ListTypeInfo) typeInfo).getListElementTypeInfo(), storageFormat); + case STRUCT -> ((StructTypeInfo) typeInfo).getAllStructFieldTypeInfos().stream().allMatch(fieldTypeInfo -> typeSupported(fieldTypeInfo, storageFormat)); + case UNION -> + // This feature (reading union types as structs) has only been verified against Avro and ORC tables. Here's a discussion: + // 1. Avro tables are supported and verified. + // 2. ORC tables are supported and verified. + // 3. The Parquet format doesn't support union types itself so there's no need to add support for it in Trino. + // 4. TODO: RCFile tables are not supported yet. + // 5. TODO: The support for Avro is done in SerDeUtils so it's possible that formats other than Avro are also supported. But verification is needed. + storageFormat.getSerde().equalsIgnoreCase(AVRO.getSerde()) || + storageFormat.getSerde().equalsIgnoreCase(ORC.getSerde()) || + ((UnionTypeInfo) typeInfo).getAllUnionObjectTypeInfos().stream().allMatch(fieldTypeInfo -> typeSupported(fieldTypeInfo, storageFormat)); + }; + } + + private static boolean typeSupported(PrimitiveCategory category) + { + return switch (category) { + case BOOLEAN, + BYTE, + SHORT, + INT, + LONG, + FLOAT, + DOUBLE, + STRING, + VARCHAR, + CHAR, + DATE, + TIMESTAMP, + TIMESTAMPLOCALTZ, + BINARY, + DECIMAL -> true; + case INTERVAL_YEAR_MONTH, + INTERVAL_DAY_TIME, + VOID, + UNKNOWN -> false; + }; + } + + public static Optional getHiveTypeForDereferences(HiveType hiveType, List dereferences) + { + TypeInfo typeInfo = hiveType.getTypeInfo(); + for (int fieldIndex : dereferences) { + if (typeInfo instanceof StructTypeInfo structTypeInfo) { + try { + typeInfo = structTypeInfo.getAllStructFieldTypeInfos().get(fieldIndex); + } + catch (RuntimeException e) { + // return empty when failed to dereference, this could happen when partition and table schema mismatch + return Optional.empty(); + } + } + else if (typeInfo instanceof UnionTypeInfo unionTypeInfo) { + try { + if (fieldIndex == 0) { + // union's tag field, defined in {@link io.trino.hive.formats.UnionToRowCoercionUtils} + return Optional.of(toHiveType(UNION_FIELD_TAG_TYPE)); + } + typeInfo = unionTypeInfo.getAllUnionObjectTypeInfos().get(fieldIndex - 1); + } + catch (RuntimeException e) { + // return empty when failed to dereference, this could happen when partition and table schema mismatch + return Optional.empty(); + } + } + else { + throw new IllegalArgumentException(lenientFormat("typeInfo: %s should be struct or union type", typeInfo)); + } + } + return Optional.of(HiveType.fromTypeInfo(typeInfo)); + } + + public static List getHiveDereferenceNames(HiveType hiveType, List dereferences) + { + ImmutableList.Builder dereferenceNames = ImmutableList.builder(); + TypeInfo typeInfo = hiveType.getTypeInfo(); + for (int i = 0; i < dereferences.size(); i++) { + int fieldIndex = dereferences.get(i); + checkArgument(fieldIndex >= 0, "fieldIndex cannot be negative"); + + if (typeInfo instanceof StructTypeInfo structTypeInfo) { + checkArgument(fieldIndex < structTypeInfo.getAllStructFieldNames().size(), + "fieldIndex should be less than the number of fields in the struct"); + + String fieldName = structTypeInfo.getAllStructFieldNames().get(fieldIndex); + dereferenceNames.add(fieldName); + typeInfo = structTypeInfo.getAllStructFieldTypeInfos().get(fieldIndex); + } + else if (typeInfo instanceof UnionTypeInfo unionTypeInfo) { + checkArgument((fieldIndex - 1) < unionTypeInfo.getAllUnionObjectTypeInfos().size(), + "fieldIndex should be less than the number of fields in the union plus tag field"); + + if (fieldIndex == 0) { + checkArgument(i == (dereferences.size() - 1), "Union's tag field should not have more subfields"); + dereferenceNames.add(UNION_FIELD_TAG_NAME); + break; + } + typeInfo = unionTypeInfo.getAllUnionObjectTypeInfos().get(fieldIndex - 1); + dereferenceNames.add(UNION_FIELD_FIELD_PREFIX + (fieldIndex - 1)); + } + else { + throw new IllegalArgumentException(lenientFormat("typeInfo: %s should be struct or union type", typeInfo)); + } + } + + return dereferenceNames.build(); + } +} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/util/HiveUtil.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/util/HiveUtil.java index e0789bc9bf38..7c3d5649e38b 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/util/HiveUtil.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/util/HiveUtil.java @@ -28,7 +28,6 @@ import io.trino.plugin.hive.HivePartitionKey; import io.trino.plugin.hive.HiveTimestampPrecision; import io.trino.plugin.hive.HiveType; -import io.trino.plugin.hive.aws.athena.PartitionProjectionService; import io.trino.plugin.hive.metastore.Column; import io.trino.plugin.hive.metastore.SortingColumn; import io.trino.plugin.hive.metastore.Table; @@ -64,7 +63,6 @@ import java.util.Map; import java.util.Optional; import java.util.OptionalInt; -import java.util.Properties; import java.util.function.Function; import static com.google.common.base.Strings.isNullOrEmpty; @@ -97,6 +95,7 @@ import static io.trino.plugin.hive.HiveType.toHiveTypes; import static io.trino.plugin.hive.metastore.SortingColumn.Order.ASCENDING; import static io.trino.plugin.hive.metastore.SortingColumn.Order.DESCENDING; +import static io.trino.plugin.hive.projection.PartitionProjectionProperties.getPartitionProjectionTrinoColumnProperties; import static io.trino.plugin.hive.util.HiveBucketing.isSupportedBucketing; import static io.trino.plugin.hive.util.HiveClassNames.HUDI_INPUT_FORMAT; import static io.trino.plugin.hive.util.HiveClassNames.HUDI_PARQUET_INPUT_FORMAT; @@ -132,7 +131,6 @@ import static java.math.RoundingMode.UNNECESSARY; import static java.nio.charset.StandardCharsets.UTF_8; import static java.util.Locale.ENGLISH; -import static java.util.Objects.requireNonNull; public final class HiveUtil { @@ -182,9 +180,16 @@ private HiveUtil() { } - public static Optional getInputFormatName(Properties schema) + public static Optional getInputFormatName(Map schema) { - return Optional.ofNullable(schema.getProperty(FILE_INPUT_FORMAT)); + return Optional.ofNullable(schema.get(FILE_INPUT_FORMAT)); + } + + public static String getSerializationLibraryName(Map schema) + { + String name = schema.get(SERIALIZATION_LIB); + checkCondition(name != null, HIVE_INVALID_METADATA, "Table or partition is missing Hive deserializer property: %s", SERIALIZATION_LIB); + return name; } private static long parseHiveDate(String value) @@ -201,9 +206,9 @@ public static long parseHiveTimestamp(String value) return HIVE_TIMESTAMP_PARSER.parseMillis(value) * MICROSECONDS_PER_MILLISECOND; } - public static String getDeserializerClassName(Properties schema) + public static String getDeserializerClassName(Map schema) { - String name = schema.getProperty(SERIALIZATION_LIB); + String name = schema.get(SERIALIZATION_LIB); checkCondition(name != null, HIVE_INVALID_METADATA, "Table or partition is missing Hive deserializer property: %s", SERIALIZATION_LIB); return name; } @@ -701,19 +706,19 @@ public static List extractStructFieldTypes(HiveType hiveType) .collect(toImmutableList()); } - public static int getHeaderCount(Properties schema) + public static int getHeaderCount(Map schema) { return getPositiveIntegerValue(schema, SKIP_HEADER_COUNT_KEY, "0"); } - public static int getFooterCount(Properties schema) + public static int getFooterCount(Map schema) { return getPositiveIntegerValue(schema, SKIP_FOOTER_COUNT_KEY, "0"); } - private static int getPositiveIntegerValue(Properties schema, String key, String defaultValue) + private static int getPositiveIntegerValue(Map schema, String key, String defaultValue) { - String value = schema.getProperty(key, defaultValue); + String value = schema.getOrDefault(key, defaultValue); try { int intValue = parseInt(value); if (intValue < 0) { @@ -726,30 +731,33 @@ private static int getPositiveIntegerValue(Properties schema, String key, String } } - public static List getColumnNames(Properties schema) + public static List getColumnNames(Map schema) { - return COLUMN_NAMES_SPLITTER.splitToList(schema.getProperty(LIST_COLUMNS, "")); + return COLUMN_NAMES_SPLITTER.splitToList(schema.getOrDefault(LIST_COLUMNS, "")); } - public static List getColumnTypes(Properties schema) + public static List getColumnTypes(Map schema) { - return toHiveTypes(schema.getProperty(LIST_COLUMN_TYPES, "")); + return toHiveTypes(schema.getOrDefault(LIST_COLUMN_TYPES, "")); } - public static OrcWriterOptions getOrcWriterOptions(Properties schema, OrcWriterOptions orcWriterOptions) + public static OrcWriterOptions getOrcWriterOptions(Map schema, OrcWriterOptions orcWriterOptions) { if (schema.containsKey(ORC_BLOOM_FILTER_COLUMNS_KEY)) { if (!schema.containsKey(ORC_BLOOM_FILTER_FPP_KEY)) { throw new TrinoException(HIVE_INVALID_METADATA, "FPP for bloom filter is missing"); } try { - double fpp = parseDouble(schema.getProperty(ORC_BLOOM_FILTER_FPP_KEY)); + // use default fpp DEFAULT_BLOOM_FILTER_FPP if fpp key does not exist in table metadata + double fpp = schema.containsKey(ORC_BLOOM_FILTER_FPP_KEY) + ? parseDouble(schema.get(ORC_BLOOM_FILTER_FPP_KEY)) + : orcWriterOptions.getBloomFilterFpp(); return orcWriterOptions - .withBloomFilterColumns(ImmutableSet.copyOf(COLUMN_NAMES_SPLITTER.splitToList(schema.getProperty(ORC_BLOOM_FILTER_COLUMNS_KEY)))) + .withBloomFilterColumns(ImmutableSet.copyOf(COLUMN_NAMES_SPLITTER.splitToList(schema.get(ORC_BLOOM_FILTER_COLUMNS_KEY)))) .withBloomFilterFpp(fpp); } catch (NumberFormatException e) { - throw new TrinoException(HIVE_UNSUPPORTED_FORMAT, format("Invalid value for %s property: %s", ORC_BLOOM_FILTER_FPP, schema.getProperty(ORC_BLOOM_FILTER_FPP_KEY))); + throw new TrinoException(HIVE_UNSUPPORTED_FORMAT, format("Invalid value for %s property: %s", ORC_BLOOM_FILTER_FPP, schema.get(ORC_BLOOM_FILTER_FPP_KEY))); } } return orcWriterOptions; @@ -810,10 +818,11 @@ public static boolean isIcebergTable(Map tableParameters) public static boolean isHudiTable(Table table) { - requireNonNull(table, "table is null"); - @Nullable - String inputFormat = table.getStorage().getStorageFormat().getInputFormatNullable(); + return isHudiTable(table.getStorage().getStorageFormat().getInputFormatNullable()); + } + public static boolean isHudiTable(String inputFormat) + { return HUDI_PARQUET_INPUT_FORMAT.equals(inputFormat) || HUDI_PARQUET_REALTIME_INPUT_FORMAT.equals(inputFormat) || HUDI_INPUT_FORMAT.equals(inputFormat) || @@ -856,7 +865,7 @@ public static Function columnMetadataGetter(Ta .setComment(handle.isHidden() ? Optional.empty() : columnComment.get(handle.getName())) .setExtraInfo(Optional.ofNullable(columnExtraInfo(handle.isPartitionKey()))) .setHidden(handle.isHidden()) - .setProperties(PartitionProjectionService.getPartitionProjectionTrinoColumnProperties(table, handle.getName())) + .setProperties(getPartitionProjectionTrinoColumnProperties(table, handle.getName())) .build(); } diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/util/HiveWriteUtils.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/util/HiveWriteUtils.java index c05d7091470b..55704d325af7 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/util/HiveWriteUtils.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/util/HiveWriteUtils.java @@ -16,10 +16,7 @@ import com.google.common.base.CharMatcher; import com.google.common.collect.ImmutableList; import io.trino.filesystem.Location; -import io.trino.hdfs.HdfsContext; -import io.trino.hdfs.HdfsEnvironment; -import io.trino.hdfs.rubix.CachingTrinoS3FileSystem; -import io.trino.hdfs.s3.TrinoS3FileSystem; +import io.trino.filesystem.TrinoFileSystem; import io.trino.plugin.hive.HiveReadOnlyException; import io.trino.plugin.hive.HiveTimestampPrecision; import io.trino.plugin.hive.HiveType; @@ -44,6 +41,7 @@ import io.trino.spi.connector.ConnectorSession; import io.trino.spi.connector.SchemaNotFoundException; import io.trino.spi.connector.SchemaTableName; +import io.trino.spi.security.ConnectorIdentity; import io.trino.spi.type.ArrayType; import io.trino.spi.type.CharType; import io.trino.spi.type.DecimalType; @@ -56,11 +54,7 @@ import io.trino.spi.type.Type; import io.trino.spi.type.VarcharType; import org.apache.hadoop.conf.Configuration; -import org.apache.hadoop.fs.FileStatus; -import org.apache.hadoop.fs.FileSystem; import org.apache.hadoop.fs.Path; -import org.apache.hadoop.fs.viewfs.ViewFileSystem; -import org.apache.hadoop.hdfs.DistributedFileSystem; import org.apache.hadoop.hive.common.type.Date; import org.apache.hadoop.hive.common.type.HiveDecimal; import org.apache.hadoop.hive.common.type.Timestamp; @@ -77,7 +71,6 @@ import org.apache.hadoop.mapred.Reporter; import org.joda.time.DateTimeZone; -import java.io.FileNotFoundException; import java.io.IOException; import java.math.BigInteger; import java.util.ArrayList; @@ -91,8 +84,6 @@ import static com.google.common.base.Verify.verify; import static com.google.common.collect.ImmutableList.toImmutableList; import static com.google.common.io.BaseEncoding.base16; -import static io.trino.hdfs.FileSystemUtils.getRawFileSystem; -import static io.trino.hdfs.s3.HiveS3Module.EMR_FS_CLASS_NAME; import static io.trino.plugin.hive.HiveErrorCode.HIVE_DATABASE_LOCATION_ERROR; import static io.trino.plugin.hive.HiveErrorCode.HIVE_FILESYSTEM_ERROR; import static io.trino.plugin.hive.HiveErrorCode.HIVE_INVALID_PARTITION_VALUE; @@ -135,7 +126,6 @@ import static java.nio.charset.StandardCharsets.UTF_8; import static java.util.Collections.unmodifiableList; import static java.util.Collections.unmodifiableMap; -import static java.util.UUID.randomUUID; import static java.util.stream.Collectors.toList; import static org.apache.hadoop.hive.conf.HiveConf.ConfVars.COMPRESSRESULT; import static org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorFactory.getPrimitiveJavaObjectInspector; @@ -438,179 +428,51 @@ private static void checkWritable( } } - public static Location getTableDefaultLocation(HdfsContext context, SemiTransactionalHiveMetastore metastore, HdfsEnvironment hdfsEnvironment, String schemaName, String tableName) + public static Location getTableDefaultLocation(SemiTransactionalHiveMetastore metastore, TrinoFileSystem fileSystem, String schemaName, String tableName) { Database database = metastore.getDatabase(schemaName) .orElseThrow(() -> new SchemaNotFoundException(schemaName)); - return getTableDefaultLocation(database, context, hdfsEnvironment, schemaName, tableName); + return getTableDefaultLocation(database, fileSystem, schemaName, tableName); } - public static Location getTableDefaultLocation(Database database, HdfsContext context, HdfsEnvironment hdfsEnvironment, String schemaName, String tableName) + public static Location getTableDefaultLocation(Database database, TrinoFileSystem fileSystem, String schemaName, String tableName) { - String location = database.getLocation() + Location location = database.getLocation().map(Location::of) .orElseThrow(() -> new TrinoException(HIVE_DATABASE_LOCATION_ERROR, format("Database '%s' location is not set", schemaName))); - Path databasePath = new Path(location); - if (!isS3FileSystem(context, hdfsEnvironment, databasePath)) { - if (!pathExists(context, hdfsEnvironment, databasePath)) { - throw new TrinoException(HIVE_DATABASE_LOCATION_ERROR, format("Database '%s' location does not exist: %s", schemaName, databasePath)); - } - if (!isDirectory(context, hdfsEnvironment, databasePath)) { - throw new TrinoException(HIVE_DATABASE_LOCATION_ERROR, format("Database '%s' location is not a directory: %s", schemaName, databasePath)); - } - } - - // Note: this results in `databaseLocation` being a "normalized location", e.g. not containing double slashes. - // TODO (https://github.com/trinodb/trino/issues/17803): We need to use normalized location until all relevant Hive connector components are migrated off Hadoop Path. - Location databaseLocation = Location.of(databasePath.toString()); - return databaseLocation.appendPath(escapeTableName(tableName)); - } - - public static boolean pathExists(HdfsContext context, HdfsEnvironment hdfsEnvironment, Path path) - { - try { - return hdfsEnvironment.getFileSystem(context, path).exists(path); - } - catch (IOException e) { - throw new TrinoException(HIVE_FILESYSTEM_ERROR, "Failed checking path: " + path, e); - } - } - - public static boolean isS3FileSystem(HdfsContext context, HdfsEnvironment hdfsEnvironment, Path path) - { - try { - FileSystem fileSystem = getRawFileSystem(hdfsEnvironment.getFileSystem(context, path)); - return fileSystem instanceof TrinoS3FileSystem || fileSystem.getClass().getName().equals(EMR_FS_CLASS_NAME) || fileSystem instanceof CachingTrinoS3FileSystem; - } - catch (IOException e) { - throw new TrinoException(HIVE_FILESYSTEM_ERROR, "Failed checking path: " + path, e); + if (!directoryExists(fileSystem, location).orElse(true)) { + throw new TrinoException(HIVE_DATABASE_LOCATION_ERROR, format("Database '%s' location does not exist: %s", schemaName, location)); } - } - public static boolean isViewFileSystem(HdfsContext context, HdfsEnvironment hdfsEnvironment, Path path) - { - try { - return getRawFileSystem(hdfsEnvironment.getFileSystem(context, path)) instanceof ViewFileSystem; - } - catch (IOException e) { - throw new TrinoException(HIVE_FILESYSTEM_ERROR, "Failed checking path: " + path, e); - } + return location.appendPath(escapeTableName(tableName)); } - private static boolean isDirectory(HdfsContext context, HdfsEnvironment hdfsEnvironment, Path path) + public static Optional directoryExists(TrinoFileSystem fileSystem, Location path) { try { - return hdfsEnvironment.getFileSystem(context, path).isDirectory(path); + return fileSystem.directoryExists(path); } catch (IOException e) { throw new TrinoException(HIVE_FILESYSTEM_ERROR, "Failed checking path: " + path, e); } } - public static boolean isHdfsEncrypted(HdfsContext context, HdfsEnvironment hdfsEnvironment, Path path) - { - try { - FileSystem fileSystem = getRawFileSystem(hdfsEnvironment.getFileSystem(context, path)); - if (fileSystem instanceof DistributedFileSystem) { - return ((DistributedFileSystem) fileSystem).getEZForPath(path) != null; - } - return false; - } - catch (IOException e) { - throw new TrinoException(HIVE_FILESYSTEM_ERROR, "Failed checking encryption status for path: " + path, e); - } - } - public static boolean isFileCreatedByQuery(String fileName, String queryId) { return fileName.startsWith(queryId) || fileName.endsWith(queryId); } - public static Location createTemporaryPath(HdfsContext context, HdfsEnvironment hdfsEnvironment, Path targetPath, String temporaryStagingDirectoryPath) + public static Optional createTemporaryPath(TrinoFileSystem fileSystem, ConnectorIdentity identity, Location targetPath, String temporaryStagingDirectoryPath) { - // use a per-user temporary directory to avoid permission problems - String temporaryPrefix = temporaryStagingDirectoryPath.replace("${USER}", context.getIdentity().getUser()); + // interpolate the username into the temporary directory path to avoid permission problems + String temporaryPrefix = temporaryStagingDirectoryPath.replace("${USER}", identity.getUser()); - // use relative temporary directory on ViewFS - if (isViewFileSystem(context, hdfsEnvironment, targetPath)) { - temporaryPrefix = ".hive-staging"; - } - - // create a temporary directory on the same filesystem - Path temporaryRoot = new Path(targetPath, temporaryPrefix); - Path temporaryPath = new Path(temporaryRoot, randomUUID().toString()); - - createDirectory(context, hdfsEnvironment, temporaryPath); - - if (hdfsEnvironment.isNewFileInheritOwnership()) { - setDirectoryOwner(context, hdfsEnvironment, temporaryPath, targetPath); - } - - return Location.of(temporaryPath.toString()); - } - - private static void setDirectoryOwner(HdfsContext context, HdfsEnvironment hdfsEnvironment, Path path, Path targetPath) - { try { - FileSystem fileSystem = hdfsEnvironment.getFileSystem(context, path); - FileStatus fileStatus; - if (!fileSystem.exists(targetPath)) { - // For new table - Path parent = targetPath.getParent(); - if (!fileSystem.exists(parent)) { - return; - } - fileStatus = fileSystem.getFileStatus(parent); - } - else { - // For existing table - fileStatus = fileSystem.getFileStatus(targetPath); - } - String owner = fileStatus.getOwner(); - String group = fileStatus.getGroup(); - fileSystem.setOwner(path, owner, group); + return fileSystem.createTemporaryDirectory(targetPath, temporaryPrefix, ".hive-staging"); } catch (IOException e) { - throw new TrinoException(HIVE_FILESYSTEM_ERROR, format("Failed to set owner on %s based on %s", path, targetPath), e); - } - } - - public static void createDirectory(HdfsContext context, HdfsEnvironment hdfsEnvironment, Path path) - { - try { - if (!hdfsEnvironment.getFileSystem(context, path).mkdirs(path, hdfsEnvironment.getNewDirectoryPermissions().orElse(null))) { - throw new IOException("mkdirs returned false"); - } - } - catch (IOException e) { - throw new TrinoException(HIVE_FILESYSTEM_ERROR, "Failed to create directory: " + path, e); - } - - if (hdfsEnvironment.getNewDirectoryPermissions().isPresent()) { - // explicitly set permission since the default umask overrides it on creation - try { - hdfsEnvironment.getFileSystem(context, path).setPermission(path, hdfsEnvironment.getNewDirectoryPermissions().get()); - } - catch (IOException e) { - throw new TrinoException(HIVE_FILESYSTEM_ERROR, "Failed to set permission on directory: " + path, e); - } - } - } - - public static void checkedDelete(FileSystem fileSystem, Path file, boolean recursive) - throws IOException - { - try { - if (!fileSystem.delete(file, recursive)) { - if (fileSystem.exists(file)) { - // only throw exception if file still exists - throw new IOException("Failed to delete " + file); - } - } - } - catch (FileNotFoundException ignored) { - // ok + throw new TrinoException(HIVE_FILESYSTEM_ERROR, e); } } diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/util/InternalHiveSplitFactory.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/util/InternalHiveSplitFactory.java index 3922f6b3983c..d63b24aabeec 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/util/InternalHiveSplitFactory.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/util/InternalHiveSplitFactory.java @@ -14,6 +14,7 @@ package io.trino.plugin.hive.util; import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; import io.airlift.units.DataSize; import io.trino.plugin.hive.AcidInfo; import io.trino.plugin.hive.HiveColumnHandle; @@ -21,94 +22,90 @@ import io.trino.plugin.hive.HiveSplit; import io.trino.plugin.hive.HiveSplit.BucketConversion; import io.trino.plugin.hive.HiveStorageFormat; +import io.trino.plugin.hive.HiveTypeName; import io.trino.plugin.hive.InternalHiveSplit; import io.trino.plugin.hive.InternalHiveSplit.InternalHiveBlock; -import io.trino.plugin.hive.TableToPartitionMapping; +import io.trino.plugin.hive.Schema; import io.trino.plugin.hive.fs.BlockLocation; import io.trino.plugin.hive.fs.TrinoFileStatus; import io.trino.plugin.hive.orc.OrcPageSourceFactory; import io.trino.plugin.hive.parquet.ParquetPageSourceFactory; import io.trino.plugin.hive.rcfile.RcFilePageSourceFactory; -import io.trino.plugin.hive.s3select.S3SelectPushdown; import io.trino.spi.HostAddress; import io.trino.spi.predicate.Domain; import io.trino.spi.predicate.TupleDomain; -import org.apache.hadoop.fs.FileStatus; -import org.apache.hadoop.fs.FileSystem; -import org.apache.hadoop.mapred.FileSplit; -import java.io.IOException; import java.util.Collection; import java.util.List; import java.util.Map; import java.util.Optional; import java.util.OptionalInt; -import java.util.Properties; import java.util.function.BooleanSupplier; import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.collect.ImmutableList.toImmutableList; import static io.airlift.slice.Slices.utf8Slice; import static io.trino.plugin.hive.HiveColumnHandle.isPathColumnHandle; +import static io.trino.plugin.hive.util.AcidTables.isFullAcidTable; +import static io.trino.plugin.hive.util.HiveUtil.getSerializationLibraryName; import static java.util.Objects.requireNonNull; public class InternalHiveSplitFactory { - private final FileSystem fileSystem; private final String partitionName; private final HiveStorageFormat storageFormat; - private final Properties strippedSchema; + private final Schema strippedSchema; private final List partitionKeys; private final Optional pathDomain; - private final TableToPartitionMapping tableToPartitionMapping; + private final Map hiveColumnCoercions; private final BooleanSupplier partitionMatchSupplier; private final Optional bucketConversion; private final Optional bucketValidation; private final long minimumTargetSplitSizeInBytes; private final Optional maxSplitFileSize; private final boolean forceLocalScheduling; - private final boolean s3SelectPushdownEnabled; public InternalHiveSplitFactory( - FileSystem fileSystem, String partitionName, HiveStorageFormat storageFormat, - Properties schema, + Map schema, List partitionKeys, TupleDomain effectivePredicate, BooleanSupplier partitionMatchSupplier, - TableToPartitionMapping tableToPartitionMapping, + Map hiveColumnCoercions, Optional bucketConversion, Optional bucketValidation, DataSize minimumTargetSplitSize, boolean forceLocalScheduling, - boolean s3SelectPushdownEnabled, Optional maxSplitFileSize) { - this.fileSystem = requireNonNull(fileSystem, "fileSystem is null"); this.partitionName = requireNonNull(partitionName, "partitionName is null"); this.storageFormat = requireNonNull(storageFormat, "storageFormat is null"); this.strippedSchema = stripUnnecessaryProperties(requireNonNull(schema, "schema is null")); this.partitionKeys = requireNonNull(partitionKeys, "partitionKeys is null"); pathDomain = getPathDomain(requireNonNull(effectivePredicate, "effectivePredicate is null")); this.partitionMatchSupplier = requireNonNull(partitionMatchSupplier, "partitionMatchSupplier is null"); - this.tableToPartitionMapping = requireNonNull(tableToPartitionMapping, "tableToPartitionMapping is null"); + this.hiveColumnCoercions = ImmutableMap.copyOf(requireNonNull(hiveColumnCoercions, "hiveColumnCoercions is null")); this.bucketConversion = requireNonNull(bucketConversion, "bucketConversion is null"); this.bucketValidation = requireNonNull(bucketValidation, "bucketValidation is null"); this.forceLocalScheduling = forceLocalScheduling; - this.s3SelectPushdownEnabled = s3SelectPushdownEnabled; this.minimumTargetSplitSizeInBytes = minimumTargetSplitSize.toBytes(); this.maxSplitFileSize = requireNonNull(maxSplitFileSize, "maxSplitFileSize is null"); checkArgument(minimumTargetSplitSizeInBytes > 0, "minimumTargetSplitSize must be > 0, found: %s", minimumTargetSplitSize); } - private static Properties stripUnnecessaryProperties(Properties schema) + private static Schema stripUnnecessaryProperties(Map schema) { // Sending the full schema with every split is costly and can be avoided for formats supported natively - schema = OrcPageSourceFactory.stripUnnecessaryProperties(schema); - schema = ParquetPageSourceFactory.stripUnnecessaryProperties(schema); - schema = RcFilePageSourceFactory.stripUnnecessaryProperties(schema); - return schema; + String serializationLibraryName = getSerializationLibraryName(schema); + boolean isFullAcidTable = isFullAcidTable(schema); + Map serdeProperties = schema; + if (RcFilePageSourceFactory.stripUnnecessaryProperties(serializationLibraryName) + || OrcPageSourceFactory.stripUnnecessaryProperties(serializationLibraryName) + || ParquetPageSourceFactory.stripUnnecessaryProperties(serializationLibraryName)) { + serdeProperties = ImmutableMap.of(); + } + return new Schema(serializationLibraryName, isFullAcidTable, serdeProperties); } public String getPartitionName() @@ -134,23 +131,6 @@ public Optional createInternalHiveSplit(TrinoFileStatus statu acidInfo); } - public Optional createInternalHiveSplit(FileSplit split) - throws IOException - { - FileStatus file = fileSystem.getFileStatus(split.getPath()); - return createInternalHiveSplit( - split.getPath().toString(), - BlockLocation.fromHiveBlockLocations(fileSystem.getFileBlockLocations(file, split.getStart(), split.getLength())), - split.getStart(), - split.getLength(), - file.getLen(), - file.getModificationTime(), - OptionalInt.empty(), - OptionalInt.empty(), - false, - Optional.empty()); - } - private Optional createInternalHiveSplit( String path, List blockLocations, @@ -220,10 +200,9 @@ private Optional createInternalHiveSplit( tableBucketNumber, splittable, forceLocalScheduling && allBlocksHaveAddress(blocks), - tableToPartitionMapping, + hiveColumnCoercions, bucketConversion, bucketValidation, - s3SelectPushdownEnabled && S3SelectPushdown.isCompressionCodecSupported(strippedSchema, path), acidInfo, partitionMatchSupplier)); } diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/util/Statistics.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/util/Statistics.java index 549bfe07ff8a..723c2791b86e 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/util/Statistics.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/util/Statistics.java @@ -121,9 +121,9 @@ public static HiveColumnStatistics merge(HiveColumnStatistics first, HiveColumnS mergeDateStatistics(first.getDateStatistics(), second.getDateStatistics()), mergeBooleanStatistics(first.getBooleanStatistics(), second.getBooleanStatistics()), reduce(first.getMaxValueSizeInBytes(), second.getMaxValueSizeInBytes(), MAX, true), - reduce(first.getTotalSizeInBytes(), second.getTotalSizeInBytes(), ADD, true), + reduce(first.getAverageColumnLength(), second.getAverageColumnLength(), ADD, true), reduce(first.getNullsCount(), second.getNullsCount(), ADD, false), - reduce(first.getDistinctValuesCount(), second.getDistinctValuesCount(), MAX, false)); + reduce(first.getDistinctValuesWithNullCount(), second.getDistinctValuesWithNullCount(), MAX, false)); } private static Optional mergeIntegerStatistics(Optional first, Optional second) @@ -279,10 +279,10 @@ private static void setColumnStatisticsForEmptyPartition(Type columnType, HiveCo result.setMaxValueSizeInBytes(0); return; case TOTAL_SIZE_IN_BYTES: - result.setTotalSizeInBytes(0); + result.setAverageColumnLength(0); return; case NUMBER_OF_DISTINCT_VALUES: - result.setDistinctValuesCount(0); + result.setDistinctValuesWithNullCount(0); return; case NUMBER_OF_NON_NULL_VALUES: result.setNullsCount(0); @@ -331,7 +331,7 @@ public static Map, ComputedStatistics> createComputedStatisticsToPa .collect(toImmutableMap(statistics -> getPartitionValues(statistics, partitionColumns, partitionColumnTypes), Function.identity())); } - private static List getPartitionValues(ComputedStatistics statistics, List partitionColumns, List partitionColumnTypes) + public static List getPartitionValues(ComputedStatistics statistics, List partitionColumns, List partitionColumnTypes) { checkArgument(statistics.getGroupingColumns().equals(partitionColumns), "Unexpected grouping. Partition columns: %s. Grouping columns: %s", partitionColumns, statistics.getGroupingColumns()); @@ -382,7 +382,9 @@ static HiveColumnStatistics createHiveColumnStatistics( // TOTAL_VALUES_SIZE_IN_BYTES if (computedStatistics.containsKey(TOTAL_SIZE_IN_BYTES)) { - result.setTotalSizeInBytes(getIntegerValue(BIGINT, computedStatistics.get(TOTAL_SIZE_IN_BYTES))); + OptionalLong totalSizeInBytes = getIntegerValue(BIGINT, computedStatistics.get(TOTAL_SIZE_IN_BYTES)); + OptionalLong numNonNullValues = getIntegerValue(BIGINT, computedStatistics.get(NUMBER_OF_NON_NULL_VALUES)); + result.setAverageColumnLength(getAverageColumnLength(totalSizeInBytes, numNonNullValues)); } // NUMBER OF NULLS @@ -395,12 +397,8 @@ static HiveColumnStatistics createHiveColumnStatistics( // number of distinct value is estimated using HLL, and can be higher than the number of non null values long numberOfNonNullValues = BIGINT.getLong(computedStatistics.get(NUMBER_OF_NON_NULL_VALUES), 0); long numberOfDistinctValues = BIGINT.getLong(computedStatistics.get(NUMBER_OF_DISTINCT_VALUES), 0); - if (numberOfDistinctValues > numberOfNonNullValues) { - result.setDistinctValuesCount(numberOfNonNullValues); - } - else { - result.setDistinctValuesCount(numberOfDistinctValues); - } + // Hive expects NDV to be one greater when column has a null + result.setDistinctValuesWithNullCount(Math.min(numberOfDistinctValues, numberOfNonNullValues) + (rowCount > numberOfNonNullValues ? 1 : 0)); } // NUMBER OF FALSE, NUMBER OF TRUE @@ -487,4 +485,17 @@ public enum ReduceOperator MIN, MAX, } + + private static OptionalDouble getAverageColumnLength(OptionalLong totalSizeInBytes, OptionalLong numNonNullValues) + { + if (totalSizeInBytes.isEmpty() || numNonNullValues.isEmpty()) { + return OptionalDouble.empty(); + } + + long nonNullsCount = numNonNullValues.getAsLong(); + if (nonNullsCount <= 0) { + return OptionalDouble.empty(); + } + return OptionalDouble.of(((double) totalSizeInBytes.getAsLong()) / nonNullsCount); + } } diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/AbstractTestHive.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/AbstractTestHive.java deleted file mode 100644 index c90eed86a59c..000000000000 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/AbstractTestHive.java +++ /dev/null @@ -1,6436 +0,0 @@ -/* - * 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.hive; - -import com.google.common.collect.ImmutableList; -import com.google.common.collect.ImmutableMap; -import com.google.common.collect.ImmutableMultimap; -import com.google.common.collect.ImmutableSet; -import com.google.common.collect.Sets; -import com.google.common.net.HostAndPort; -import io.airlift.json.JsonCodec; -import io.airlift.log.Logger; -import io.airlift.slice.Slice; -import io.airlift.stats.CounterStat; -import io.airlift.units.DataSize; -import io.airlift.units.Duration; -import io.trino.filesystem.Location; -import io.trino.filesystem.TrinoFileSystem; -import io.trino.filesystem.hdfs.HdfsFileSystemFactory; -import io.trino.hdfs.HdfsConfig; -import io.trino.hdfs.HdfsContext; -import io.trino.hdfs.HdfsEnvironment; -import io.trino.hdfs.HdfsNamenodeStats; -import io.trino.hdfs.authentication.NoHdfsAuthentication; -import io.trino.operator.GroupByHashPageIndexerFactory; -import io.trino.plugin.base.CatalogName; -import io.trino.plugin.base.metrics.LongCount; -import io.trino.plugin.hive.LocationService.WriteInfo; -import io.trino.plugin.hive.aws.athena.PartitionProjectionService; -import io.trino.plugin.hive.fs.DirectoryLister; -import io.trino.plugin.hive.fs.RemoteIterator; -import io.trino.plugin.hive.fs.TransactionScopeCachingDirectoryListerFactory; -import io.trino.plugin.hive.fs.TrinoFileStatus; -import io.trino.plugin.hive.fs.TrinoFileStatusRemoteIterator; -import io.trino.plugin.hive.line.LinePageSource; -import io.trino.plugin.hive.metastore.Column; -import io.trino.plugin.hive.metastore.HiveColumnStatistics; -import io.trino.plugin.hive.metastore.HiveMetastore; -import io.trino.plugin.hive.metastore.HiveMetastoreFactory; -import io.trino.plugin.hive.metastore.HivePrincipal; -import io.trino.plugin.hive.metastore.HivePrivilegeInfo; -import io.trino.plugin.hive.metastore.HivePrivilegeInfo.HivePrivilege; -import io.trino.plugin.hive.metastore.Partition; -import io.trino.plugin.hive.metastore.PartitionWithStatistics; -import io.trino.plugin.hive.metastore.PrincipalPrivileges; -import io.trino.plugin.hive.metastore.SemiTransactionalHiveMetastore; -import io.trino.plugin.hive.metastore.SortingColumn; -import io.trino.plugin.hive.metastore.StorageFormat; -import io.trino.plugin.hive.metastore.Table; -import io.trino.plugin.hive.metastore.cache.CachingHiveMetastore; -import io.trino.plugin.hive.metastore.cache.CachingHiveMetastoreConfig; -import io.trino.plugin.hive.metastore.thrift.BridgingHiveMetastore; -import io.trino.plugin.hive.metastore.thrift.ThriftMetastoreConfig; -import io.trino.plugin.hive.orc.OrcPageSource; -import io.trino.plugin.hive.parquet.ParquetPageSource; -import io.trino.plugin.hive.rcfile.RcFilePageSource; -import io.trino.plugin.hive.security.SqlStandardAccessControlMetadata; -import io.trino.spi.Page; -import io.trino.spi.TrinoException; -import io.trino.spi.block.Block; -import io.trino.spi.connector.Assignment; -import io.trino.spi.connector.CatalogSchemaTableName; -import io.trino.spi.connector.ColumnHandle; -import io.trino.spi.connector.ColumnMetadata; -import io.trino.spi.connector.ConnectorBucketNodeMap; -import io.trino.spi.connector.ConnectorInsertTableHandle; -import io.trino.spi.connector.ConnectorMaterializedViewDefinition; -import io.trino.spi.connector.ConnectorMetadata; -import io.trino.spi.connector.ConnectorNodePartitioningProvider; -import io.trino.spi.connector.ConnectorOutputTableHandle; -import io.trino.spi.connector.ConnectorPageSink; -import io.trino.spi.connector.ConnectorPageSinkProvider; -import io.trino.spi.connector.ConnectorPageSource; -import io.trino.spi.connector.ConnectorPageSourceProvider; -import io.trino.spi.connector.ConnectorPartitioningHandle; -import io.trino.spi.connector.ConnectorSession; -import io.trino.spi.connector.ConnectorSplit; -import io.trino.spi.connector.ConnectorSplitManager; -import io.trino.spi.connector.ConnectorSplitSource; -import io.trino.spi.connector.ConnectorTableHandle; -import io.trino.spi.connector.ConnectorTableLayout; -import io.trino.spi.connector.ConnectorTableMetadata; -import io.trino.spi.connector.ConnectorTableProperties; -import io.trino.spi.connector.ConnectorTransactionHandle; -import io.trino.spi.connector.ConnectorViewDefinition; -import io.trino.spi.connector.ConnectorViewDefinition.ViewColumn; -import io.trino.spi.connector.Constraint; -import io.trino.spi.connector.ConstraintApplicationResult; -import io.trino.spi.connector.DiscretePredicates; -import io.trino.spi.connector.DynamicFilter; -import io.trino.spi.connector.ProjectionApplicationResult; -import io.trino.spi.connector.RecordCursor; -import io.trino.spi.connector.RecordPageSource; -import io.trino.spi.connector.SchemaTableName; -import io.trino.spi.connector.SchemaTablePrefix; -import io.trino.spi.connector.SortingProperty; -import io.trino.spi.connector.TableColumnsMetadata; -import io.trino.spi.connector.TableNotFoundException; -import io.trino.spi.connector.TableScanRedirectApplicationResult; -import io.trino.spi.connector.ViewNotFoundException; -import io.trino.spi.expression.ConnectorExpression; -import io.trino.spi.expression.FieldDereference; -import io.trino.spi.expression.Variable; -import io.trino.spi.metrics.Metrics; -import io.trino.spi.predicate.Domain; -import io.trino.spi.predicate.NullableValue; -import io.trino.spi.predicate.Range; -import io.trino.spi.predicate.TupleDomain; -import io.trino.spi.predicate.ValueSet; -import io.trino.spi.statistics.ColumnStatistics; -import io.trino.spi.statistics.TableStatistics; -import io.trino.spi.type.ArrayType; -import io.trino.spi.type.CharType; -import io.trino.spi.type.MapType; -import io.trino.spi.type.NamedTypeSignature; -import io.trino.spi.type.RowFieldName; -import io.trino.spi.type.RowType; -import io.trino.spi.type.SqlDate; -import io.trino.spi.type.SqlTimestamp; -import io.trino.spi.type.SqlTimestampWithTimeZone; -import io.trino.spi.type.SqlVarbinary; -import io.trino.spi.type.TestingTypeManager; -import io.trino.spi.type.Type; -import io.trino.spi.type.TypeId; -import io.trino.spi.type.TypeOperators; -import io.trino.spi.type.VarcharType; -import io.trino.sql.gen.JoinCompiler; -import io.trino.testing.MaterializedResult; -import io.trino.testing.MaterializedRow; -import io.trino.testing.TestingConnectorSession; -import io.trino.testing.TestingNodeManager; -import io.trino.type.BlockTypeOperators; -import org.apache.hadoop.fs.FileStatus; -import org.apache.hadoop.fs.FileSystem; -import org.apache.hadoop.fs.Path; -import org.apache.hadoop.hive.metastore.TableType; -import org.assertj.core.api.InstanceOfAssertFactories; -import org.joda.time.DateTime; -import org.testng.annotations.AfterClass; -import org.testng.annotations.BeforeClass; -import org.testng.annotations.Test; - -import java.io.IOException; -import java.io.OutputStream; -import java.math.BigDecimal; -import java.time.LocalDate; -import java.time.LocalDateTime; -import java.util.Arrays; -import java.util.Collection; -import java.util.HashMap; -import java.util.HashSet; -import java.util.List; -import java.util.Map; -import java.util.Optional; -import java.util.OptionalDouble; -import java.util.OptionalInt; -import java.util.OptionalLong; -import java.util.Set; -import java.util.UUID; -import java.util.concurrent.ExecutorService; -import java.util.concurrent.ScheduledExecutorService; -import java.util.concurrent.ThreadLocalRandom; -import java.util.concurrent.atomic.AtomicInteger; -import java.util.concurrent.atomic.AtomicReference; -import java.util.function.Function; -import java.util.function.Predicate; -import java.util.stream.IntStream; -import java.util.stream.LongStream; - -import static com.google.common.base.MoreObjects.toStringHelper; -import static com.google.common.base.Preconditions.checkArgument; -import static com.google.common.base.Preconditions.checkState; -import static com.google.common.base.Verify.verify; -import static com.google.common.collect.ImmutableList.toImmutableList; -import static com.google.common.collect.ImmutableMap.toImmutableMap; -import static com.google.common.collect.ImmutableSet.toImmutableSet; -import static com.google.common.collect.Iterables.concat; -import static com.google.common.collect.Iterables.getOnlyElement; -import static com.google.common.collect.Lists.newArrayList; -import static com.google.common.collect.Lists.reverse; -import static com.google.common.collect.Maps.uniqueIndex; -import static com.google.common.collect.MoreCollectors.onlyElement; -import static com.google.common.collect.Sets.difference; -import static com.google.common.collect.Streams.stream; -import static com.google.common.hash.Hashing.sha256; -import static com.google.common.io.MoreFiles.deleteRecursively; -import static com.google.common.io.RecursiveDeleteOption.ALLOW_INSECURE; -import static io.airlift.concurrent.MoreFutures.getFutureValue; -import static io.airlift.concurrent.Threads.daemonThreadsNamed; -import static io.airlift.slice.Slices.utf8Slice; -import static io.airlift.testing.Assertions.assertGreaterThan; -import static io.airlift.testing.Assertions.assertGreaterThanOrEqual; -import static io.airlift.testing.Assertions.assertInstanceOf; -import static io.airlift.testing.Assertions.assertLessThanOrEqual; -import static io.airlift.units.DataSize.Unit.KILOBYTE; -import static io.trino.parquet.reader.ParquetReader.PARQUET_CODEC_METRIC_PREFIX; -import static io.trino.plugin.hive.AbstractTestHive.TransactionDeleteInsertTestTag.COMMIT; -import static io.trino.plugin.hive.AbstractTestHive.TransactionDeleteInsertTestTag.ROLLBACK_AFTER_APPEND_PAGE; -import static io.trino.plugin.hive.AbstractTestHive.TransactionDeleteInsertTestTag.ROLLBACK_AFTER_BEGIN_INSERT; -import static io.trino.plugin.hive.AbstractTestHive.TransactionDeleteInsertTestTag.ROLLBACK_AFTER_DELETE; -import static io.trino.plugin.hive.AbstractTestHive.TransactionDeleteInsertTestTag.ROLLBACK_AFTER_FINISH_INSERT; -import static io.trino.plugin.hive.AbstractTestHive.TransactionDeleteInsertTestTag.ROLLBACK_AFTER_SINK_FINISH; -import static io.trino.plugin.hive.AbstractTestHive.TransactionDeleteInsertTestTag.ROLLBACK_RIGHT_AWAY; -import static io.trino.plugin.hive.HiveBasicStatistics.createEmptyStatistics; -import static io.trino.plugin.hive.HiveBasicStatistics.createZeroStatistics; -import static io.trino.plugin.hive.HiveColumnHandle.BUCKET_COLUMN_NAME; -import static io.trino.plugin.hive.HiveColumnHandle.ColumnType.PARTITION_KEY; -import static io.trino.plugin.hive.HiveColumnHandle.ColumnType.REGULAR; -import static io.trino.plugin.hive.HiveColumnHandle.bucketColumnHandle; -import static io.trino.plugin.hive.HiveColumnHandle.createBaseColumn; -import static io.trino.plugin.hive.HiveErrorCode.HIVE_INVALID_BUCKET_FILES; -import static io.trino.plugin.hive.HiveErrorCode.HIVE_INVALID_PARTITION_VALUE; -import static io.trino.plugin.hive.HiveErrorCode.HIVE_PARTITION_SCHEMA_MISMATCH; -import static io.trino.plugin.hive.HiveMetadata.PRESTO_QUERY_ID_NAME; -import static io.trino.plugin.hive.HiveMetadata.PRESTO_VERSION_NAME; -import static io.trino.plugin.hive.HiveStorageFormat.AVRO; -import static io.trino.plugin.hive.HiveStorageFormat.CSV; -import static io.trino.plugin.hive.HiveStorageFormat.JSON; -import static io.trino.plugin.hive.HiveStorageFormat.ORC; -import static io.trino.plugin.hive.HiveStorageFormat.PARQUET; -import static io.trino.plugin.hive.HiveStorageFormat.RCBINARY; -import static io.trino.plugin.hive.HiveStorageFormat.RCTEXT; -import static io.trino.plugin.hive.HiveStorageFormat.REGEX; -import static io.trino.plugin.hive.HiveStorageFormat.SEQUENCEFILE; -import static io.trino.plugin.hive.HiveStorageFormat.TEXTFILE; -import static io.trino.plugin.hive.HiveTableProperties.BUCKETED_BY_PROPERTY; -import static io.trino.plugin.hive.HiveTableProperties.BUCKET_COUNT_PROPERTY; -import static io.trino.plugin.hive.HiveTableProperties.EXTERNAL_LOCATION_PROPERTY; -import static io.trino.plugin.hive.HiveTableProperties.PARTITIONED_BY_PROPERTY; -import static io.trino.plugin.hive.HiveTableProperties.SORTED_BY_PROPERTY; -import static io.trino.plugin.hive.HiveTableProperties.STORAGE_FORMAT_PROPERTY; -import static io.trino.plugin.hive.HiveTableProperties.TRANSACTIONAL; -import static io.trino.plugin.hive.HiveTestUtils.HDFS_CONFIGURATION; -import static io.trino.plugin.hive.HiveTestUtils.HDFS_ENVIRONMENT; -import static io.trino.plugin.hive.HiveTestUtils.HDFS_FILE_SYSTEM_STATS; -import static io.trino.plugin.hive.HiveTestUtils.PAGE_SORTER; -import static io.trino.plugin.hive.HiveTestUtils.SESSION; -import static io.trino.plugin.hive.HiveTestUtils.arrayType; -import static io.trino.plugin.hive.HiveTestUtils.getDefaultHiveFileWriterFactories; -import static io.trino.plugin.hive.HiveTestUtils.getDefaultHivePageSourceFactories; -import static io.trino.plugin.hive.HiveTestUtils.getDefaultHiveRecordCursorProviders; -import static io.trino.plugin.hive.HiveTestUtils.getHiveSession; -import static io.trino.plugin.hive.HiveTestUtils.getHiveSessionProperties; -import static io.trino.plugin.hive.HiveTestUtils.getTypes; -import static io.trino.plugin.hive.HiveTestUtils.mapType; -import static io.trino.plugin.hive.HiveTestUtils.rowType; -import static io.trino.plugin.hive.HiveType.HIVE_INT; -import static io.trino.plugin.hive.HiveType.HIVE_LONG; -import static io.trino.plugin.hive.HiveType.HIVE_STRING; -import static io.trino.plugin.hive.HiveType.toHiveType; -import static io.trino.plugin.hive.LocationHandle.WriteMode.STAGE_AND_MOVE_TO_TARGET_DIRECTORY; -import static io.trino.plugin.hive.TestingThriftHiveMetastoreBuilder.testingThriftHiveMetastoreBuilder; -import static io.trino.plugin.hive.acid.AcidTransaction.NO_ACID_TRANSACTION; -import static io.trino.plugin.hive.metastore.HiveColumnStatistics.createBinaryColumnStatistics; -import static io.trino.plugin.hive.metastore.HiveColumnStatistics.createBooleanColumnStatistics; -import static io.trino.plugin.hive.metastore.HiveColumnStatistics.createDateColumnStatistics; -import static io.trino.plugin.hive.metastore.HiveColumnStatistics.createDecimalColumnStatistics; -import static io.trino.plugin.hive.metastore.HiveColumnStatistics.createDoubleColumnStatistics; -import static io.trino.plugin.hive.metastore.HiveColumnStatistics.createIntegerColumnStatistics; -import static io.trino.plugin.hive.metastore.HiveColumnStatistics.createStringColumnStatistics; -import static io.trino.plugin.hive.metastore.PrincipalPrivileges.NO_PRIVILEGES; -import static io.trino.plugin.hive.metastore.SortingColumn.Order.ASCENDING; -import static io.trino.plugin.hive.metastore.SortingColumn.Order.DESCENDING; -import static io.trino.plugin.hive.metastore.StorageFormat.fromHiveStorageFormat; -import static io.trino.plugin.hive.orc.OrcPageSource.ORC_CODEC_METRIC_PREFIX; -import static io.trino.plugin.hive.util.HiveBucketing.BucketingVersion.BUCKETING_V1; -import static io.trino.plugin.hive.util.HiveUtil.DELTA_LAKE_PROVIDER; -import static io.trino.plugin.hive.util.HiveUtil.ICEBERG_TABLE_TYPE_NAME; -import static io.trino.plugin.hive.util.HiveUtil.ICEBERG_TABLE_TYPE_VALUE; -import static io.trino.plugin.hive.util.HiveUtil.SPARK_TABLE_PROVIDER_KEY; -import static io.trino.plugin.hive.util.HiveUtil.columnExtraInfo; -import static io.trino.plugin.hive.util.HiveUtil.toPartitionValues; -import static io.trino.plugin.hive.util.HiveWriteUtils.createDirectory; -import static io.trino.plugin.hive.util.HiveWriteUtils.getTableDefaultLocation; -import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED; -import static io.trino.spi.StandardErrorCode.TRANSACTION_CONFLICT; -import static io.trino.spi.connector.MetadataProvider.NOOP_METADATA_PROVIDER; -import static io.trino.spi.connector.RetryMode.NO_RETRIES; -import static io.trino.spi.connector.SortOrder.ASC_NULLS_FIRST; -import static io.trino.spi.connector.SortOrder.DESC_NULLS_LAST; -import static io.trino.spi.security.PrincipalType.USER; -import static io.trino.spi.type.BigintType.BIGINT; -import static io.trino.spi.type.BooleanType.BOOLEAN; -import static io.trino.spi.type.CharType.createCharType; -import static io.trino.spi.type.DateType.DATE; -import static io.trino.spi.type.DecimalType.createDecimalType; -import static io.trino.spi.type.DoubleType.DOUBLE; -import static io.trino.spi.type.HyperLogLogType.HYPER_LOG_LOG; -import static io.trino.spi.type.IntegerType.INTEGER; -import static io.trino.spi.type.RealType.REAL; -import static io.trino.spi.type.SmallintType.SMALLINT; -import static io.trino.spi.type.TimestampType.TIMESTAMP_MILLIS; -import static io.trino.spi.type.TimestampWithTimeZoneType.TIMESTAMP_TZ_MILLIS; -import static io.trino.spi.type.TinyintType.TINYINT; -import static io.trino.spi.type.VarbinaryType.VARBINARY; -import static io.trino.spi.type.VarcharType.VARCHAR; -import static io.trino.spi.type.VarcharType.createUnboundedVarcharType; -import static io.trino.spi.type.VarcharType.createVarcharType; -import static io.trino.testing.DateTimeTestingUtils.sqlTimestampOf; -import static io.trino.testing.MaterializedResult.materializeSourceDataStream; -import static io.trino.testing.QueryAssertions.assertEqualsIgnoreOrder; -import static io.trino.testing.TestingNames.randomNameSuffix; -import static io.trino.testing.TestingPageSinkId.TESTING_PAGE_SINK_ID; -import static io.trino.testing.assertions.TrinoExceptionAssert.assertTrinoExceptionThrownBy; -import static io.trino.type.InternalTypeManager.TESTING_TYPE_MANAGER; -import static java.lang.Float.floatToRawIntBits; -import static java.lang.Math.toIntExact; -import static java.lang.String.format; -import static java.nio.charset.StandardCharsets.UTF_8; -import static java.nio.file.Files.createTempDirectory; -import static java.util.Locale.ENGLISH; -import static java.util.Objects.requireNonNull; -import static java.util.concurrent.Executors.newCachedThreadPool; -import static java.util.concurrent.Executors.newScheduledThreadPool; -import static java.util.concurrent.TimeUnit.MILLISECONDS; -import static java.util.concurrent.TimeUnit.MINUTES; -import static java.util.concurrent.TimeUnit.SECONDS; -import static java.util.stream.Collectors.toList; -import static org.apache.hadoop.hive.common.FileUtils.makePartName; -import static org.apache.hadoop.hive.metastore.TableType.MANAGED_TABLE; -import static org.assertj.core.api.Assertions.assertThat; -import static org.assertj.core.api.Assertions.assertThatThrownBy; -import static org.joda.time.DateTimeZone.UTC; -import static org.testng.Assert.assertEquals; -import static org.testng.Assert.assertFalse; -import static org.testng.Assert.assertNotNull; -import static org.testng.Assert.assertNull; -import static org.testng.Assert.assertTrue; -import static org.testng.Assert.fail; - -// staging directory is shared mutable state -@Test(singleThreaded = true) -public abstract class AbstractTestHive -{ - private static final Logger log = Logger.get(AbstractTestHive.class); - - protected static final String TEMPORARY_TABLE_PREFIX = "tmp_trino_test_"; - - protected static final String INVALID_DATABASE = "totally_invalid_database_name"; - protected static final String INVALID_TABLE = "totally_invalid_table_name"; - protected static final String INVALID_COLUMN = "totally_invalid_column_name"; - - protected static final String TEST_SERVER_VERSION = "test_version"; - - private static final Type ARRAY_TYPE = arrayType(createUnboundedVarcharType()); - private static final Type MAP_TYPE = mapType(createUnboundedVarcharType(), BIGINT); - private static final Type ROW_TYPE = rowType(ImmutableList.of( - new NamedTypeSignature(Optional.of(new RowFieldName("f_string")), createUnboundedVarcharType().getTypeSignature()), - new NamedTypeSignature(Optional.of(new RowFieldName("f_bigint")), BIGINT.getTypeSignature()), - new NamedTypeSignature(Optional.of(new RowFieldName("f_boolean")), BOOLEAN.getTypeSignature()))); - - private static final List CREATE_TABLE_COLUMNS = ImmutableList.builder() - .add(new ColumnMetadata("id", BIGINT)) - .add(new ColumnMetadata("t_string", createUnboundedVarcharType())) - .add(new ColumnMetadata("t_tinyint", TINYINT)) - .add(new ColumnMetadata("t_smallint", SMALLINT)) - .add(new ColumnMetadata("t_integer", INTEGER)) - .add(new ColumnMetadata("t_bigint", BIGINT)) - .add(new ColumnMetadata("t_float", REAL)) - .add(new ColumnMetadata("t_double", DOUBLE)) - .add(new ColumnMetadata("t_boolean", BOOLEAN)) - .add(new ColumnMetadata("t_array", ARRAY_TYPE)) - .add(new ColumnMetadata("t_map", MAP_TYPE)) - .add(new ColumnMetadata("t_row", ROW_TYPE)) - .build(); - - private static final MaterializedResult CREATE_TABLE_DATA = - MaterializedResult.resultBuilder(SESSION, BIGINT, createUnboundedVarcharType(), TINYINT, SMALLINT, INTEGER, BIGINT, REAL, DOUBLE, BOOLEAN, ARRAY_TYPE, MAP_TYPE, ROW_TYPE) - .row(1L, "hello", (byte) 45, (short) 345, 234, 123L, -754.1985f, 43.5, true, ImmutableList.of("apple", "banana"), ImmutableMap.of("one", 1L, "two", 2L), ImmutableList.of("true", 1L, true)) - .row(2L, null, null, null, null, null, null, null, null, null, null, null) - .row(3L, "bye", (byte) 46, (short) 346, 345, 456L, 754.2008f, 98.1, false, ImmutableList.of("ape", "bear"), ImmutableMap.of("three", 3L, "four", 4L), ImmutableList.of("false", 0L, false)) - .build(); - - protected static final List CREATE_TABLE_COLUMNS_PARTITIONED = ImmutableList.builder() - .addAll(CREATE_TABLE_COLUMNS) - .add(new ColumnMetadata("ds", createUnboundedVarcharType())) - .build(); - - protected static final Set COLUMN_NAMES_PARTITIONED = CREATE_TABLE_COLUMNS_PARTITIONED.stream().map(ColumnMetadata::getName).collect(toImmutableSet()); - - protected static final Predicate PARTITION_COLUMN_FILTER = columnName -> columnName.equals("ds") || columnName.startsWith("part_"); - - private static final MaterializedResult CREATE_TABLE_PARTITIONED_DATA = new MaterializedResult( - CREATE_TABLE_DATA.getMaterializedRows().stream() - .map(row -> new MaterializedRow(row.getPrecision(), newArrayList(concat(row.getFields(), ImmutableList.of("2015-07-0" + row.getField(0)))))) - .collect(toList()), - ImmutableList.builder() - .addAll(CREATE_TABLE_DATA.getTypes()) - .add(createUnboundedVarcharType()) - .build()); - - private static final String CREATE_TABLE_PARTITIONED_DATA_2ND_PARTITION_VALUE = "2015-07-04"; - - private static final MaterializedResult CREATE_TABLE_PARTITIONED_DATA_2ND = - MaterializedResult.resultBuilder(SESSION, BIGINT, createUnboundedVarcharType(), TINYINT, SMALLINT, INTEGER, BIGINT, REAL, DOUBLE, BOOLEAN, ARRAY_TYPE, MAP_TYPE, ROW_TYPE, createUnboundedVarcharType()) - .row(4L, "hello", (byte) 45, (short) 345, 234, 123L, 754.1985f, 43.5, true, ImmutableList.of("apple", "banana"), ImmutableMap.of("one", 1L, "two", 2L), ImmutableList.of("true", 1L, true), CREATE_TABLE_PARTITIONED_DATA_2ND_PARTITION_VALUE) - .row(5L, null, null, null, null, null, null, null, null, null, null, null, CREATE_TABLE_PARTITIONED_DATA_2ND_PARTITION_VALUE) - .row(6L, "bye", (byte) 46, (short) 346, 345, 456L, -754.2008f, 98.1, false, ImmutableList.of("ape", "bear"), ImmutableMap.of("three", 3L, "four", 4L), ImmutableList.of("false", 0L, false), CREATE_TABLE_PARTITIONED_DATA_2ND_PARTITION_VALUE) - .build(); - - private static final List MISMATCH_SCHEMA_PRIMITIVE_COLUMN_BEFORE = ImmutableList.builder() - .add(new ColumnMetadata("tinyint_to_smallint", TINYINT)) - .add(new ColumnMetadata("tinyint_to_integer", TINYINT)) - .add(new ColumnMetadata("tinyint_to_bigint", TINYINT)) - .add(new ColumnMetadata("smallint_to_integer", SMALLINT)) - .add(new ColumnMetadata("smallint_to_bigint", SMALLINT)) - .add(new ColumnMetadata("integer_to_bigint", INTEGER)) - .add(new ColumnMetadata("integer_to_varchar", INTEGER)) - .add(new ColumnMetadata("varchar_to_integer", createUnboundedVarcharType())) - .add(new ColumnMetadata("float_to_double", REAL)) - .add(new ColumnMetadata("varchar_to_drop_in_row", createUnboundedVarcharType())) - .build(); - - private static final List MISMATCH_SCHEMA_TABLE_BEFORE = ImmutableList.builder() - .addAll(MISMATCH_SCHEMA_PRIMITIVE_COLUMN_BEFORE) - .add(new ColumnMetadata("struct_to_struct", toRowType(MISMATCH_SCHEMA_PRIMITIVE_COLUMN_BEFORE))) - .add(new ColumnMetadata("list_to_list", arrayType(toRowType(MISMATCH_SCHEMA_PRIMITIVE_COLUMN_BEFORE)))) - .add(new ColumnMetadata("map_to_map", mapType(MISMATCH_SCHEMA_PRIMITIVE_COLUMN_BEFORE.get(1).getType(), toRowType(MISMATCH_SCHEMA_PRIMITIVE_COLUMN_BEFORE)))) - .add(new ColumnMetadata("ds", createUnboundedVarcharType())) - .build(); - - private static RowType toRowType(List columns) - { - return rowType(columns.stream() - .map(col -> new NamedTypeSignature(Optional.of(new RowFieldName(format("f_%s", col.getName()))), col.getType().getTypeSignature())) - .collect(toImmutableList())); - } - - private static final MaterializedResult MISMATCH_SCHEMA_PRIMITIVE_FIELDS_DATA_BEFORE = - MaterializedResult.resultBuilder(SESSION, TINYINT, TINYINT, TINYINT, SMALLINT, SMALLINT, INTEGER, INTEGER, createUnboundedVarcharType(), REAL, createUnboundedVarcharType()) - .row((byte) -11, (byte) 12, (byte) -13, (short) 14, (short) 15, -16, 17, "2147483647", 18.0f, "2016-08-01") - .row((byte) 21, (byte) -22, (byte) 23, (short) -24, (short) 25, 26, -27, "asdf", -28.0f, "2016-08-02") - .row((byte) -31, (byte) -32, (byte) 33, (short) 34, (short) -35, 36, 37, "-923", 39.5f, "2016-08-03") - .row(null, (byte) 42, (byte) 43, (short) 44, (short) -45, 46, 47, "2147483648", 49.5f, "2016-08-03") - .build(); - - private static final MaterializedResult MISMATCH_SCHEMA_TABLE_DATA_BEFORE = - MaterializedResult.resultBuilder(SESSION, MISMATCH_SCHEMA_TABLE_BEFORE.stream().map(ColumnMetadata::getType).collect(toImmutableList())) - .rows(MISMATCH_SCHEMA_PRIMITIVE_FIELDS_DATA_BEFORE.getMaterializedRows() - .stream() - .map(materializedRow -> { - List result = materializedRow.getFields(); - List rowResult = materializedRow.getFields(); - result.add(rowResult); - result.add(Arrays.asList(rowResult, null, rowResult)); - result.add(ImmutableMap.of(rowResult.get(1), rowResult)); - result.add(rowResult.get(9)); - return new MaterializedRow(materializedRow.getPrecision(), result); - }).collect(toImmutableList())) - .build(); - - private static final List MISMATCH_SCHEMA_PRIMITIVE_COLUMN_AFTER = ImmutableList.builder() - .add(new ColumnMetadata("tinyint_to_smallint", SMALLINT)) - .add(new ColumnMetadata("tinyint_to_integer", INTEGER)) - .add(new ColumnMetadata("tinyint_to_bigint", BIGINT)) - .add(new ColumnMetadata("smallint_to_integer", INTEGER)) - .add(new ColumnMetadata("smallint_to_bigint", BIGINT)) - .add(new ColumnMetadata("integer_to_bigint", BIGINT)) - .add(new ColumnMetadata("integer_to_varchar", createUnboundedVarcharType())) - .add(new ColumnMetadata("varchar_to_integer", INTEGER)) - .add(new ColumnMetadata("float_to_double", DOUBLE)) - .add(new ColumnMetadata("varchar_to_drop_in_row", createUnboundedVarcharType())) - .build(); - - private static final Type MISMATCH_SCHEMA_ROW_TYPE_APPEND = toRowType(ImmutableList.builder() - .addAll(MISMATCH_SCHEMA_PRIMITIVE_COLUMN_AFTER) - .add(new ColumnMetadata(format("%s_append", MISMATCH_SCHEMA_PRIMITIVE_COLUMN_AFTER.get(0).getName()), MISMATCH_SCHEMA_PRIMITIVE_COLUMN_AFTER.get(0).getType())) - .build()); - private static final Type MISMATCH_SCHEMA_ROW_TYPE_DROP = toRowType(MISMATCH_SCHEMA_PRIMITIVE_COLUMN_AFTER.subList(0, MISMATCH_SCHEMA_PRIMITIVE_COLUMN_AFTER.size() - 1)); - - private static final List MISMATCH_SCHEMA_TABLE_AFTER = ImmutableList.builder() - .addAll(MISMATCH_SCHEMA_PRIMITIVE_COLUMN_AFTER) - .add(new ColumnMetadata("struct_to_struct", MISMATCH_SCHEMA_ROW_TYPE_APPEND)) - .add(new ColumnMetadata("list_to_list", arrayType(MISMATCH_SCHEMA_ROW_TYPE_APPEND))) - .add(new ColumnMetadata("map_to_map", mapType(MISMATCH_SCHEMA_PRIMITIVE_COLUMN_AFTER.get(1).getType(), MISMATCH_SCHEMA_ROW_TYPE_DROP))) - .add(new ColumnMetadata("ds", createUnboundedVarcharType())) - .build(); - - private static final MaterializedResult MISMATCH_SCHEMA_PRIMITIVE_FIELDS_DATA_AFTER = - MaterializedResult.resultBuilder(SESSION, SMALLINT, INTEGER, BIGINT, INTEGER, BIGINT, BIGINT, createUnboundedVarcharType(), INTEGER, DOUBLE, createUnboundedVarcharType()) - .row((short) -11, 12, -13L, 14, 15L, -16L, "17", 2147483647, 18.0, "2016-08-01") - .row((short) 21, -22, 23L, -24, 25L, 26L, "-27", null, -28.0, "2016-08-02") - .row((short) -31, -32, 33L, 34, -35L, 36L, "37", -923, 39.5, "2016-08-03") - .row(null, 42, 43L, 44, -45L, 46L, "47", null, 49.5, "2016-08-03") - .build(); - - private static final MaterializedResult MISMATCH_SCHEMA_TABLE_DATA_AFTER = - MaterializedResult.resultBuilder(SESSION, MISMATCH_SCHEMA_TABLE_AFTER.stream().map(ColumnMetadata::getType).collect(toImmutableList())) - .rows(MISMATCH_SCHEMA_PRIMITIVE_FIELDS_DATA_AFTER.getMaterializedRows() - .stream() - .map(materializedRow -> { - List result = materializedRow.getFields(); - List appendFieldRowResult = materializedRow.getFields(); - appendFieldRowResult.add(null); - List dropFieldRowResult = materializedRow.getFields().subList(0, materializedRow.getFields().size() - 1); - result.add(appendFieldRowResult); - result.add(Arrays.asList(appendFieldRowResult, null, appendFieldRowResult)); - result.add(ImmutableMap.of(result.get(1), dropFieldRowResult)); - result.add(result.get(9)); - return new MaterializedRow(materializedRow.getPrecision(), result); - }).collect(toImmutableList())) - .build(); - - protected Set createTableFormats = difference( - ImmutableSet.copyOf(HiveStorageFormat.values()), - // exclude formats that change table schema with serde and read-only formats - ImmutableSet.of(AVRO, CSV, REGEX)); - - private static final TypeOperators TYPE_OPERATORS = new TypeOperators(); - private static final BlockTypeOperators BLOCK_TYPE_OPERATORS = new BlockTypeOperators(TYPE_OPERATORS); - private static final JoinCompiler JOIN_COMPILER = new JoinCompiler(TYPE_OPERATORS); - - protected static final List STATISTICS_TABLE_COLUMNS = ImmutableList.builder() - .add(new ColumnMetadata("t_boolean", BOOLEAN)) - .add(new ColumnMetadata("t_bigint", BIGINT)) - .add(new ColumnMetadata("t_integer", INTEGER)) - .add(new ColumnMetadata("t_smallint", SMALLINT)) - .add(new ColumnMetadata("t_tinyint", TINYINT)) - .add(new ColumnMetadata("t_double", DOUBLE)) - .add(new ColumnMetadata("t_float", REAL)) - .add(new ColumnMetadata("t_string", createUnboundedVarcharType())) - .add(new ColumnMetadata("t_varchar", createVarcharType(100))) - .add(new ColumnMetadata("t_char", createCharType(5))) - .add(new ColumnMetadata("t_varbinary", VARBINARY)) - .add(new ColumnMetadata("t_date", DATE)) - .add(new ColumnMetadata("t_timestamp", TIMESTAMP_MILLIS)) - .add(new ColumnMetadata("t_short_decimal", createDecimalType(5, 2))) - .add(new ColumnMetadata("t_long_decimal", createDecimalType(20, 3))) - .build(); - - protected static final List STATISTICS_PARTITIONED_TABLE_COLUMNS = ImmutableList.builder() - .addAll(STATISTICS_TABLE_COLUMNS) - .add(new ColumnMetadata("ds", VARCHAR)) - .build(); - - protected static final PartitionStatistics ZERO_TABLE_STATISTICS = new PartitionStatistics(createZeroStatistics(), ImmutableMap.of()); - protected static final PartitionStatistics BASIC_STATISTICS_1 = new PartitionStatistics(new HiveBasicStatistics(0, 20, 3, 0), ImmutableMap.of()); - protected static final PartitionStatistics BASIC_STATISTICS_2 = new PartitionStatistics(new HiveBasicStatistics(0, 30, 2, 0), ImmutableMap.of()); - - protected static final PartitionStatistics STATISTICS_1 = - new PartitionStatistics( - BASIC_STATISTICS_1.getBasicStatistics(), - ImmutableMap.builder() - .put("t_boolean", createBooleanColumnStatistics(OptionalLong.of(5), OptionalLong.of(6), OptionalLong.of(3))) - .put("t_bigint", createIntegerColumnStatistics(OptionalLong.of(1234L), OptionalLong.of(5678L), OptionalLong.of(2), OptionalLong.of(5))) - .put("t_integer", createIntegerColumnStatistics(OptionalLong.of(123L), OptionalLong.of(567L), OptionalLong.of(3), OptionalLong.of(4))) - .put("t_smallint", createIntegerColumnStatistics(OptionalLong.of(12L), OptionalLong.of(56L), OptionalLong.of(2), OptionalLong.of(6))) - .put("t_tinyint", createIntegerColumnStatistics(OptionalLong.of(1L), OptionalLong.of(2L), OptionalLong.of(1), OptionalLong.of(3))) - .put("t_double", createDoubleColumnStatistics(OptionalDouble.of(1234.25), OptionalDouble.of(5678.58), OptionalLong.of(7), OptionalLong.of(8))) - .put("t_float", createDoubleColumnStatistics(OptionalDouble.of(123.25), OptionalDouble.of(567.58), OptionalLong.of(9), OptionalLong.of(10))) - .put("t_string", createStringColumnStatistics(OptionalLong.of(10), OptionalLong.of(50), OptionalLong.of(3), OptionalLong.of(7))) - .put("t_varchar", createStringColumnStatistics(OptionalLong.of(100), OptionalLong.of(230), OptionalLong.of(5), OptionalLong.of(3))) - .put("t_char", createStringColumnStatistics(OptionalLong.of(5), OptionalLong.of(50), OptionalLong.of(1), OptionalLong.of(4))) - .put("t_varbinary", createBinaryColumnStatistics(OptionalLong.of(4), OptionalLong.of(50), OptionalLong.of(1))) - .put("t_date", createDateColumnStatistics(Optional.of(LocalDate.ofEpochDay(1)), Optional.of(LocalDate.ofEpochDay(2)), OptionalLong.of(7), OptionalLong.of(6))) - .put("t_timestamp", createIntegerColumnStatistics(OptionalLong.of(1234567L), OptionalLong.of(71234567L), OptionalLong.of(7), OptionalLong.of(5))) - .put("t_short_decimal", createDecimalColumnStatistics(Optional.of(new BigDecimal(10)), Optional.of(new BigDecimal(12)), OptionalLong.of(3), OptionalLong.of(5))) - .put("t_long_decimal", createDecimalColumnStatistics(Optional.of(new BigDecimal("12345678901234567.123")), Optional.of(new BigDecimal("81234567890123456.123")), OptionalLong.of(2), OptionalLong.of(1))) - .buildOrThrow()); - - protected static final PartitionStatistics STATISTICS_1_1 = - new PartitionStatistics( - new HiveBasicStatistics(OptionalLong.of(0), OptionalLong.of(15), OptionalLong.empty(), OptionalLong.of(0)), - STATISTICS_1.getColumnStatistics().entrySet() - .stream() - .filter(entry -> entry.getKey().hashCode() % 2 == 0) - .collect(toImmutableMap(Map.Entry::getKey, Map.Entry::getValue))); - - protected static final PartitionStatistics STATISTICS_1_2 = - new PartitionStatistics( - new HiveBasicStatistics(OptionalLong.of(0), OptionalLong.of(15), OptionalLong.of(3), OptionalLong.of(0)), - STATISTICS_1.getColumnStatistics().entrySet() - .stream() - .filter(entry -> entry.getKey().hashCode() % 2 == 1) - .collect(toImmutableMap(Map.Entry::getKey, Map.Entry::getValue))); - - protected static final PartitionStatistics STATISTICS_2 = - new PartitionStatistics( - BASIC_STATISTICS_2.getBasicStatistics(), - ImmutableMap.builder() - .put("t_boolean", createBooleanColumnStatistics(OptionalLong.of(4), OptionalLong.of(3), OptionalLong.of(2))) - .put("t_bigint", createIntegerColumnStatistics(OptionalLong.of(2345L), OptionalLong.of(6789L), OptionalLong.of(4), OptionalLong.of(7))) - .put("t_integer", createIntegerColumnStatistics(OptionalLong.of(234L), OptionalLong.of(678L), OptionalLong.of(5), OptionalLong.of(6))) - .put("t_smallint", createIntegerColumnStatistics(OptionalLong.of(23L), OptionalLong.of(65L), OptionalLong.of(7), OptionalLong.of(5))) - .put("t_tinyint", createIntegerColumnStatistics(OptionalLong.of(3L), OptionalLong.of(12L), OptionalLong.of(2), OptionalLong.of(3))) - .put("t_double", createDoubleColumnStatistics(OptionalDouble.of(2345.25), OptionalDouble.of(6785.58), OptionalLong.of(6), OptionalLong.of(3))) - .put("t_float", createDoubleColumnStatistics(OptionalDouble.of(235.25), OptionalDouble.of(676.58), OptionalLong.of(7), OptionalLong.of(11))) - .put("t_string", createStringColumnStatistics(OptionalLong.of(301), OptionalLong.of(600), OptionalLong.of(2), OptionalLong.of(6))) - .put("t_varchar", createStringColumnStatistics(OptionalLong.of(99), OptionalLong.of(223), OptionalLong.of(7), OptionalLong.of(1))) - .put("t_char", createStringColumnStatistics(OptionalLong.of(6), OptionalLong.of(60), OptionalLong.of(0), OptionalLong.of(3))) - .put("t_varbinary", createBinaryColumnStatistics(OptionalLong.of(2), OptionalLong.of(10), OptionalLong.of(2))) - .put("t_date", createDateColumnStatistics(Optional.of(LocalDate.ofEpochDay(2)), Optional.of(LocalDate.ofEpochDay(3)), OptionalLong.of(8), OptionalLong.of(7))) - .put("t_timestamp", createIntegerColumnStatistics(OptionalLong.of(2345671L), OptionalLong.of(12345677L), OptionalLong.of(9), OptionalLong.of(1))) - .put("t_short_decimal", createDecimalColumnStatistics(Optional.of(new BigDecimal(11)), Optional.of(new BigDecimal(14)), OptionalLong.of(5), OptionalLong.of(7))) - .put("t_long_decimal", createDecimalColumnStatistics(Optional.of(new BigDecimal("71234567890123456.123")), Optional.of(new BigDecimal("78123456789012345.123")), OptionalLong.of(2), OptionalLong.of(1))) - .buildOrThrow()); - - protected static final PartitionStatistics STATISTICS_EMPTY_OPTIONAL_FIELDS = - new PartitionStatistics( - new HiveBasicStatistics(OptionalLong.of(0), OptionalLong.of(20), OptionalLong.empty(), OptionalLong.of(0)), - ImmutableMap.builder() - .put("t_boolean", createBooleanColumnStatistics(OptionalLong.of(4), OptionalLong.of(3), OptionalLong.of(2))) - .put("t_bigint", createIntegerColumnStatistics(OptionalLong.empty(), OptionalLong.empty(), OptionalLong.of(4), OptionalLong.of(7))) - .put("t_integer", createIntegerColumnStatistics(OptionalLong.empty(), OptionalLong.empty(), OptionalLong.of(5), OptionalLong.of(6))) - .put("t_smallint", createIntegerColumnStatistics(OptionalLong.empty(), OptionalLong.empty(), OptionalLong.of(7), OptionalLong.of(5))) - .put("t_tinyint", createIntegerColumnStatistics(OptionalLong.empty(), OptionalLong.empty(), OptionalLong.of(2), OptionalLong.of(3))) - .put("t_double", createDoubleColumnStatistics(OptionalDouble.empty(), OptionalDouble.empty(), OptionalLong.of(6), OptionalLong.of(3))) - .put("t_float", createDoubleColumnStatistics(OptionalDouble.empty(), OptionalDouble.empty(), OptionalLong.of(7), OptionalLong.of(11))) - .put("t_string", createStringColumnStatistics(OptionalLong.of(0), OptionalLong.of(0), OptionalLong.of(2), OptionalLong.of(6))) - .put("t_varchar", createStringColumnStatistics(OptionalLong.of(0), OptionalLong.of(0), OptionalLong.of(7), OptionalLong.of(1))) - .put("t_char", createStringColumnStatistics(OptionalLong.of(0), OptionalLong.of(0), OptionalLong.of(0), OptionalLong.of(3))) - .put("t_varbinary", createBinaryColumnStatistics(OptionalLong.of(0), OptionalLong.of(0), OptionalLong.of(2))) - // https://issues.apache.org/jira/browse/HIVE-20098 - // .put("t_date", createDateColumnStatistics(Optional.empty(), Optional.empty(), OptionalLong.of(8), OptionalLong.of(7))) - .put("t_timestamp", createIntegerColumnStatistics(OptionalLong.empty(), OptionalLong.empty(), OptionalLong.of(9), OptionalLong.of(1))) - .put("t_short_decimal", createDecimalColumnStatistics(Optional.empty(), Optional.empty(), OptionalLong.of(5), OptionalLong.of(7))) - .put("t_long_decimal", createDecimalColumnStatistics(Optional.empty(), Optional.empty(), OptionalLong.of(2), OptionalLong.of(1))) - .buildOrThrow()); - - protected String database; - protected SchemaTableName tablePartitionFormat; - protected SchemaTableName tableUnpartitioned; - protected SchemaTableName tablePartitionedWithNull; - protected SchemaTableName tableOffline; - protected SchemaTableName tableNotReadable; - protected SchemaTableName view; - protected SchemaTableName invalidTable; - protected SchemaTableName tableBucketedStringInt; - protected SchemaTableName tableBucketedBigintBoolean; - protected SchemaTableName tableBucketedDoubleFloat; - protected SchemaTableName tablePartitionSchemaChange; - protected SchemaTableName tablePartitionSchemaChangeNonCanonical; - protected SchemaTableName tableBucketEvolution; - - protected ConnectorTableHandle invalidTableHandle; - - protected ColumnHandle dsColumn; - protected ColumnHandle fileFormatColumn; - protected ColumnHandle dummyColumn; - protected ColumnHandle intColumn; - protected ColumnHandle invalidColumnHandle; - protected ColumnHandle pStringColumn; - protected ColumnHandle pIntegerColumn; - - protected ConnectorTableProperties tablePartitionFormatProperties; - protected ConnectorTableProperties tableUnpartitionedProperties; - protected List tablePartitionFormatPartitions; - protected List tableUnpartitionedPartitions; - - protected HdfsEnvironment hdfsEnvironment; - protected LocationService locationService; - - protected CountingDirectoryLister countingDirectoryLister; - protected HiveMetadataFactory metadataFactory; - protected HiveTransactionManager transactionManager; - protected HiveMetastore metastoreClient; - protected ConnectorSplitManager splitManager; - protected ConnectorPageSourceProvider pageSourceProvider; - protected ConnectorPageSinkProvider pageSinkProvider; - protected ConnectorNodePartitioningProvider nodePartitioningProvider; - protected ExecutorService executor; - - private ScheduledExecutorService heartbeatService; - private java.nio.file.Path temporaryStagingDirectory; - - protected final Set materializedViews = Sets.newConcurrentHashSet(); - - @BeforeClass(alwaysRun = true) - public void setupClass() - throws Exception - { - executor = newCachedThreadPool(daemonThreadsNamed("hive-%s")); - heartbeatService = newScheduledThreadPool(1); - // Use separate staging directory for each test class to prevent intermittent failures coming from test parallelism - temporaryStagingDirectory = createTempDirectory("trino-staging-"); - } - - @AfterClass(alwaysRun = true) - public void tearDown() - { - if (executor != null) { - executor.shutdownNow(); - executor = null; - } - if (heartbeatService != null) { - heartbeatService.shutdownNow(); - heartbeatService = null; - } - if (temporaryStagingDirectory != null) { - try { - deleteRecursively(temporaryStagingDirectory, ALLOW_INSECURE); - } - catch (Exception e) { - log.warn(e, "Error deleting %s", temporaryStagingDirectory); - } - } - } - - protected void setupHive(String databaseName) - { - database = databaseName; - tablePartitionFormat = new SchemaTableName(database, "trino_test_partition_format"); - tableUnpartitioned = new SchemaTableName(database, "trino_test_unpartitioned"); - tablePartitionedWithNull = new SchemaTableName(database, "trino_test_partitioned_with_null"); - tableOffline = new SchemaTableName(database, "trino_test_offline"); - tableNotReadable = new SchemaTableName(database, "trino_test_not_readable"); - view = new SchemaTableName(database, "trino_test_view"); - invalidTable = new SchemaTableName(database, INVALID_TABLE); - tableBucketedStringInt = new SchemaTableName(database, "trino_test_bucketed_by_string_int"); - tableBucketedBigintBoolean = new SchemaTableName(database, "trino_test_bucketed_by_bigint_boolean"); - tableBucketedDoubleFloat = new SchemaTableName(database, "trino_test_bucketed_by_double_float"); - tablePartitionSchemaChange = new SchemaTableName(database, "trino_test_partition_schema_change"); - tablePartitionSchemaChangeNonCanonical = new SchemaTableName(database, "trino_test_partition_schema_change_non_canonical"); - tableBucketEvolution = new SchemaTableName(database, "trino_test_bucket_evolution"); - - invalidTableHandle = new HiveTableHandle(database, INVALID_TABLE, ImmutableMap.of(), ImmutableList.of(), ImmutableList.of(), Optional.empty()); - - dsColumn = createBaseColumn("ds", -1, HIVE_STRING, VARCHAR, PARTITION_KEY, Optional.empty()); - fileFormatColumn = createBaseColumn("file_format", -1, HIVE_STRING, VARCHAR, PARTITION_KEY, Optional.empty()); - dummyColumn = createBaseColumn("dummy", -1, HIVE_INT, INTEGER, PARTITION_KEY, Optional.empty()); - intColumn = createBaseColumn("t_int", -1, HIVE_INT, INTEGER, PARTITION_KEY, Optional.empty()); - invalidColumnHandle = createBaseColumn(INVALID_COLUMN, 0, HIVE_STRING, VARCHAR, REGULAR, Optional.empty()); - pStringColumn = createBaseColumn("p_string", -1, HIVE_STRING, VARCHAR, PARTITION_KEY, Optional.empty()); - pIntegerColumn = createBaseColumn("p_integer", -1, HIVE_INT, INTEGER, PARTITION_KEY, Optional.empty()); - - List partitionColumns = ImmutableList.of(dsColumn, fileFormatColumn, dummyColumn); - tablePartitionFormatPartitions = ImmutableList.builder() - .add(new HivePartition(tablePartitionFormat, - "ds=2012-12-29/file_format=textfile/dummy=1", - ImmutableMap.builder() - .put(dsColumn, NullableValue.of(createUnboundedVarcharType(), utf8Slice("2012-12-29"))) - .put(fileFormatColumn, NullableValue.of(createUnboundedVarcharType(), utf8Slice("textfile"))) - .put(dummyColumn, NullableValue.of(INTEGER, 1L)) - .buildOrThrow())) - .add(new HivePartition(tablePartitionFormat, - "ds=2012-12-29/file_format=sequencefile/dummy=2", - ImmutableMap.builder() - .put(dsColumn, NullableValue.of(createUnboundedVarcharType(), utf8Slice("2012-12-29"))) - .put(fileFormatColumn, NullableValue.of(createUnboundedVarcharType(), utf8Slice("sequencefile"))) - .put(dummyColumn, NullableValue.of(INTEGER, 2L)) - .buildOrThrow())) - .add(new HivePartition(tablePartitionFormat, - "ds=2012-12-29/file_format=rctext/dummy=3", - ImmutableMap.builder() - .put(dsColumn, NullableValue.of(createUnboundedVarcharType(), utf8Slice("2012-12-29"))) - .put(fileFormatColumn, NullableValue.of(createUnboundedVarcharType(), utf8Slice("rctext"))) - .put(dummyColumn, NullableValue.of(INTEGER, 3L)) - .buildOrThrow())) - .add(new HivePartition(tablePartitionFormat, - "ds=2012-12-29/file_format=rcbinary/dummy=4", - ImmutableMap.builder() - .put(dsColumn, NullableValue.of(createUnboundedVarcharType(), utf8Slice("2012-12-29"))) - .put(fileFormatColumn, NullableValue.of(createUnboundedVarcharType(), utf8Slice("rcbinary"))) - .put(dummyColumn, NullableValue.of(INTEGER, 4L)) - .buildOrThrow())) - .build(); - tableUnpartitionedPartitions = ImmutableList.of(new HivePartition(tableUnpartitioned)); - tablePartitionFormatProperties = new ConnectorTableProperties( - TupleDomain.withColumnDomains(ImmutableMap.of( - dsColumn, Domain.create(ValueSet.ofRanges(Range.equal(createUnboundedVarcharType(), utf8Slice("2012-12-29"))), false), - fileFormatColumn, Domain.create(ValueSet.ofRanges(Range.equal(createUnboundedVarcharType(), utf8Slice("textfile")), Range.equal(createUnboundedVarcharType(), utf8Slice("sequencefile")), Range.equal(createUnboundedVarcharType(), utf8Slice("rctext")), Range.equal(createUnboundedVarcharType(), utf8Slice("rcbinary"))), false), - dummyColumn, Domain.create(ValueSet.ofRanges(Range.equal(INTEGER, 1L), Range.equal(INTEGER, 2L), Range.equal(INTEGER, 3L), Range.equal(INTEGER, 4L)), false))), - Optional.empty(), - Optional.of(new DiscretePredicates(partitionColumns, ImmutableList.of( - TupleDomain.withColumnDomains(ImmutableMap.of( - dsColumn, Domain.create(ValueSet.ofRanges(Range.equal(createUnboundedVarcharType(), utf8Slice("2012-12-29"))), false), - fileFormatColumn, Domain.create(ValueSet.ofRanges(Range.equal(createUnboundedVarcharType(), utf8Slice("textfile"))), false), - dummyColumn, Domain.create(ValueSet.ofRanges(Range.equal(INTEGER, 1L)), false))), - TupleDomain.withColumnDomains(ImmutableMap.of( - dsColumn, Domain.create(ValueSet.ofRanges(Range.equal(createUnboundedVarcharType(), utf8Slice("2012-12-29"))), false), - fileFormatColumn, Domain.create(ValueSet.ofRanges(Range.equal(createUnboundedVarcharType(), utf8Slice("sequencefile"))), false), - dummyColumn, Domain.create(ValueSet.ofRanges(Range.equal(INTEGER, 2L)), false))), - TupleDomain.withColumnDomains(ImmutableMap.of( - dsColumn, Domain.create(ValueSet.ofRanges(Range.equal(createUnboundedVarcharType(), utf8Slice("2012-12-29"))), false), - fileFormatColumn, Domain.create(ValueSet.ofRanges(Range.equal(createUnboundedVarcharType(), utf8Slice("rctext"))), false), - dummyColumn, Domain.create(ValueSet.ofRanges(Range.equal(INTEGER, 3L)), false))), - TupleDomain.withColumnDomains(ImmutableMap.of( - dsColumn, Domain.create(ValueSet.ofRanges(Range.equal(createUnboundedVarcharType(), utf8Slice("2012-12-29"))), false), - fileFormatColumn, Domain.create(ValueSet.ofRanges(Range.equal(createUnboundedVarcharType(), utf8Slice("rcbinary"))), false), - dummyColumn, Domain.create(ValueSet.ofRanges(Range.equal(INTEGER, 4L)), false)))))), - ImmutableList.of()); - tableUnpartitionedProperties = new ConnectorTableProperties(); - } - - protected final void setup(HostAndPort metastoreAddress, String databaseName) - { - HiveConfig hiveConfig = getHiveConfig() - .setParquetTimeZone("UTC") - .setRcfileTimeZone("UTC"); - - hdfsEnvironment = HDFS_ENVIRONMENT; - HiveMetastore metastore = CachingHiveMetastore.builder() - .delegate(new BridgingHiveMetastore(testingThriftHiveMetastoreBuilder() - .metastoreClient(metastoreAddress) - .hiveConfig(hiveConfig) - .thriftMetastoreConfig(new ThriftMetastoreConfig() - .setAssumeCanonicalPartitionKeys(true)) - .hdfsEnvironment(hdfsEnvironment) - .build())) - .executor(executor) - .metadataCacheEnabled(true) - .statsCacheEnabled(true) - .cacheTtl(new Duration(1, MINUTES)) - .refreshInterval(new Duration(15, SECONDS)) - .maximumSize(10000) - .cacheMissing(new CachingHiveMetastoreConfig().isCacheMissing()) - .partitionCacheEnabled(new CachingHiveMetastoreConfig().isPartitionCacheEnabled()) - .build(); - - setup(databaseName, hiveConfig, metastore, hdfsEnvironment); - } - - protected final void setup(String databaseName, HiveConfig hiveConfig, HiveMetastore hiveMetastore, HdfsEnvironment hdfsConfiguration) - { - setupHive(databaseName); - - metastoreClient = hiveMetastore; - hdfsEnvironment = hdfsConfiguration; - HivePartitionManager partitionManager = new HivePartitionManager(hiveConfig); - locationService = new HiveLocationService(hdfsEnvironment, hiveConfig); - JsonCodec partitionUpdateCodec = JsonCodec.jsonCodec(PartitionUpdate.class); - countingDirectoryLister = new CountingDirectoryLister(); - metadataFactory = new HiveMetadataFactory( - new CatalogName("hive"), - HiveMetastoreFactory.ofInstance(metastoreClient), - getDefaultHiveFileWriterFactories(hiveConfig, hdfsEnvironment), - new HdfsFileSystemFactory(hdfsEnvironment, HDFS_FILE_SYSTEM_STATS), - hdfsEnvironment, - partitionManager, - 10, - 10, - 10, - 100_000, - false, - false, - false, - true, - true, - false, - false, - 1000, - Optional.empty(), - true, - TESTING_TYPE_MANAGER, - NOOP_METADATA_PROVIDER, - locationService, - partitionUpdateCodec, - executor, - heartbeatService, - TEST_SERVER_VERSION, - (session, tableHandle) -> { - if (!tableHandle.getTableName().contains("apply_redirection_tester")) { - return Optional.empty(); - } - return Optional.of(new TableScanRedirectApplicationResult( - new CatalogSchemaTableName("hive", databaseName, "mock_redirection_target"), - ImmutableMap.of(), - TupleDomain.all())); - }, - ImmutableSet.of( - new PartitionsSystemTableProvider(partitionManager, TESTING_TYPE_MANAGER), - new PropertiesSystemTableProvider()), - metastore -> new NoneHiveMaterializedViewMetadata() - { - @Override - public List listMaterializedViews(ConnectorSession session, Optional schemaName) - { - return materializedViews.stream() - .filter(schemaName - .>map(name -> mvName -> mvName.getSchemaName().equals(name)) - .orElse(mvName -> true)) - .collect(toImmutableList()); - } - - @Override - public Optional getMaterializedView(ConnectorSession session, SchemaTableName viewName) - { - if (!viewName.getTableName().contains("materialized_view_tester")) { - return Optional.empty(); - } - return Optional.of(new ConnectorMaterializedViewDefinition( - "dummy_view_sql", - Optional.empty(), - Optional.empty(), - Optional.empty(), - ImmutableList.of(new ConnectorMaterializedViewDefinition.Column("abc", TypeId.of("type"))), - Optional.empty(), - Optional.of("alice"), - ImmutableMap.of())); - } - }, - SqlStandardAccessControlMetadata::new, - countingDirectoryLister, - new TransactionScopeCachingDirectoryListerFactory(hiveConfig), - new PartitionProjectionService(hiveConfig, ImmutableMap.of(), new TestingTypeManager()), - true, - HiveTimestampPrecision.DEFAULT_PRECISION); - transactionManager = new HiveTransactionManager(metadataFactory); - splitManager = new HiveSplitManager( - transactionManager, - partitionManager, - new HdfsFileSystemFactory(hdfsEnvironment, HDFS_FILE_SYSTEM_STATS), - new HdfsNamenodeStats(), - hdfsEnvironment, - executor, - new CounterStat(), - 100, - hiveConfig.getMaxOutstandingSplitsSize(), - hiveConfig.getMinPartitionBatchSize(), - hiveConfig.getMaxPartitionBatchSize(), - hiveConfig.getMaxInitialSplits(), - hiveConfig.getSplitLoaderConcurrency(), - hiveConfig.getMaxSplitsPerSecond(), - false, - TESTING_TYPE_MANAGER, - hiveConfig.getMaxPartitionsPerScan()); - pageSinkProvider = new HivePageSinkProvider( - getDefaultHiveFileWriterFactories(hiveConfig, hdfsEnvironment), - new HdfsFileSystemFactory(hdfsEnvironment, HDFS_FILE_SYSTEM_STATS), - hdfsEnvironment, - PAGE_SORTER, - HiveMetastoreFactory.ofInstance(metastoreClient), - new GroupByHashPageIndexerFactory(JOIN_COMPILER, BLOCK_TYPE_OPERATORS), - TESTING_TYPE_MANAGER, - getHiveConfig(), - getSortingFileWriterConfig(), - locationService, - partitionUpdateCodec, - new TestingNodeManager("fake-environment"), - new HiveEventClient(), - getHiveSessionProperties(hiveConfig), - new HiveWriterStats()); - pageSourceProvider = new HivePageSourceProvider( - TESTING_TYPE_MANAGER, - hdfsEnvironment, - hiveConfig, - getDefaultHivePageSourceFactories(hdfsEnvironment, hiveConfig), - getDefaultHiveRecordCursorProviders(hiveConfig, hdfsEnvironment), - new GenericHiveRecordCursorProvider(hdfsEnvironment, hiveConfig)); - nodePartitioningProvider = new HiveNodePartitioningProvider( - new TestingNodeManager("fake-environment"), - TESTING_TYPE_MANAGER); - } - - /** - * Allow subclass to change default configuration. - */ - protected HiveConfig getHiveConfig() - { - return new HiveConfig() - .setTemporaryStagingDirectoryPath(temporaryStagingDirectory.resolve("temp_path_").toAbsolutePath().toString()); - } - - protected SortingFileWriterConfig getSortingFileWriterConfig() - { - return new SortingFileWriterConfig() - .setMaxOpenSortFiles(10) - .setWriterSortBufferSize(DataSize.of(100, KILOBYTE)); - } - - protected ConnectorSession newSession() - { - return newSession(ImmutableMap.of()); - } - - protected ConnectorSession newSession(Map propertyValues) - { - return TestingConnectorSession.builder() - .setPropertyMetadata(getHiveSessionProperties(getHiveConfig()).getSessionProperties()) - .setPropertyValues(propertyValues) - .build(); - } - - protected Transaction newTransaction() - { - return new HiveTransaction(transactionManager); - } - - protected interface Transaction - extends AutoCloseable - { - ConnectorMetadata getMetadata(); - - SemiTransactionalHiveMetastore getMetastore(); - - ConnectorTransactionHandle getTransactionHandle(); - - void commit(); - - void rollback(); - - @Override - void close(); - } - - static class HiveTransaction - implements Transaction - { - private final HiveTransactionManager transactionManager; - private final ConnectorTransactionHandle transactionHandle; - private boolean closed; - - public HiveTransaction(HiveTransactionManager transactionManager) - { - this.transactionManager = requireNonNull(transactionManager, "transactionManager is null"); - this.transactionHandle = new HiveTransactionHandle(false); - transactionManager.begin(transactionHandle); - getMetastore().testOnlyThrowOnCleanupFailures(); - } - - @Override - public ConnectorMetadata getMetadata() - { - return transactionManager.get(transactionHandle, SESSION.getIdentity()); - } - - @Override - public SemiTransactionalHiveMetastore getMetastore() - { - return transactionManager.get(transactionHandle, SESSION.getIdentity()).getMetastore(); - } - - @Override - public ConnectorTransactionHandle getTransactionHandle() - { - return transactionHandle; - } - - @Override - public void commit() - { - checkState(!closed); - closed = true; - transactionManager.commit(transactionHandle); - } - - @Override - public void rollback() - { - checkState(!closed); - closed = true; - transactionManager.rollback(transactionHandle); - } - - @Override - public void close() - { - if (!closed) { - try { - getMetastore().testOnlyCheckIsReadOnly(); // transactions in this test with writes in it must explicitly commit or rollback - } - finally { - rollback(); - } - } - } - } - - @Test - public void testGetDatabaseNames() - { - try (Transaction transaction = newTransaction()) { - ConnectorMetadata metadata = transaction.getMetadata(); - List databases = metadata.listSchemaNames(newSession()); - assertTrue(databases.contains(database)); - } - } - - @Test - public void testGetTableNames() - { - try (Transaction transaction = newTransaction()) { - ConnectorMetadata metadata = transaction.getMetadata(); - List tables = metadata.listTables(newSession(), Optional.of(database)); - assertTrue(tables.contains(tablePartitionFormat)); - assertTrue(tables.contains(tableUnpartitioned)); - } - } - - @Test - public void testGetAllTableNames() - { - try (Transaction transaction = newTransaction()) { - ConnectorMetadata metadata = transaction.getMetadata(); - List tables = metadata.listTables(newSession(), Optional.empty()); - assertTrue(tables.contains(tablePartitionFormat)); - assertTrue(tables.contains(tableUnpartitioned)); - } - } - - @Test - public void testGetAllTableColumns() - { - try (Transaction transaction = newTransaction()) { - ConnectorMetadata metadata = transaction.getMetadata(); - Map> allColumns = listTableColumns(metadata, newSession(), new SchemaTablePrefix()); - assertTrue(allColumns.containsKey(tablePartitionFormat)); - assertTrue(allColumns.containsKey(tableUnpartitioned)); - } - } - - @Test - public void testGetAllTableColumnsInSchema() - { - try (Transaction transaction = newTransaction()) { - ConnectorMetadata metadata = transaction.getMetadata(); - Map> allColumns = listTableColumns(metadata, newSession(), new SchemaTablePrefix(database)); - assertTrue(allColumns.containsKey(tablePartitionFormat)); - assertTrue(allColumns.containsKey(tableUnpartitioned)); - } - } - - @Test - public void testListUnknownSchema() - { - try (Transaction transaction = newTransaction()) { - ConnectorMetadata metadata = transaction.getMetadata(); - ConnectorSession session = newSession(); - assertNull(metadata.getTableHandle(session, new SchemaTableName(INVALID_DATABASE, INVALID_TABLE))); - assertEquals(metadata.listTables(session, Optional.of(INVALID_DATABASE)), ImmutableList.of()); - assertEquals(listTableColumns(metadata, session, new SchemaTablePrefix(INVALID_DATABASE, INVALID_TABLE)), ImmutableMap.of()); - assertEquals(metadata.listViews(session, Optional.of(INVALID_DATABASE)), ImmutableList.of()); - assertEquals(metadata.getViews(session, Optional.of(INVALID_DATABASE)), ImmutableMap.of()); - assertEquals(metadata.getView(session, new SchemaTableName(INVALID_DATABASE, INVALID_TABLE)), Optional.empty()); - } - } - - @Test - public void testGetPartitions() - throws Exception - { - try (Transaction transaction = newTransaction()) { - ConnectorMetadata metadata = transaction.getMetadata(); - ConnectorTableHandle tableHandle = getTableHandle(metadata, tablePartitionFormat); - tableHandle = applyFilter(metadata, tableHandle, Constraint.alwaysTrue()); - ConnectorTableProperties properties = metadata.getTableProperties(newSession(), tableHandle); - assertExpectedTableProperties(properties, tablePartitionFormatProperties); - assertExpectedPartitions(tableHandle, tablePartitionFormatPartitions); - } - } - - @Test - public void testGetPartitionsWithBindings() - { - try (Transaction transaction = newTransaction()) { - ConnectorMetadata metadata = transaction.getMetadata(); - ConnectorTableHandle tableHandle = getTableHandle(metadata, tablePartitionFormat); - Constraint constraint = new Constraint(TupleDomain.withColumnDomains(ImmutableMap.of(intColumn, Domain.singleValue(BIGINT, 5L)))); - tableHandle = applyFilter(metadata, tableHandle, constraint); - ConnectorTableProperties properties = metadata.getTableProperties(newSession(), tableHandle); - assertExpectedTableProperties(properties, tablePartitionFormatProperties); - assertExpectedPartitions(tableHandle, tablePartitionFormatPartitions); - } - } - - @Test - public void testGetPartitionsWithFilter() - { - try (Transaction transaction = newTransaction()) { - ConnectorMetadata metadata = transaction.getMetadata(); - ConnectorTableHandle tableHandle = getTableHandle(metadata, tablePartitionedWithNull); - - Domain varcharSomeValue = Domain.singleValue(VARCHAR, utf8Slice("abc")); - Domain varcharOnlyNull = Domain.onlyNull(VARCHAR); - Domain varcharNotNull = Domain.notNull(VARCHAR); - - Domain integerSomeValue = Domain.singleValue(INTEGER, 123L); - Domain integerOnlyNull = Domain.onlyNull(INTEGER); - Domain integerNotNull = Domain.notNull(INTEGER); - - // all - assertThat(getPartitionNamesByFilter(metadata, tableHandle, new Constraint(TupleDomain.all()))) - .containsOnly( - "p_string=__HIVE_DEFAULT_PARTITION__/p_integer=__HIVE_DEFAULT_PARTITION__", - "p_string=abc/p_integer=123", - "p_string=def/p_integer=456"); - - // is some value - assertThat(getPartitionNamesByFilter(metadata, tableHandle, pStringColumn, varcharSomeValue)) - .containsOnly("p_string=abc/p_integer=123"); - assertThat(getPartitionNamesByFilter(metadata, tableHandle, pIntegerColumn, integerSomeValue)) - .containsOnly("p_string=abc/p_integer=123"); - - // IS NULL - assertThat(getPartitionNamesByFilter(metadata, tableHandle, pStringColumn, varcharOnlyNull)) - .containsOnly("p_string=__HIVE_DEFAULT_PARTITION__/p_integer=__HIVE_DEFAULT_PARTITION__"); - assertThat(getPartitionNamesByFilter(metadata, tableHandle, pIntegerColumn, integerOnlyNull)) - .containsOnly("p_string=__HIVE_DEFAULT_PARTITION__/p_integer=__HIVE_DEFAULT_PARTITION__"); - - // IS NOT NULL - assertThat(getPartitionNamesByFilter(metadata, tableHandle, pStringColumn, varcharNotNull)) - .containsOnly("p_string=abc/p_integer=123", "p_string=def/p_integer=456"); - assertThat(getPartitionNamesByFilter(metadata, tableHandle, pIntegerColumn, integerNotNull)) - .containsOnly("p_string=abc/p_integer=123", "p_string=def/p_integer=456"); - - // IS NULL OR is some value - assertThat(getPartitionNamesByFilter(metadata, tableHandle, pStringColumn, varcharOnlyNull.union(varcharSomeValue))) - .containsOnly("p_string=__HIVE_DEFAULT_PARTITION__/p_integer=__HIVE_DEFAULT_PARTITION__", "p_string=abc/p_integer=123"); - assertThat(getPartitionNamesByFilter(metadata, tableHandle, pIntegerColumn, integerOnlyNull.union(integerSomeValue))) - .containsOnly("p_string=__HIVE_DEFAULT_PARTITION__/p_integer=__HIVE_DEFAULT_PARTITION__", "p_string=abc/p_integer=123"); - - // IS NOT NULL AND is NOT some value - assertThat(getPartitionNamesByFilter(metadata, tableHandle, pStringColumn, varcharSomeValue.complement().intersect(varcharNotNull))) - .containsOnly("p_string=def/p_integer=456"); - assertThat(getPartitionNamesByFilter(metadata, tableHandle, pIntegerColumn, integerSomeValue.complement().intersect(integerNotNull))) - .containsOnly("p_string=def/p_integer=456"); - - // IS NULL OR is NOT some value - assertThat(getPartitionNamesByFilter(metadata, tableHandle, pStringColumn, varcharSomeValue.complement())) - .containsOnly("p_string=__HIVE_DEFAULT_PARTITION__/p_integer=__HIVE_DEFAULT_PARTITION__", "p_string=def/p_integer=456"); - assertThat(getPartitionNamesByFilter(metadata, tableHandle, pIntegerColumn, integerSomeValue.complement())) - .containsOnly("p_string=__HIVE_DEFAULT_PARTITION__/p_integer=__HIVE_DEFAULT_PARTITION__", "p_string=def/p_integer=456"); - } - } - - private Set getPartitionNamesByFilter(ConnectorMetadata metadata, ConnectorTableHandle tableHandle, ColumnHandle columnHandle, Domain domain) - { - return getPartitionNamesByFilter(metadata, tableHandle, new Constraint(TupleDomain.withColumnDomains(ImmutableMap.of(columnHandle, domain)))); - } - - private Set getPartitionNamesByFilter(ConnectorMetadata metadata, ConnectorTableHandle tableHandle, Constraint constraint) - { - return applyFilter(metadata, tableHandle, constraint) - .getPartitions().orElseThrow(() -> new IllegalStateException("No partitions")) - .stream() - .map(HivePartition::getPartitionId) - .collect(toImmutableSet()); - } - - @Test - public void testMismatchSchemaTable() - throws Exception - { - for (HiveStorageFormat storageFormat : createTableFormats) { - // TODO: fix coercion for JSON - if (storageFormat == JSON) { - continue; - } - SchemaTableName temporaryMismatchSchemaTable = temporaryTable("mismatch_schema"); - try { - doTestMismatchSchemaTable( - temporaryMismatchSchemaTable, - storageFormat, - MISMATCH_SCHEMA_TABLE_BEFORE, - MISMATCH_SCHEMA_TABLE_DATA_BEFORE, - MISMATCH_SCHEMA_TABLE_AFTER, - MISMATCH_SCHEMA_TABLE_DATA_AFTER); - } - finally { - dropTable(temporaryMismatchSchemaTable); - } - } - } - - protected void doTestMismatchSchemaTable( - SchemaTableName schemaTableName, - HiveStorageFormat storageFormat, - List tableBefore, - MaterializedResult dataBefore, - List tableAfter, - MaterializedResult dataAfter) - throws Exception - { - String schemaName = schemaTableName.getSchemaName(); - String tableName = schemaTableName.getTableName(); - - doCreateEmptyTable(schemaTableName, storageFormat, tableBefore); - - // insert the data - try (Transaction transaction = newTransaction()) { - ConnectorSession session = newSession(); - ConnectorMetadata metadata = transaction.getMetadata(); - ConnectorTableHandle tableHandle = getTableHandle(metadata, schemaTableName); - - ConnectorInsertTableHandle insertTableHandle = metadata.beginInsert(session, tableHandle, ImmutableList.of(), NO_RETRIES); - ConnectorPageSink sink = pageSinkProvider.createPageSink(transaction.getTransactionHandle(), session, insertTableHandle, TESTING_PAGE_SINK_ID); - sink.appendPage(dataBefore.toPage()); - Collection fragments = getFutureValue(sink.finish()); - - metadata.finishInsert(session, insertTableHandle, fragments, ImmutableList.of()); - - transaction.commit(); - } - - // load the table and verify the data - try (Transaction transaction = newTransaction()) { - ConnectorSession session = newSession(); - ConnectorMetadata metadata = transaction.getMetadata(); - metadata.beginQuery(session); - ConnectorTableHandle tableHandle = getTableHandle(metadata, schemaTableName); - - List columnHandles = metadata.getColumnHandles(session, tableHandle).values().stream() - .filter(columnHandle -> !((HiveColumnHandle) columnHandle).isHidden()) - .collect(toList()); - - MaterializedResult result = readTable(transaction, tableHandle, columnHandles, session, TupleDomain.all(), OptionalInt.empty(), Optional.empty()); - assertEqualsIgnoreOrder(result.getMaterializedRows(), dataBefore.getMaterializedRows()); - transaction.commit(); - } - - // alter the table schema - try (Transaction transaction = newTransaction()) { - ConnectorSession session = newSession(); - PrincipalPrivileges principalPrivileges = testingPrincipalPrivilege(session); - Table oldTable = transaction.getMetastore().getTable(schemaName, tableName).get(); - List dataColumns = tableAfter.stream() - .filter(columnMetadata -> !columnMetadata.getName().equals("ds")) - .map(columnMetadata -> new Column(columnMetadata.getName(), toHiveType(columnMetadata.getType()), Optional.empty())) - .collect(toList()); - Table.Builder newTable = Table.builder(oldTable) - .setDataColumns(dataColumns); - - transaction.getMetastore().replaceTable(schemaName, tableName, newTable.build(), principalPrivileges); - - transaction.commit(); - } - - // load the altered table and verify the data - try (Transaction transaction = newTransaction()) { - ConnectorSession session = newSession(); - ConnectorMetadata metadata = transaction.getMetadata(); - metadata.beginQuery(session); - ConnectorTableHandle tableHandle = getTableHandle(metadata, schemaTableName); - List columnHandles = metadata.getColumnHandles(session, tableHandle).values().stream() - .filter(columnHandle -> !((HiveColumnHandle) columnHandle).isHidden()) - .collect(toList()); - - MaterializedResult result = readTable(transaction, tableHandle, columnHandles, session, TupleDomain.all(), OptionalInt.empty(), Optional.empty()); - assertEqualsIgnoreOrder(result.getMaterializedRows(), dataAfter.getMaterializedRows()); - - transaction.commit(); - } - - // insertions to the partitions with type mismatches should fail - try (Transaction transaction = newTransaction()) { - ConnectorSession session = newSession(); - ConnectorMetadata metadata = transaction.getMetadata(); - ConnectorTableHandle tableHandle = getTableHandle(metadata, schemaTableName); - - ConnectorInsertTableHandle insertTableHandle = metadata.beginInsert(session, tableHandle, ImmutableList.of(), NO_RETRIES); - ConnectorPageSink sink = pageSinkProvider.createPageSink(transaction.getTransactionHandle(), session, insertTableHandle, TESTING_PAGE_SINK_ID); - sink.appendPage(dataAfter.toPage()); - Collection fragments = getFutureValue(sink.finish()); - - metadata.finishInsert(session, insertTableHandle, fragments, ImmutableList.of()); - - transaction.commit(); - - fail("expected exception"); - } - catch (TrinoException e) { - // expected - assertEquals(e.getErrorCode(), HIVE_PARTITION_SCHEMA_MISMATCH.toErrorCode()); - } - } - - protected void assertExpectedTableProperties(ConnectorTableProperties actualProperties, ConnectorTableProperties expectedProperties) - { - assertEquals(actualProperties.getPredicate(), expectedProperties.getPredicate()); - assertEquals(actualProperties.getDiscretePredicates().isPresent(), expectedProperties.getDiscretePredicates().isPresent()); - actualProperties.getDiscretePredicates().ifPresent(actual -> { - DiscretePredicates expected = expectedProperties.getDiscretePredicates().get(); - assertEquals(actual.getColumns(), expected.getColumns()); - assertEqualsIgnoreOrder(actual.getPredicates(), expected.getPredicates()); - }); - assertEquals(actualProperties.getLocalProperties(), expectedProperties.getLocalProperties()); - } - - protected void assertExpectedPartitions(ConnectorTableHandle table, Iterable expectedPartitions) - { - Iterable actualPartitions = ((HiveTableHandle) table).getPartitions().orElseThrow(AssertionError::new); - Map actualById = uniqueIndex(actualPartitions, HivePartition::getPartitionId); - Map expectedById = uniqueIndex(expectedPartitions, HivePartition::getPartitionId); - - assertThat(actualById).isEqualTo(expectedById); - - // HivePartition.equals doesn't compare all the fields, so let's check them - for (Map.Entry expected : expectedById.entrySet()) { - HivePartition actualPartition = actualById.get(expected.getKey()); - HivePartition expectedPartition = expected.getValue(); - assertEquals(actualPartition.getPartitionId(), expectedPartition.getPartitionId()); - assertEquals(actualPartition.getKeys(), expectedPartition.getKeys()); - assertEquals(actualPartition.getTableName(), expectedPartition.getTableName()); - } - } - - @Test - public void testGetPartitionNamesUnpartitioned() - { - try (Transaction transaction = newTransaction()) { - ConnectorMetadata metadata = transaction.getMetadata(); - ConnectorTableHandle tableHandle = getTableHandle(metadata, tableUnpartitioned); - tableHandle = applyFilter(metadata, tableHandle, Constraint.alwaysTrue()); - ConnectorTableProperties properties = metadata.getTableProperties(newSession(), tableHandle); - assertExpectedTableProperties(properties, new ConnectorTableProperties()); - assertExpectedPartitions(tableHandle, tableUnpartitionedPartitions); - } - } - - @Test - public void testGetTableSchemaPartitionFormat() - { - try (Transaction transaction = newTransaction()) { - ConnectorMetadata metadata = transaction.getMetadata(); - ConnectorTableMetadata tableMetadata = metadata.getTableMetadata(newSession(), getTableHandle(metadata, tablePartitionFormat)); - Map map = uniqueIndex(tableMetadata.getColumns(), ColumnMetadata::getName); - - assertPrimitiveField(map, "t_string", createUnboundedVarcharType(), false); - assertPrimitiveField(map, "t_tinyint", TINYINT, false); - assertPrimitiveField(map, "t_smallint", SMALLINT, false); - assertPrimitiveField(map, "t_int", INTEGER, false); - assertPrimitiveField(map, "t_bigint", BIGINT, false); - assertPrimitiveField(map, "t_float", REAL, false); - assertPrimitiveField(map, "t_double", DOUBLE, false); - assertPrimitiveField(map, "t_boolean", BOOLEAN, false); - assertPrimitiveField(map, "ds", createUnboundedVarcharType(), true); - assertPrimitiveField(map, "file_format", createUnboundedVarcharType(), true); - assertPrimitiveField(map, "dummy", INTEGER, true); - } - } - - @Test - public void testGetTableSchemaUnpartitioned() - { - try (Transaction transaction = newTransaction()) { - ConnectorMetadata metadata = transaction.getMetadata(); - ConnectorTableHandle tableHandle = getTableHandle(metadata, tableUnpartitioned); - ConnectorTableMetadata tableMetadata = metadata.getTableMetadata(newSession(), tableHandle); - Map map = uniqueIndex(tableMetadata.getColumns(), ColumnMetadata::getName); - - assertPrimitiveField(map, "t_string", createUnboundedVarcharType(), false); - assertPrimitiveField(map, "t_tinyint", TINYINT, false); - } - } - - @Test - public void testGetTableSchemaOffline() - { - try (Transaction transaction = newTransaction()) { - ConnectorMetadata metadata = transaction.getMetadata(); - Map> columns = listTableColumns(metadata, newSession(), tableOffline.toSchemaTablePrefix()); - assertEquals(columns.size(), 1); - Map map = uniqueIndex(getOnlyElement(columns.values()), ColumnMetadata::getName); - - assertPrimitiveField(map, "t_string", createUnboundedVarcharType(), false); - } - } - - @Test - public void testGetTableSchemaNotReadablePartition() - { - try (Transaction transaction = newTransaction()) { - ConnectorMetadata metadata = transaction.getMetadata(); - ConnectorTableHandle tableHandle = getTableHandle(metadata, tableNotReadable); - ConnectorTableMetadata tableMetadata = metadata.getTableMetadata(newSession(), tableHandle); - Map map = uniqueIndex(tableMetadata.getColumns(), ColumnMetadata::getName); - - assertPrimitiveField(map, "t_string", createUnboundedVarcharType(), false); - } - } - - @Test - public void testGetTableSchemaException() - { - try (Transaction transaction = newTransaction()) { - ConnectorMetadata metadata = transaction.getMetadata(); - assertNull(metadata.getTableHandle(newSession(), invalidTable)); - } - } - - @Test - public void testGetTableStatsBucketedStringInt() - { - assertTableStatsComputed( - tableBucketedStringInt, - ImmutableSet.of( - "t_bigint", - "t_boolean", - "t_double", - "t_float", - "t_int", - "t_smallint", - "t_string", - "t_tinyint", - "ds")); - } - - @Test - public void testGetTableStatsUnpartitioned() - { - assertTableStatsComputed( - tableUnpartitioned, - ImmutableSet.of("t_string", "t_tinyint")); - } - - private void assertTableStatsComputed( - SchemaTableName tableName, - Set expectedColumnStatsColumns) - { - try (Transaction transaction = newTransaction()) { - ConnectorMetadata metadata = transaction.getMetadata(); - ConnectorSession session = newSession(); - ConnectorTableHandle tableHandle = getTableHandle(metadata, tableName); - - // first check if table handle with only one projected column will return this column stats - String firstColumnName = expectedColumnStatsColumns.iterator().next(); - verifyTableStatisticsWithColumns(metadata, session, applyProjection(metadata, session, tableHandle, firstColumnName), ImmutableSet.of(firstColumnName)); - - verifyTableStatisticsWithColumns(metadata, session, tableHandle, expectedColumnStatsColumns); - } - } - - private static ConnectorTableHandle applyProjection(ConnectorMetadata metadata, ConnectorSession session, ConnectorTableHandle tableHandle, String columnName) - { - Map columnHandles = metadata.getColumnHandles(session, tableHandle); - HiveColumnHandle firstColumn = (HiveColumnHandle) columnHandles.get(columnName); - return metadata.applyProjection( - session, - tableHandle, - ImmutableList.of(new Variable("c1", firstColumn.getBaseType())), - ImmutableMap.of("c1", firstColumn)) - .orElseThrow() - .getHandle(); - } - - private static void verifyTableStatisticsWithColumns( - ConnectorMetadata metadata, - ConnectorSession session, - ConnectorTableHandle tableHandle, - Set expectedColumnStatsColumns) - { - TableStatistics tableStatistics = metadata.getTableStatistics(session, tableHandle); - - assertFalse(tableStatistics.getRowCount().isUnknown(), "row count is unknown"); - - Map columnsStatistics = tableStatistics - .getColumnStatistics() - .entrySet() - .stream() - .collect( - toImmutableMap( - entry -> ((HiveColumnHandle) entry.getKey()).getName(), - Map.Entry::getValue)); - - assertEquals(columnsStatistics.keySet(), expectedColumnStatsColumns, "columns with statistics"); - - Map columnHandles = metadata.getColumnHandles(session, tableHandle); - columnsStatistics.forEach((columnName, columnStatistics) -> { - ColumnHandle columnHandle = columnHandles.get(columnName); - Type columnType = metadata.getColumnMetadata(session, tableHandle, columnHandle).getType(); - - assertFalse( - columnStatistics.getNullsFraction().isUnknown(), - "unknown nulls fraction for " + columnName); - - assertFalse( - columnStatistics.getDistinctValuesCount().isUnknown(), - "unknown distinct values count for " + columnName); - - if (columnType instanceof VarcharType) { - assertFalse( - columnStatistics.getDataSize().isUnknown(), - "unknown data size for " + columnName); - } - else { - assertTrue( - columnStatistics.getDataSize().isUnknown(), - "unknown data size for" + columnName); - } - }); - } - - @Test - public void testGetPartitionSplitsBatch() - { - try (Transaction transaction = newTransaction()) { - ConnectorMetadata metadata = transaction.getMetadata(); - ConnectorSession session = newSession(); - metadata.beginQuery(session); - - ConnectorTableHandle tableHandle = getTableHandle(metadata, tablePartitionFormat); - ConnectorSplitSource splitSource = getSplits(splitManager, transaction, session, tableHandle); - - assertEquals(getSplitCount(splitSource), tablePartitionFormatPartitions.size()); - } - } - - @Test - public void testGetPartitionSplitsBatchUnpartitioned() - { - try (Transaction transaction = newTransaction()) { - ConnectorMetadata metadata = transaction.getMetadata(); - ConnectorSession session = newSession(); - metadata.beginQuery(session); - - ConnectorTableHandle tableHandle = getTableHandle(metadata, tableUnpartitioned); - ConnectorSplitSource splitSource = getSplits(splitManager, transaction, session, tableHandle); - - assertEquals(getSplitCount(splitSource), 1); - } - } - - @Test - public void testPerTransactionDirectoryListerCache() - throws Exception - { - long initListCount = countingDirectoryLister.getListCount(); - SchemaTableName tableName = temporaryTable("per_transaction_listing_cache_test"); - List columns = ImmutableList.of(new Column("test", HIVE_STRING, Optional.empty())); - createEmptyTable(tableName, ORC, columns, ImmutableList.of()); - try { - try (Transaction transaction = newTransaction()) { - ConnectorMetadata metadata = transaction.getMetadata(); - ConnectorSession session = newSession(); - metadata.beginQuery(session); - - ConnectorTableHandle tableHandle = getTableHandle(metadata, tableName); - getAllSplits(getSplits(splitManager, transaction, session, tableHandle)); - - // directory should be listed initially - assertEquals(countingDirectoryLister.getListCount(), initListCount + 1); - - // directory content should be cached - getAllSplits(getSplits(splitManager, transaction, session, tableHandle)); - assertEquals(countingDirectoryLister.getListCount(), initListCount + 1); - } - - try (Transaction transaction = newTransaction()) { - ConnectorMetadata metadata = transaction.getMetadata(); - ConnectorSession session = newSession(); - metadata.beginQuery(session); - - ConnectorTableHandle tableHandle = getTableHandle(metadata, tableName); - getAllSplits(getSplits(splitManager, transaction, session, tableHandle)); - - // directory should be listed again in new transaction - assertEquals(countingDirectoryLister.getListCount(), initListCount + 2); - } - } - finally { - dropTable(tableName); - } - } - - @Test(expectedExceptions = TableNotFoundException.class) - public void testGetPartitionSplitsBatchInvalidTable() - { - try (Transaction transaction = newTransaction()) { - getSplits(splitManager, transaction, newSession(), invalidTableHandle); - } - } - - @Test - public void testGetPartitionTableOffline() - { - try (Transaction transaction = newTransaction()) { - ConnectorMetadata metadata = transaction.getMetadata(); - try { - getTableHandle(metadata, tableOffline); - fail("expected TableOfflineException"); - } - catch (TableOfflineException e) { - assertEquals(e.getTableName(), tableOffline); - } - } - } - - @Test - public void testGetPartitionSplitsTableNotReadablePartition() - { - try (Transaction transaction = newTransaction()) { - ConnectorMetadata metadata = transaction.getMetadata(); - ConnectorSession session = newSession(); - - ConnectorTableHandle tableHandle = getTableHandle(metadata, tableNotReadable); - assertNotNull(tableHandle); - - try { - getSplitCount(getSplits(splitManager, transaction, session, tableHandle)); - fail("Expected HiveNotReadableException"); - } - catch (HiveNotReadableException e) { - assertThat(e).hasMessageMatching("Table '.*\\.trino_test_not_readable' is not readable: reason for not readable"); - assertEquals(e.getTableName(), tableNotReadable); - assertEquals(e.getPartition(), Optional.empty()); - } - } - } - - @Test - public void testBucketedTableStringInt() - throws Exception - { - try (Transaction transaction = newTransaction()) { - ConnectorMetadata metadata = transaction.getMetadata(); - ConnectorSession session = newSession(); - metadata.beginQuery(session); - - ConnectorTableHandle tableHandle = getTableHandle(metadata, tableBucketedStringInt); - List columnHandles = ImmutableList.copyOf(metadata.getColumnHandles(session, tableHandle).values()); - Map columnIndex = indexColumns(columnHandles); - - assertTableIsBucketed(tableHandle, transaction, session); - - String testString = "test"; - Integer testInt = 13; - Short testSmallint = 12; - - // Reverse the order of bindings as compared to bucketing order - ImmutableMap bindings = ImmutableMap.builder() - .put(columnHandles.get(columnIndex.get("t_int")), NullableValue.of(INTEGER, (long) testInt)) - .put(columnHandles.get(columnIndex.get("t_string")), NullableValue.of(createUnboundedVarcharType(), utf8Slice(testString))) - .put(columnHandles.get(columnIndex.get("t_smallint")), NullableValue.of(SMALLINT, (long) testSmallint)) - .buildOrThrow(); - - MaterializedResult result = readTable(transaction, tableHandle, columnHandles, session, TupleDomain.fromFixedValues(bindings), OptionalInt.of(1), Optional.empty()); - - boolean rowFound = false; - for (MaterializedRow row : result) { - if (testString.equals(row.getField(columnIndex.get("t_string"))) && - testInt.equals(row.getField(columnIndex.get("t_int"))) && - testSmallint.equals(row.getField(columnIndex.get("t_smallint")))) { - rowFound = true; - } - } - assertTrue(rowFound); - } - } - - @SuppressWarnings("ConstantConditions") - @Test - public void testBucketedTableBigintBoolean() - throws Exception - { - try (Transaction transaction = newTransaction()) { - ConnectorMetadata metadata = transaction.getMetadata(); - ConnectorSession session = newSession(); - metadata.beginQuery(session); - - ConnectorTableHandle tableHandle = getTableHandle(metadata, tableBucketedBigintBoolean); - List columnHandles = ImmutableList.copyOf(metadata.getColumnHandles(session, tableHandle).values()); - Map columnIndex = indexColumns(columnHandles); - - assertTableIsBucketed(tableHandle, transaction, session); - ConnectorTableProperties properties = metadata.getTableProperties( - newSession(ImmutableMap.of("propagate_table_scan_sorting_properties", true)), - tableHandle); - // trino_test_bucketed_by_bigint_boolean does not define sorting, therefore local properties is empty - assertTrue(properties.getLocalProperties().isEmpty()); - assertTrue(metadata.getTableProperties(newSession(), tableHandle).getLocalProperties().isEmpty()); - - String testString = "test"; - Long testBigint = 89L; - Boolean testBoolean = true; - - ImmutableMap bindings = ImmutableMap.builder() - .put(columnHandles.get(columnIndex.get("t_string")), NullableValue.of(createUnboundedVarcharType(), utf8Slice(testString))) - .put(columnHandles.get(columnIndex.get("t_bigint")), NullableValue.of(BIGINT, testBigint)) - .put(columnHandles.get(columnIndex.get("t_boolean")), NullableValue.of(BOOLEAN, testBoolean)) - .buildOrThrow(); - - MaterializedResult result = readTable(transaction, tableHandle, columnHandles, session, TupleDomain.fromFixedValues(bindings), OptionalInt.of(1), Optional.empty()); - - boolean rowFound = false; - for (MaterializedRow row : result) { - if (testString.equals(row.getField(columnIndex.get("t_string"))) && - testBigint.equals(row.getField(columnIndex.get("t_bigint"))) && - testBoolean.equals(row.getField(columnIndex.get("t_boolean")))) { - rowFound = true; - break; - } - } - assertTrue(rowFound); - } - } - - @Test - public void testBucketedTableDoubleFloat() - throws Exception - { - try (Transaction transaction = newTransaction()) { - ConnectorMetadata metadata = transaction.getMetadata(); - ConnectorSession session = newSession(); - metadata.beginQuery(session); - - ConnectorTableHandle tableHandle = getTableHandle(metadata, tableBucketedDoubleFloat); - List columnHandles = ImmutableList.copyOf(metadata.getColumnHandles(session, tableHandle).values()); - Map columnIndex = indexColumns(columnHandles); - - assertTableIsBucketed(tableHandle, transaction, session); - - float testFloatValue = 87.1f; - double testDoubleValue = 88.2; - - ImmutableMap bindings = ImmutableMap.builder() - .put(columnHandles.get(columnIndex.get("t_float")), NullableValue.of(REAL, (long) floatToRawIntBits(testFloatValue))) - .put(columnHandles.get(columnIndex.get("t_double")), NullableValue.of(DOUBLE, testDoubleValue)) - .buildOrThrow(); - - MaterializedResult result = readTable(transaction, tableHandle, columnHandles, session, TupleDomain.fromFixedValues(bindings), OptionalInt.of(1), Optional.empty()); - assertThat(result).anyMatch(row -> testFloatValue == (float) row.getField(columnIndex.get("t_float")) - && testDoubleValue == (double) row.getField(columnIndex.get("t_double"))); - } - } - - @Test - public void testBucketedTableEvolutionWithDifferentReadBucketCount() - throws Exception - { - for (HiveStorageFormat storageFormat : createTableFormats) { - SchemaTableName temporaryBucketEvolutionTable = temporaryTable("bucket_evolution"); - try { - doTestBucketedTableEvolutionWithDifferentReadCount(storageFormat, temporaryBucketEvolutionTable); - } - finally { - dropTable(temporaryBucketEvolutionTable); - } - } - } - - private void doTestBucketedTableEvolutionWithDifferentReadCount(HiveStorageFormat storageFormat, SchemaTableName tableName) - throws Exception - { - int rowCount = 100; - int bucketCount = 16; - - // Produce a table with a partition with bucket count different but compatible with the table bucket count - createEmptyTable( - tableName, - storageFormat, - ImmutableList.of( - new Column("id", HIVE_LONG, Optional.empty()), - new Column("name", HIVE_STRING, Optional.empty())), - ImmutableList.of(new Column("pk", HIVE_STRING, Optional.empty())), - Optional.of(new HiveBucketProperty(ImmutableList.of("id"), BUCKETING_V1, 4, ImmutableList.of()))); - // write a 4-bucket partition - MaterializedResult.Builder bucket8Builder = MaterializedResult.resultBuilder(SESSION, BIGINT, VARCHAR, VARCHAR); - IntStream.range(0, rowCount).forEach(i -> bucket8Builder.row((long) i, String.valueOf(i), "four")); - insertData(tableName, bucket8Builder.build()); - - // Alter the bucket count to 16 - alterBucketProperty(tableName, Optional.of(new HiveBucketProperty(ImmutableList.of("id"), BUCKETING_V1, bucketCount, ImmutableList.of()))); - - MaterializedResult result; - try (Transaction transaction = newTransaction()) { - ConnectorMetadata metadata = transaction.getMetadata(); - ConnectorSession session = newSession(); - metadata.beginQuery(session); - - ConnectorTableHandle tableHandle = getTableHandle(metadata, tableName); - - // read entire table - List columnHandles = ImmutableList.copyOf(metadata.getColumnHandles(session, tableHandle).values()); - - List splits = getAllSplits(getSplits(splitManager, transaction, session, tableHandle)); - assertEquals(splits.size(), 16); - - ImmutableList.Builder allRows = ImmutableList.builder(); - for (ConnectorSplit split : splits) { - try (ConnectorPageSource pageSource = pageSourceProvider.createPageSource(transaction.getTransactionHandle(), session, split, tableHandle, columnHandles, DynamicFilter.EMPTY)) { - MaterializedResult intermediateResult = materializeSourceDataStream(session, pageSource, getTypes(columnHandles)); - allRows.addAll(intermediateResult.getMaterializedRows()); - } - } - result = new MaterializedResult(allRows.build(), getTypes(columnHandles)); - - assertEquals(result.getRowCount(), rowCount); - - Map columnIndex = indexColumns(columnHandles); - int nameColumnIndex = columnIndex.get("name"); - int bucketColumnIndex = columnIndex.get(BUCKET_COLUMN_NAME); - for (MaterializedRow row : result.getMaterializedRows()) { - String name = (String) row.getField(nameColumnIndex); - int bucket = (int) row.getField(bucketColumnIndex); - - assertEquals(bucket, Integer.parseInt(name) % bucketCount); - } - } - } - - @Test - public void testBucketedTableEvolution() - throws Exception - { - for (HiveStorageFormat storageFormat : createTableFormats) { - SchemaTableName temporaryBucketEvolutionTable = temporaryTable("bucket_evolution"); - try { - doTestBucketedTableEvolution(storageFormat, temporaryBucketEvolutionTable); - } - finally { - dropTable(temporaryBucketEvolutionTable); - } - } - } - - private void doTestBucketedTableEvolution(HiveStorageFormat storageFormat, SchemaTableName tableName) - throws Exception - { - int rowCount = 100; - - // - // Produce a table with 8 buckets. - // The table has 3 partitions of 3 different bucket count (4, 8, 16). - createEmptyTable( - tableName, - storageFormat, - ImmutableList.of( - new Column("id", HIVE_LONG, Optional.empty()), - new Column("name", HIVE_STRING, Optional.empty())), - ImmutableList.of(new Column("pk", HIVE_STRING, Optional.empty())), - Optional.of(new HiveBucketProperty(ImmutableList.of("id"), BUCKETING_V1, 4, ImmutableList.of()))); - // write a 4-bucket partition - MaterializedResult.Builder bucket4Builder = MaterializedResult.resultBuilder(SESSION, BIGINT, VARCHAR, VARCHAR); - IntStream.range(0, rowCount).forEach(i -> bucket4Builder.row((long) i, String.valueOf(i), "four")); - insertData(tableName, bucket4Builder.build()); - // write a 16-bucket partition - alterBucketProperty(tableName, Optional.of(new HiveBucketProperty(ImmutableList.of("id"), BUCKETING_V1, 16, ImmutableList.of()))); - MaterializedResult.Builder bucket16Builder = MaterializedResult.resultBuilder(SESSION, BIGINT, VARCHAR, VARCHAR); - IntStream.range(0, rowCount).forEach(i -> bucket16Builder.row((long) i, String.valueOf(i), "sixteen")); - insertData(tableName, bucket16Builder.build()); - // write an 8-bucket partition - alterBucketProperty(tableName, Optional.of(new HiveBucketProperty(ImmutableList.of("id"), BUCKETING_V1, 8, ImmutableList.of()))); - MaterializedResult.Builder bucket8Builder = MaterializedResult.resultBuilder(SESSION, BIGINT, VARCHAR, VARCHAR); - IntStream.range(0, rowCount).forEach(i -> bucket8Builder.row((long) i, String.valueOf(i), "eight")); - insertData(tableName, bucket8Builder.build()); - - try (Transaction transaction = newTransaction()) { - ConnectorMetadata metadata = transaction.getMetadata(); - ConnectorSession session = newSession(); - metadata.beginQuery(session); - - ConnectorTableHandle tableHandle = getTableHandle(metadata, tableName); - - // read entire table - List columnHandles = ImmutableList.copyOf(metadata.getColumnHandles(session, tableHandle).values()); - MaterializedResult result = readTable( - transaction, - tableHandle, - columnHandles, - session, - TupleDomain.all(), - OptionalInt.empty(), - Optional.empty()); - assertBucketTableEvolutionResult(result, columnHandles, ImmutableSet.of(0, 1, 2, 3, 4, 5, 6, 7), rowCount); - - // read single bucket (table/logical bucket) - result = readTable( - transaction, - tableHandle, - columnHandles, - session, - TupleDomain.fromFixedValues(ImmutableMap.of(bucketColumnHandle(), NullableValue.of(INTEGER, 6L))), - OptionalInt.empty(), - Optional.empty()); - assertBucketTableEvolutionResult(result, columnHandles, ImmutableSet.of(6), rowCount); - - // read single bucket, without selecting the bucketing column (i.e. id column) - columnHandles = metadata.getColumnHandles(session, tableHandle).values().stream() - .filter(columnHandle -> !"id".equals(((HiveColumnHandle) columnHandle).getName())) - .collect(toImmutableList()); - result = readTable( - transaction, - tableHandle, - columnHandles, - session, - TupleDomain.fromFixedValues(ImmutableMap.of(bucketColumnHandle(), NullableValue.of(INTEGER, 6L))), - OptionalInt.empty(), - Optional.empty()); - assertBucketTableEvolutionResult(result, columnHandles, ImmutableSet.of(6), rowCount); - } - } - - private static void assertBucketTableEvolutionResult(MaterializedResult result, List columnHandles, Set bucketIds, int rowCount) - { - // Assert that only elements in the specified bucket shows up, and each element shows up 3 times. - int bucketCount = 8; - Set expectedIds = LongStream.range(0, rowCount) - .filter(x -> bucketIds.contains(toIntExact(x % bucketCount))) - .boxed() - .collect(toImmutableSet()); - - // assert that content from all three buckets are the same - Map columnIndex = indexColumns(columnHandles); - OptionalInt idColumnIndex = columnIndex.containsKey("id") ? OptionalInt.of(columnIndex.get("id")) : OptionalInt.empty(); - int nameColumnIndex = columnIndex.get("name"); - int bucketColumnIndex = columnIndex.get(BUCKET_COLUMN_NAME); - Map idCount = new HashMap<>(); - for (MaterializedRow row : result.getMaterializedRows()) { - String name = (String) row.getField(nameColumnIndex); - int bucket = (int) row.getField(bucketColumnIndex); - idCount.compute(Long.parseLong(name), (key, oldValue) -> oldValue == null ? 1 : oldValue + 1); - assertEquals(bucket, Integer.parseInt(name) % bucketCount); - if (idColumnIndex.isPresent()) { - long id = (long) row.getField(idColumnIndex.getAsInt()); - assertEquals(Integer.parseInt(name), id); - } - } - assertEquals( - (int) idCount.values().stream() - .distinct() - .collect(onlyElement()), - 3); - assertEquals(idCount.keySet(), expectedIds); - } - - @Test - public void testBucketedSortedTableEvolution() - throws Exception - { - SchemaTableName temporaryTable = temporaryTable("test_bucket_sorting_evolution"); - try { - doTestBucketedSortedTableEvolution(temporaryTable); - } - finally { - dropTable(temporaryTable); - } - } - - private void doTestBucketedSortedTableEvolution(SchemaTableName tableName) - throws Exception - { - int rowCount = 100; - // Create table and populate it with 3 partitions with different sort orders but same bucketing - createEmptyTable( - tableName, - ORC, - ImmutableList.of( - new Column("id", HIVE_LONG, Optional.empty()), - new Column("name", HIVE_STRING, Optional.empty())), - ImmutableList.of(new Column("pk", HIVE_STRING, Optional.empty())), - Optional.of(new HiveBucketProperty( - ImmutableList.of("id"), - BUCKETING_V1, - 4, - ImmutableList.of(new SortingColumn("id", ASCENDING), new SortingColumn("name", ASCENDING))))); - // write a 4-bucket partition sorted by id, name - MaterializedResult.Builder sortedByIdNameBuilder = MaterializedResult.resultBuilder(SESSION, BIGINT, VARCHAR, VARCHAR); - IntStream.range(0, rowCount).forEach(i -> sortedByIdNameBuilder.row((long) i, String.valueOf(i), "sorted_by_id_name")); - insertData(tableName, sortedByIdNameBuilder.build()); - - // write a 4-bucket partition sorted by name - alterBucketProperty(tableName, Optional.of(new HiveBucketProperty( - ImmutableList.of("id"), - BUCKETING_V1, - 4, - ImmutableList.of(new SortingColumn("name", ASCENDING))))); - MaterializedResult.Builder sortedByNameBuilder = MaterializedResult.resultBuilder(SESSION, BIGINT, VARCHAR, VARCHAR); - IntStream.range(0, rowCount).forEach(i -> sortedByNameBuilder.row((long) i, String.valueOf(i), "sorted_by_name")); - insertData(tableName, sortedByNameBuilder.build()); - - // write a 4-bucket partition sorted by id - alterBucketProperty(tableName, Optional.of(new HiveBucketProperty( - ImmutableList.of("id"), - BUCKETING_V1, - 4, - ImmutableList.of(new SortingColumn("id", ASCENDING))))); - MaterializedResult.Builder sortedByIdBuilder = MaterializedResult.resultBuilder(SESSION, BIGINT, VARCHAR, VARCHAR); - IntStream.range(0, rowCount).forEach(i -> sortedByIdBuilder.row((long) i, String.valueOf(i), "sorted_by_id")); - insertData(tableName, sortedByIdBuilder.build()); - - ConnectorTableHandle tableHandle; - try (Transaction transaction = newTransaction()) { - ConnectorMetadata metadata = transaction.getMetadata(); - ConnectorSession session = newSession(); - metadata.beginQuery(session); - tableHandle = getTableHandle(metadata, tableName); - - // read entire table - List columnHandles = ImmutableList.copyOf(metadata.getColumnHandles(session, tableHandle).values()); - MaterializedResult result = readTable(transaction, tableHandle, columnHandles, session, TupleDomain.all(), OptionalInt.empty(), Optional.empty()); - assertEquals(result.getRowCount(), 300); - } - - try (Transaction transaction = newTransaction()) { - ConnectorMetadata metadata = transaction.getMetadata(); - ConnectorSession session = newSession(ImmutableMap.of("propagate_table_scan_sorting_properties", true)); - metadata.beginQuery(session); - Map columnHandles = metadata.getColumnHandles(session, tableHandle); - // verify local sorting property - ConnectorTableProperties properties = metadata.getTableProperties(session, tableHandle); - assertEquals(properties.getLocalProperties(), ImmutableList.of( - new SortingProperty<>(columnHandles.get("id"), ASC_NULLS_FIRST))); - - // read on a entire table should fail with exception - assertThatThrownBy(() -> readTable(transaction, tableHandle, ImmutableList.copyOf(columnHandles.values()), session, TupleDomain.all(), OptionalInt.empty(), Optional.empty())) - .isInstanceOf(TrinoException.class) - .hasMessage("Hive table (%s) sorting by [id] is not compatible with partition (pk=sorted_by_name) sorting by [name]." + - " This restriction can be avoided by disabling propagate_table_scan_sorting_properties.", tableName); - - // read only the partitions with sorting that is compatible to table sorting - MaterializedResult result = readTable( - transaction, - tableHandle, - ImmutableList.copyOf(columnHandles.values()), - session, - TupleDomain.withColumnDomains(ImmutableMap.of( - columnHandles.get("pk"), - Domain.create(ValueSet.of(VARCHAR, utf8Slice("sorted_by_id_name"), utf8Slice("sorted_by_id")), false))), - OptionalInt.empty(), - Optional.empty()); - assertEquals(result.getRowCount(), 200); - } - } - - @Test - public void testBucketedTableValidation() - throws Exception - { - for (HiveStorageFormat storageFormat : createTableFormats) { - SchemaTableName table = temporaryTable("bucket_validation"); - try { - doTestBucketedTableValidation(storageFormat, table); - } - finally { - dropTable(table); - } - } - } - - private void doTestBucketedTableValidation(HiveStorageFormat storageFormat, SchemaTableName tableName) - throws Exception - { - prepareInvalidBuckets(storageFormat, tableName); - - // read succeeds when validation is disabled - try (Transaction transaction = newTransaction()) { - ConnectorMetadata metadata = transaction.getMetadata(); - ConnectorSession session = newSession(ImmutableMap.of("validate_bucketing", false)); - metadata.beginQuery(session); - ConnectorTableHandle tableHandle = getTableHandle(metadata, tableName); - List columnHandles = filterNonHiddenColumnHandles(metadata.getColumnHandles(session, tableHandle).values()); - MaterializedResult result = readTable(transaction, tableHandle, columnHandles, session, TupleDomain.all(), OptionalInt.empty(), Optional.of(storageFormat)); - assertEquals(result.getRowCount(), 87); // fewer rows due to deleted file - } - - // read fails due to validation failure - assertReadFailsWithMessageMatching(storageFormat, tableName, "Hive table is corrupt\\. File '.*/000002_0_.*' is for bucket 2, but contains a row for bucket 5."); - } - - private void prepareInvalidBuckets(HiveStorageFormat storageFormat, SchemaTableName tableName) - throws Exception - { - createEmptyTable( - tableName, - storageFormat, - ImmutableList.of( - new Column("id", HIVE_LONG, Optional.empty()), - new Column("name", HIVE_STRING, Optional.empty())), - ImmutableList.of(), - Optional.of(new HiveBucketProperty(ImmutableList.of("id"), BUCKETING_V1, 8, ImmutableList.of()))); - - MaterializedResult.Builder dataBuilder = MaterializedResult.resultBuilder(SESSION, BIGINT, VARCHAR); - for (long id = 0; id < 100; id++) { - dataBuilder.row(id, String.valueOf(id)); - } - insertData(tableName, dataBuilder.build()); - - try (Transaction transaction = newTransaction()) { - Set files = listAllDataFiles(transaction, tableName.getSchemaName(), tableName.getTableName()); - - Path bucket2 = files.stream() - .map(Path::new) - .filter(path -> path.getName().startsWith("000002_0_")) - .collect(onlyElement()); - - Path bucket5 = files.stream() - .map(Path::new) - .filter(path -> path.getName().startsWith("000005_0_")) - .collect(onlyElement()); - - HdfsContext context = new HdfsContext(newSession()); - FileSystem fileSystem = hdfsEnvironment.getFileSystem(context, bucket2); - fileSystem.delete(bucket2, false); - fileSystem.rename(bucket5, bucket2); - } - } - - protected void assertReadFailsWithMessageMatching(HiveStorageFormat storageFormat, SchemaTableName tableName, String regex) - { - try (Transaction transaction = newTransaction()) { - ConnectorMetadata metadata = transaction.getMetadata(); - ConnectorSession session = newSession(); - metadata.beginQuery(session); - ConnectorTableHandle tableHandle = getTableHandle(metadata, tableName); - List columnHandles = filterNonHiddenColumnHandles(metadata.getColumnHandles(session, tableHandle).values()); - assertTrinoExceptionThrownBy( - () -> readTable(transaction, tableHandle, columnHandles, session, TupleDomain.all(), OptionalInt.empty(), Optional.of(storageFormat))) - .hasErrorCode(HIVE_INVALID_BUCKET_FILES) - .hasMessageMatching(regex); - } - } - - private void assertTableIsBucketed(ConnectorTableHandle tableHandle, Transaction transaction, ConnectorSession session) - { - // the bucketed test tables should have ~32 splits - List splits = getAllSplits(tableHandle, transaction, session); - assertThat(splits.size()).as("splits.size()") - .isBetween(31, 32); - - // verify all paths are unique - Set paths = new HashSet<>(); - for (ConnectorSplit split : splits) { - assertTrue(paths.add(((HiveSplit) split).getPath())); - } - } - - @Test - public void testGetRecords() - throws Exception - { - try (Transaction transaction = newTransaction()) { - ConnectorMetadata metadata = transaction.getMetadata(); - ConnectorSession session = newSession(); - metadata.beginQuery(session); - - ConnectorTableHandle tableHandle = getTableHandle(metadata, tablePartitionFormat); - ConnectorTableMetadata tableMetadata = metadata.getTableMetadata(session, tableHandle); - List columnHandles = ImmutableList.copyOf(metadata.getColumnHandles(session, tableHandle).values()); - Map columnIndex = indexColumns(columnHandles); - - List splits = getAllSplits(tableHandle, transaction, session); - assertEquals(splits.size(), tablePartitionFormatPartitions.size()); - - for (ConnectorSplit split : splits) { - HiveSplit hiveSplit = (HiveSplit) split; - - List partitionKeys = hiveSplit.getPartitionKeys(); - String ds = partitionKeys.get(0).getValue(); - String fileFormat = partitionKeys.get(1).getValue(); - HiveStorageFormat fileType = HiveStorageFormat.valueOf(fileFormat.toUpperCase(ENGLISH)); - int dummyPartition = Integer.parseInt(partitionKeys.get(2).getValue()); - - long rowNumber = 0; - long completedBytes = 0; - try (ConnectorPageSource pageSource = pageSourceProvider.createPageSource(transaction.getTransactionHandle(), session, hiveSplit, tableHandle, columnHandles, DynamicFilter.EMPTY)) { - MaterializedResult result = materializeSourceDataStream(session, pageSource, getTypes(columnHandles)); - - assertPageSourceType(pageSource, fileType); - - for (MaterializedRow row : result) { - try { - assertValueTypes(row, tableMetadata.getColumns()); - } - catch (RuntimeException e) { - throw new RuntimeException("row " + rowNumber, e); - } - - rowNumber++; - Object value; - - value = row.getField(columnIndex.get("t_string")); - if (rowNumber % 19 == 0) { - assertNull(value); - } - else if (rowNumber % 19 == 1) { - assertEquals(value, ""); - } - else { - assertEquals(value, "test"); - } - - assertEquals(row.getField(columnIndex.get("t_tinyint")), (byte) (1 + rowNumber)); - assertEquals(row.getField(columnIndex.get("t_smallint")), (short) (2 + rowNumber)); - assertEquals(row.getField(columnIndex.get("t_int")), 3 + (int) rowNumber); - - if (rowNumber % 13 == 0) { - assertNull(row.getField(columnIndex.get("t_bigint"))); - } - else { - assertEquals(row.getField(columnIndex.get("t_bigint")), 4 + rowNumber); - } - - assertEquals((Float) row.getField(columnIndex.get("t_float")), 5.1f + rowNumber, 0.001); - assertEquals(row.getField(columnIndex.get("t_double")), 6.2 + rowNumber); - - if (rowNumber % 3 == 2) { - assertNull(row.getField(columnIndex.get("t_boolean"))); - } - else { - assertEquals(row.getField(columnIndex.get("t_boolean")), rowNumber % 3 != 0); - } - - assertEquals(row.getField(columnIndex.get("ds")), ds); - assertEquals(row.getField(columnIndex.get("file_format")), fileFormat); - assertEquals(row.getField(columnIndex.get("dummy")), dummyPartition); - - long newCompletedBytes = pageSource.getCompletedBytes(); - assertTrue(newCompletedBytes >= completedBytes); - assertTrue(newCompletedBytes <= hiveSplit.getLength()); - completedBytes = newCompletedBytes; - } - - assertTrue(completedBytes <= hiveSplit.getLength()); - assertEquals(rowNumber, 100); - } - } - } - } - - @Test - public void testGetPartialRecords() - throws Exception - { - try (Transaction transaction = newTransaction()) { - ConnectorMetadata metadata = transaction.getMetadata(); - ConnectorSession session = newSession(); - metadata.beginQuery(session); - - ConnectorTableHandle tableHandle = getTableHandle(metadata, tablePartitionFormat); - List columnHandles = ImmutableList.copyOf(metadata.getColumnHandles(session, tableHandle).values()); - Map columnIndex = indexColumns(columnHandles); - - List splits = getAllSplits(tableHandle, transaction, session); - assertEquals(splits.size(), tablePartitionFormatPartitions.size()); - - for (ConnectorSplit split : splits) { - HiveSplit hiveSplit = (HiveSplit) split; - - List partitionKeys = hiveSplit.getPartitionKeys(); - String ds = partitionKeys.get(0).getValue(); - String fileFormat = partitionKeys.get(1).getValue(); - HiveStorageFormat fileType = HiveStorageFormat.valueOf(fileFormat.toUpperCase(ENGLISH)); - int dummyPartition = Integer.parseInt(partitionKeys.get(2).getValue()); - - long rowNumber = 0; - try (ConnectorPageSource pageSource = pageSourceProvider.createPageSource(transaction.getTransactionHandle(), session, hiveSplit, tableHandle, columnHandles, DynamicFilter.EMPTY)) { - assertPageSourceType(pageSource, fileType); - MaterializedResult result = materializeSourceDataStream(session, pageSource, getTypes(columnHandles)); - for (MaterializedRow row : result) { - rowNumber++; - - assertEquals(row.getField(columnIndex.get("t_double")), 6.2 + rowNumber); - assertEquals(row.getField(columnIndex.get("ds")), ds); - assertEquals(row.getField(columnIndex.get("file_format")), fileFormat); - assertEquals(row.getField(columnIndex.get("dummy")), dummyPartition); - } - } - assertEquals(rowNumber, 100); - } - } - } - - @Test - public void testGetRecordsUnpartitioned() - throws Exception - { - try (Transaction transaction = newTransaction()) { - ConnectorMetadata metadata = transaction.getMetadata(); - ConnectorSession session = newSession(); - metadata.beginQuery(session); - - ConnectorTableHandle tableHandle = getTableHandle(metadata, tableUnpartitioned); - List columnHandles = ImmutableList.copyOf(metadata.getColumnHandles(session, tableHandle).values()); - Map columnIndex = indexColumns(columnHandles); - - List splits = getAllSplits(tableHandle, transaction, session); - assertThat(splits).hasSameSizeAs(tableUnpartitionedPartitions); - - for (ConnectorSplit split : splits) { - HiveSplit hiveSplit = (HiveSplit) split; - - assertEquals(hiveSplit.getPartitionKeys(), ImmutableList.of()); - - long rowNumber = 0; - try (ConnectorPageSource pageSource = pageSourceProvider.createPageSource(transaction.getTransactionHandle(), session, split, tableHandle, columnHandles, DynamicFilter.EMPTY)) { - assertPageSourceType(pageSource, TEXTFILE); - MaterializedResult result = materializeSourceDataStream(session, pageSource, getTypes(columnHandles)); - - for (MaterializedRow row : result) { - rowNumber++; - - if (rowNumber % 19 == 0) { - assertNull(row.getField(columnIndex.get("t_string"))); - } - else if (rowNumber % 19 == 1) { - assertEquals(row.getField(columnIndex.get("t_string")), ""); - } - else { - assertEquals(row.getField(columnIndex.get("t_string")), "unpartitioned"); - } - - assertEquals(row.getField(columnIndex.get("t_tinyint")), (byte) (1 + rowNumber)); - } - } - assertEquals(rowNumber, 100); - } - } - } - - @Test(expectedExceptions = TrinoException.class, expectedExceptionsMessageRegExp = ".*The column 't_data' in table '.*\\.trino_test_partition_schema_change' is declared as type 'double', but partition 'ds=2012-12-29' declared column 't_data' as type 'string'.") - public void testPartitionSchemaMismatch() - throws Exception - { - try (Transaction transaction = newTransaction()) { - ConnectorMetadata metadata = transaction.getMetadata(); - ConnectorTableHandle table = getTableHandle(metadata, tablePartitionSchemaChange); - ConnectorSession session = newSession(); - metadata.beginQuery(session); - readTable(transaction, table, ImmutableList.of(dsColumn), session, TupleDomain.all(), OptionalInt.empty(), Optional.empty()); - } - } - - // TODO coercion of non-canonical values should be supported - @Test(enabled = false) - public void testPartitionSchemaNonCanonical() - throws Exception - { - try (Transaction transaction = newTransaction()) { - ConnectorSession session = newSession(); - ConnectorMetadata metadata = transaction.getMetadata(); - - ConnectorTableHandle table = getTableHandle(metadata, tablePartitionSchemaChangeNonCanonical); - ColumnHandle column = metadata.getColumnHandles(session, table).get("t_boolean"); - - Constraint constraint = new Constraint(TupleDomain.fromFixedValues(ImmutableMap.of(column, NullableValue.of(BOOLEAN, false)))); - table = applyFilter(metadata, table, constraint); - HivePartition partition = getOnlyElement(((HiveTableHandle) table).getPartitions().orElseThrow(AssertionError::new)); - assertEquals(getPartitionId(partition), "t_boolean=0"); - - ConnectorSplitSource splitSource = getSplits(splitManager, transaction, session, table); - ConnectorSplit split = getOnlyElement(getAllSplits(splitSource)); - - ImmutableList columnHandles = ImmutableList.of(column); - try (ConnectorPageSource ignored = pageSourceProvider.createPageSource(transaction.getTransactionHandle(), session, split, table, columnHandles, DynamicFilter.EMPTY)) { - fail("expected exception"); - } - catch (TrinoException e) { - assertEquals(e.getErrorCode(), HIVE_INVALID_PARTITION_VALUE.toErrorCode()); - } - } - } - - @Test - public void testTypesTextFile() - throws Exception - { - assertGetRecords("trino_test_types_textfile", TEXTFILE); - } - - @Test - public void testTypesSequenceFile() - throws Exception - { - assertGetRecords("trino_test_types_sequencefile", SEQUENCEFILE); - } - - @Test - public void testTypesRcText() - throws Exception - { - assertGetRecords("trino_test_types_rctext", RCTEXT); - } - - @Test - public void testTypesRcBinary() - throws Exception - { - assertGetRecords("trino_test_types_rcbinary", RCBINARY); - } - - @Test - public void testTypesOrc() - throws Exception - { - assertGetRecords("trino_test_types_orc", ORC); - } - - @Test - public void testTypesParquet() - throws Exception - { - assertGetRecords("trino_test_types_parquet", PARQUET); - } - - @Test - public void testEmptyTextFile() - throws Exception - { - assertEmptyFile(TEXTFILE); - } - - @Test - public void testEmptySequenceFile() - throws Exception - { - assertEmptyFile(SEQUENCEFILE); - } - - @Test - public void testEmptyRcTextFile() - throws Exception - { - assertEmptyFile(RCTEXT); - } - - @Test - public void testEmptyRcBinaryFile() - throws Exception - { - assertEmptyFile(RCBINARY); - } - - @Test - public void testEmptyOrcFile() - throws Exception - { - assertEmptyFile(ORC); - } - - private void assertEmptyFile(HiveStorageFormat format) - throws Exception - { - SchemaTableName tableName = temporaryTable("empty_file"); - try { - List columns = ImmutableList.of(new Column("test", HIVE_STRING, Optional.empty())); - createEmptyTable(tableName, format, columns, ImmutableList.of()); - - try (Transaction transaction = newTransaction()) { - ConnectorSession session = newSession(); - ConnectorMetadata metadata = transaction.getMetadata(); - metadata.beginQuery(session); - - ConnectorTableHandle tableHandle = getTableHandle(metadata, tableName); - List columnHandles = filterNonHiddenColumnHandles(metadata.getColumnHandles(session, tableHandle).values()); - - Table table = transaction.getMetastore() - .getTable(tableName.getSchemaName(), tableName.getTableName()) - .orElseThrow(AssertionError::new); - - // verify directory is empty - HdfsContext context = new HdfsContext(session); - Path location = new Path(table.getStorage().getLocation()); - assertTrue(listDirectory(context, location).isEmpty()); - - // read table with empty directory - readTable(transaction, tableHandle, columnHandles, session, TupleDomain.all(), OptionalInt.of(0), Optional.of(ORC)); - - // create empty file - FileSystem fileSystem = hdfsEnvironment.getFileSystem(context, location); - assertTrue(fileSystem.createNewFile(new Path(location, "empty-file"))); - assertEquals(listDirectory(context, location), ImmutableList.of("empty-file")); - - // read table with empty file - MaterializedResult result = readTable(transaction, tableHandle, columnHandles, session, TupleDomain.all(), OptionalInt.of(0), Optional.empty()); - assertEquals(result.getRowCount(), 0); - } - } - finally { - dropTable(tableName); - } - } - - @Test - public void testRenameTable() - { - SchemaTableName temporaryRenameTableOld = temporaryTable("rename_old"); - SchemaTableName temporaryRenameTableNew = temporaryTable("rename_new"); - try { - createDummyTable(temporaryRenameTableOld); - - try (Transaction transaction = newTransaction()) { - ConnectorSession session = newSession(); - ConnectorMetadata metadata = transaction.getMetadata(); - - metadata.renameTable(session, getTableHandle(metadata, temporaryRenameTableOld), temporaryRenameTableNew); - transaction.commit(); - } - - try (Transaction transaction = newTransaction()) { - ConnectorSession session = newSession(); - ConnectorMetadata metadata = transaction.getMetadata(); - - assertNull(metadata.getTableHandle(session, temporaryRenameTableOld)); - assertNotNull(metadata.getTableHandle(session, temporaryRenameTableNew)); - } - } - finally { - dropTable(temporaryRenameTableOld); - dropTable(temporaryRenameTableNew); - } - } - - @Test - public void testTableCreation() - throws Exception - { - for (HiveStorageFormat storageFormat : createTableFormats) { - SchemaTableName temporaryCreateTable = temporaryTable("create"); - try { - doCreateTable(temporaryCreateTable, storageFormat); - } - finally { - dropTable(temporaryCreateTable); - } - } - } - - @Test - public void testTableCreationWithTrailingSpaceInLocation() - throws Exception - { - SchemaTableName tableName = temporaryTable("test_table_creation_with_trailing_space_in_location_" + randomNameSuffix()); - String tableDefaultLocationWithTrailingSpace = null; - try { - try (Transaction transaction = newTransaction()) { - ConnectorSession session = newSession(); - SemiTransactionalHiveMetastore metastore = transaction.getMetastore(); - - // Write data - tableDefaultLocationWithTrailingSpace = getTableDefaultLocation(new HdfsContext(session), metastore, HDFS_ENVIRONMENT, tableName.getSchemaName(), tableName.getTableName()) + " "; - Path dataFilePath = new Path(tableDefaultLocationWithTrailingSpace, "foo.txt"); - FileSystem fs = hdfsEnvironment.getFileSystem(new HdfsContext(session), new Path(tableDefaultLocationWithTrailingSpace)); - try (OutputStream outputStream = fs.create(dataFilePath)) { - outputStream.write("hello\u0001world\nbye\u0001world".getBytes(UTF_8)); - } - - // create table - ConnectorTableMetadata tableMetadata = new ConnectorTableMetadata( - tableName, - ImmutableList.builder() - .add(new ColumnMetadata("t_string1", VARCHAR)) - .add(new ColumnMetadata("t_string2", VARCHAR)) - .build(), - ImmutableMap.builder() - .putAll(createTableProperties(TEXTFILE, ImmutableList.of())) - .put(EXTERNAL_LOCATION_PROPERTY, tableDefaultLocationWithTrailingSpace) - .buildOrThrow()); - - ConnectorMetadata metadata = transaction.getMetadata(); - metadata.createTable(session, tableMetadata, false); - - transaction.commit(); - } - - try (Transaction transaction = newTransaction()) { - ConnectorSession session = newSession(); - ConnectorMetadata metadata = transaction.getMetadata(); - metadata.beginQuery(session); - - // verify the data - ConnectorTableHandle tableHandle = getTableHandle(metadata, tableName); - List columnHandles = filterNonHiddenColumnHandles(metadata.getColumnHandles(session, tableHandle).values()); - MaterializedResult result = readTable(transaction, tableHandle, columnHandles, session, TupleDomain.all(), OptionalInt.empty(), Optional.of(TEXTFILE)); - assertEqualsIgnoreOrder( - result.getMaterializedRows(), - MaterializedResult.resultBuilder(SESSION, VARCHAR, VARCHAR) - .row("hello", "world") - .row("bye", "world") - .build()); - } - } - finally { - dropTable(tableName); - if (tableDefaultLocationWithTrailingSpace != null) { - FileSystem fs = hdfsEnvironment.getFileSystem(new HdfsContext(SESSION), new Path(tableDefaultLocationWithTrailingSpace)); - fs.delete(new Path(tableDefaultLocationWithTrailingSpace), true); - } - } - } - - @Test - public void testTableCreationRollback() - throws Exception - { - SchemaTableName temporaryCreateRollbackTable = temporaryTable("create_rollback"); - try { - Location stagingPathRoot; - try (Transaction transaction = newTransaction()) { - ConnectorSession session = newSession(); - ConnectorMetadata metadata = transaction.getMetadata(); - - // begin creating the table - ConnectorTableMetadata tableMetadata = new ConnectorTableMetadata(temporaryCreateRollbackTable, CREATE_TABLE_COLUMNS, createTableProperties(RCBINARY)); - - ConnectorOutputTableHandle outputHandle = metadata.beginCreateTable(session, tableMetadata, Optional.empty(), NO_RETRIES); - - // write the data - ConnectorPageSink sink = pageSinkProvider.createPageSink(transaction.getTransactionHandle(), session, outputHandle, TESTING_PAGE_SINK_ID); - sink.appendPage(CREATE_TABLE_DATA.toPage()); - getFutureValue(sink.finish()); - - // verify we have data files - stagingPathRoot = getStagingPathRoot(outputHandle); - HdfsContext context = new HdfsContext(session); - assertFalse(listAllDataFiles(context, stagingPathRoot).isEmpty()); - - // rollback the table - transaction.rollback(); - } - - // verify all files have been deleted - HdfsContext context = new HdfsContext(newSession()); - assertTrue(listAllDataFiles(context, stagingPathRoot).isEmpty()); - - // verify table is not in the metastore - try (Transaction transaction = newTransaction()) { - ConnectorSession session = newSession(); - ConnectorMetadata metadata = transaction.getMetadata(); - assertNull(metadata.getTableHandle(session, temporaryCreateRollbackTable)); - } - } - finally { - dropTable(temporaryCreateRollbackTable); - } - } - - @Test - public void testTableCreationIgnoreExisting() - { - List columns = ImmutableList.of(new Column("dummy", HiveType.valueOf("uniontype"), Optional.empty())); - SchemaTableName schemaTableName = temporaryTable("create"); - ConnectorSession session = newSession(); - String schemaName = schemaTableName.getSchemaName(); - String tableName = schemaTableName.getTableName(); - PrincipalPrivileges privileges = testingPrincipalPrivilege(session); - Location targetPath; - try { - try (Transaction transaction = newTransaction()) { - LocationService locationService = getLocationService(); - targetPath = locationService.forNewTable(transaction.getMetastore(), session, schemaName, tableName); - Table table = createSimpleTable(schemaTableName, columns, session, targetPath, "q1"); - transaction.getMetastore() - .createTable(session, table, privileges, Optional.empty(), Optional.empty(), false, ZERO_TABLE_STATISTICS, false); - Optional
tableHandle = transaction.getMetastore().getTable(schemaName, tableName); - assertTrue(tableHandle.isPresent()); - transaction.commit(); - } - - // try creating it again from another transaction with ignoreExisting=false - try (Transaction transaction = newTransaction()) { - Table table = createSimpleTable(schemaTableName, columns, session, targetPath.appendSuffix("_2"), "q2"); - transaction.getMetastore() - .createTable(session, table, privileges, Optional.empty(), Optional.empty(), false, ZERO_TABLE_STATISTICS, false); - transaction.commit(); - fail("Expected exception"); - } - catch (TrinoException e) { - assertInstanceOf(e, TableAlreadyExistsException.class); - } - - // try creating it again from another transaction with ignoreExisting=true - try (Transaction transaction = newTransaction()) { - Table table = createSimpleTable(schemaTableName, columns, session, targetPath.appendSuffix("_3"), "q3"); - transaction.getMetastore() - .createTable(session, table, privileges, Optional.empty(), Optional.empty(), true, ZERO_TABLE_STATISTICS, false); - transaction.commit(); - } - - // at this point the table should exist, now try creating the table again with a different table definition - columns = ImmutableList.of(new Column("new_column", HiveType.valueOf("string"), Optional.empty())); - try (Transaction transaction = newTransaction()) { - Table table = createSimpleTable(schemaTableName, columns, session, targetPath.appendSuffix("_4"), "q4"); - transaction.getMetastore() - .createTable(session, table, privileges, Optional.empty(), Optional.empty(), true, ZERO_TABLE_STATISTICS, false); - transaction.commit(); - fail("Expected exception"); - } - catch (TrinoException e) { - assertEquals(e.getErrorCode(), TRANSACTION_CONFLICT.toErrorCode()); - assertEquals(e.getMessage(), format("Table already exists with a different schema: '%s'", schemaTableName.getTableName())); - } - } - finally { - dropTable(schemaTableName); - } - } - - private static Table createSimpleTable(SchemaTableName schemaTableName, List columns, ConnectorSession session, Location targetPath, String queryId) - { - String tableOwner = session.getUser(); - String schemaName = schemaTableName.getSchemaName(); - String tableName = schemaTableName.getTableName(); - return Table.builder() - .setDatabaseName(schemaName) - .setTableName(tableName) - .setOwner(Optional.of(tableOwner)) - .setTableType(TableType.MANAGED_TABLE.name()) - .setParameters(ImmutableMap.of( - PRESTO_VERSION_NAME, TEST_SERVER_VERSION, - PRESTO_QUERY_ID_NAME, queryId)) - .setDataColumns(columns) - .withStorage(storage -> storage - .setLocation(targetPath.toString()) - .setStorageFormat(fromHiveStorageFormat(ORC)) - .setSerdeParameters(ImmutableMap.of())) - .build(); - } - - @Test - public void testBucketSortedTables() - throws Exception - { - SchemaTableName table = temporaryTable("create_sorted"); - try { - doTestBucketSortedTables(table); - } - finally { - dropTable(table); - } - } - - private void doTestBucketSortedTables(SchemaTableName table) - throws IOException - { - int bucketCount = 3; - int expectedRowCount = 0; - - try (Transaction transaction = newTransaction()) { - ConnectorSession session = newSession(); - ConnectorMetadata metadata = transaction.getMetadata(); - - // begin creating the table - ConnectorTableMetadata tableMetadata = new ConnectorTableMetadata( - table, - ImmutableList.builder() - .add(new ColumnMetadata("id", VARCHAR)) - .add(new ColumnMetadata("value_asc", VARCHAR)) - .add(new ColumnMetadata("value_desc", BIGINT)) - .add(new ColumnMetadata("ds", VARCHAR)) - .build(), - ImmutableMap.builder() - .put(STORAGE_FORMAT_PROPERTY, RCBINARY) - .put(PARTITIONED_BY_PROPERTY, ImmutableList.of("ds")) - .put(BUCKETED_BY_PROPERTY, ImmutableList.of("id")) - .put(BUCKET_COUNT_PROPERTY, bucketCount) - .put(SORTED_BY_PROPERTY, ImmutableList.builder() - .add(new SortingColumn("value_asc", ASCENDING)) - .add(new SortingColumn("value_desc", DESCENDING)) - .build()) - .buildOrThrow()); - - ConnectorOutputTableHandle outputHandle = metadata.beginCreateTable(session, tableMetadata, Optional.empty(), NO_RETRIES); - - // write the data - ConnectorPageSink sink = pageSinkProvider.createPageSink(transaction.getTransactionHandle(), session, outputHandle, TESTING_PAGE_SINK_ID); - List types = tableMetadata.getColumns().stream() - .map(ColumnMetadata::getType) - .collect(toList()); - ThreadLocalRandom random = ThreadLocalRandom.current(); - for (int i = 0; i < 50; i++) { - MaterializedResult.Builder builder = MaterializedResult.resultBuilder(session, types); - for (int j = 0; j < 1000; j++) { - builder.row( - sha256().hashLong(random.nextLong()).toString(), - "test" + random.nextInt(100), - random.nextLong(100_000), - "2018-04-01"); - expectedRowCount++; - } - sink.appendPage(builder.build().toPage()); - } - - HdfsContext context = new HdfsContext(session); - HiveConfig config = getHiveConfig(); - // verify we have enough temporary files per bucket to require multiple passes - Location stagingPathRoot; - if (config.isTemporaryStagingDirectoryEnabled()) { - stagingPathRoot = Location.of(config.getTemporaryStagingDirectoryPath() - .replace("${USER}", context.getIdentity().getUser())); - } - else { - stagingPathRoot = getStagingPathRoot(outputHandle); - } - assertThat(listAllDataFiles(context, stagingPathRoot)) - .filteredOn(file -> file.contains(".tmp-sort.")) - .size().isGreaterThan(bucketCount * getSortingFileWriterConfig().getMaxOpenSortFiles() * 2); - - // finish the write - Collection fragments = getFutureValue(sink.finish()); - - // verify there are no temporary files - for (String file : listAllDataFiles(context, stagingPathRoot)) { - assertThat(file).doesNotContain(".tmp-sort."); - } - - // finish creating table - metadata.finishCreateTable(session, outputHandle, fragments, ImmutableList.of()); - - transaction.commit(); - } - - // verify that bucket files are sorted - try (Transaction transaction = newTransaction()) { - ConnectorMetadata metadata = transaction.getMetadata(); - ConnectorSession session = newSession(); - metadata.beginQuery(session); - - ConnectorTableHandle tableHandle = getTableHandle(metadata, table); - List columnHandles = ImmutableList.copyOf(metadata.getColumnHandles(session, tableHandle).values()); - - // verify local sorting property - ConnectorTableProperties properties = metadata.getTableProperties( - newSession(ImmutableMap.of( - "propagate_table_scan_sorting_properties", true, - "bucket_execution_enabled", false)), - tableHandle); - Map columnIndex = indexColumns(columnHandles); - assertEquals(properties.getLocalProperties(), ImmutableList.of( - new SortingProperty<>(columnHandles.get(columnIndex.get("value_asc")), ASC_NULLS_FIRST), - new SortingProperty<>(columnHandles.get(columnIndex.get("value_desc")), DESC_NULLS_LAST))); - assertThat(metadata.getTableProperties(newSession(), tableHandle).getLocalProperties()).isEmpty(); - - List splits = getAllSplits(tableHandle, transaction, session); - assertThat(splits).hasSize(bucketCount); - - int actualRowCount = 0; - for (ConnectorSplit split : splits) { - try (ConnectorPageSource pageSource = pageSourceProvider.createPageSource(transaction.getTransactionHandle(), session, split, tableHandle, columnHandles, DynamicFilter.EMPTY)) { - String lastValueAsc = null; - long lastValueDesc = -1; - - while (!pageSource.isFinished()) { - Page page = pageSource.getNextPage(); - if (page == null) { - continue; - } - for (int i = 0; i < page.getPositionCount(); i++) { - Block blockAsc = page.getBlock(1); - Block blockDesc = page.getBlock(2); - assertFalse(blockAsc.isNull(i)); - assertFalse(blockDesc.isNull(i)); - - String valueAsc = VARCHAR.getSlice(blockAsc, i).toStringUtf8(); - if (lastValueAsc != null) { - assertGreaterThanOrEqual(valueAsc, lastValueAsc); - if (valueAsc.equals(lastValueAsc)) { - long valueDesc = BIGINT.getLong(blockDesc, i); - if (lastValueDesc != -1) { - assertLessThanOrEqual(valueDesc, lastValueDesc); - } - lastValueDesc = valueDesc; - } - else { - lastValueDesc = -1; - } - } - lastValueAsc = valueAsc; - actualRowCount++; - } - } - } - } - assertThat(actualRowCount).isEqualTo(expectedRowCount); - } - } - - @Test - public void testInsert() - throws Exception - { - for (HiveStorageFormat storageFormat : createTableFormats) { - SchemaTableName temporaryInsertTable = temporaryTable("insert"); - try { - doInsert(storageFormat, temporaryInsertTable); - } - finally { - dropTable(temporaryInsertTable); - } - } - } - - @Test - public void testInsertOverwriteUnpartitioned() - throws Exception - { - SchemaTableName table = temporaryTable("insert_overwrite"); - try { - doInsertOverwriteUnpartitioned(table); - } - finally { - dropTable(table); - } - } - - @Test - public void testInsertIntoNewPartition() - throws Exception - { - for (HiveStorageFormat storageFormat : createTableFormats) { - SchemaTableName temporaryInsertIntoNewPartitionTable = temporaryTable("insert_new_partitioned"); - try { - doInsertIntoNewPartition(storageFormat, temporaryInsertIntoNewPartitionTable); - } - finally { - dropTable(temporaryInsertIntoNewPartitionTable); - } - } - } - - @Test - public void testInsertIntoExistingPartition() - throws Exception - { - for (HiveStorageFormat storageFormat : createTableFormats) { - SchemaTableName temporaryInsertIntoExistingPartitionTable = temporaryTable("insert_existing_partitioned"); - try { - doInsertIntoExistingPartition(storageFormat, temporaryInsertIntoExistingPartitionTable); - } - finally { - dropTable(temporaryInsertIntoExistingPartitionTable); - } - } - } - - @Test - public void testInsertIntoExistingPartitionEmptyStatistics() - throws Exception - { - for (HiveStorageFormat storageFormat : createTableFormats) { - SchemaTableName temporaryInsertIntoExistingPartitionTable = temporaryTable("insert_existing_partitioned_empty_statistics"); - try { - doInsertIntoExistingPartitionEmptyStatistics(storageFormat, temporaryInsertIntoExistingPartitionTable); - } - finally { - dropTable(temporaryInsertIntoExistingPartitionTable); - } - } - } - - @Test - public void testInsertUnsupportedWriteType() - throws Exception - { - SchemaTableName temporaryInsertUnsupportedWriteType = temporaryTable("insert_unsupported_type"); - try { - doInsertUnsupportedWriteType(ORC, temporaryInsertUnsupportedWriteType); - } - finally { - dropTable(temporaryInsertUnsupportedWriteType); - } - } - - @Test - public void testMetadataDelete() - throws Exception - { - for (HiveStorageFormat storageFormat : createTableFormats) { - SchemaTableName temporaryMetadataDeleteTable = temporaryTable("metadata_delete"); - try { - doTestMetadataDelete(storageFormat, temporaryMetadataDeleteTable); - } - finally { - dropTable(temporaryMetadataDeleteTable); - } - } - } - - @Test - public void testEmptyTableCreation() - throws Exception - { - for (HiveStorageFormat storageFormat : createTableFormats) { - SchemaTableName temporaryCreateEmptyTable = temporaryTable("create_empty"); - try { - doCreateEmptyTable(temporaryCreateEmptyTable, storageFormat, CREATE_TABLE_COLUMNS); - } - finally { - dropTable(temporaryCreateEmptyTable); - } - } - } - - @Test - public void testCreateEmptyTableShouldNotCreateStagingDirectory() - throws IOException - { - for (HiveStorageFormat storageFormat : createTableFormats) { - SchemaTableName temporaryCreateEmptyTable = temporaryTable("create_empty"); - try { - List columns = ImmutableList.of(new Column("test", HIVE_STRING, Optional.empty())); - try (Transaction transaction = newTransaction()) { - String temporaryStagingPrefix = "hive-temporary-staging-prefix-" + UUID.randomUUID().toString().toLowerCase(ENGLISH).replace("-", ""); - ConnectorSession session = newSession(); - String tableOwner = session.getUser(); - String schemaName = temporaryCreateEmptyTable.getSchemaName(); - String tableName = temporaryCreateEmptyTable.getTableName(); - HiveConfig hiveConfig = getHiveConfig() - .setTemporaryStagingDirectoryPath(temporaryStagingPrefix) - .setTemporaryStagingDirectoryEnabled(true); - LocationService locationService = new HiveLocationService(hdfsEnvironment, hiveConfig); - Location targetPath = locationService.forNewTable(transaction.getMetastore(), session, schemaName, tableName); - Table.Builder tableBuilder = Table.builder() - .setDatabaseName(schemaName) - .setTableName(tableName) - .setOwner(Optional.of(tableOwner)) - .setTableType(MANAGED_TABLE.name()) - .setParameters(ImmutableMap.of( - PRESTO_VERSION_NAME, TEST_SERVER_VERSION, - PRESTO_QUERY_ID_NAME, session.getQueryId())) - .setDataColumns(columns); - tableBuilder.getStorageBuilder() - .setLocation(targetPath.toString()) - .setStorageFormat(StorageFormat.create(storageFormat.getSerde(), storageFormat.getInputFormat(), storageFormat.getOutputFormat())); - transaction.getMetastore().createTable( - session, - tableBuilder.build(), - testingPrincipalPrivilege(tableOwner, session.getUser()), - Optional.empty(), - Optional.empty(), - true, - ZERO_TABLE_STATISTICS, - false); - transaction.commit(); - - HdfsContext context = new HdfsContext(session); - Path temporaryRoot = new Path(targetPath.toString(), temporaryStagingPrefix); - FileSystem fileSystem = hdfsEnvironment.getFileSystem(context, temporaryRoot); - assertFalse(fileSystem.exists(temporaryRoot), format("Temporary staging directory %s is created.", temporaryRoot)); - } - } - finally { - dropTable(temporaryCreateEmptyTable); - } - } - } - - @Test - public void testViewCreation() - { - SchemaTableName temporaryCreateView = temporaryTable("create_view"); - try { - verifyViewCreation(temporaryCreateView); - } - finally { - try (Transaction transaction = newTransaction()) { - ConnectorMetadata metadata = transaction.getMetadata(); - metadata.dropView(newSession(), temporaryCreateView); - transaction.commit(); - } - catch (RuntimeException e) { - // this usually occurs because the view was not created - } - } - } - - @Test - public void testCreateTableUnsupportedType() - { - for (HiveStorageFormat storageFormat : createTableFormats) { - try (Transaction transaction = newTransaction()) { - ConnectorSession session = newSession(); - ConnectorMetadata metadata = transaction.getMetadata(); - List columns = ImmutableList.of(new ColumnMetadata("dummy", HYPER_LOG_LOG)); - ConnectorTableMetadata tableMetadata = new ConnectorTableMetadata(invalidTable, columns, createTableProperties(storageFormat)); - metadata.beginCreateTable(session, tableMetadata, Optional.empty(), NO_RETRIES); - fail("create table with unsupported type should fail for storage format " + storageFormat); - } - catch (TrinoException e) { - assertEquals(e.getErrorCode(), NOT_SUPPORTED.toErrorCode()); - } - } - } - - @Test - public void testHideDeltaLakeTables() - { - ConnectorSession session = newSession(); - SchemaTableName tableName = temporaryTable("trino_delta_lake_table"); - - Table.Builder table = Table.builder() - .setDatabaseName(tableName.getSchemaName()) - .setTableName(tableName.getTableName()) - .setOwner(Optional.of(session.getUser())) - .setTableType(MANAGED_TABLE.name()) - .setPartitionColumns(List.of(new Column("a_partition_column", HIVE_INT, Optional.empty()))) - .setDataColumns(List.of(new Column("a_column", HIVE_STRING, Optional.empty()))) - .setParameter(SPARK_TABLE_PROVIDER_KEY, DELTA_LAKE_PROVIDER); - table.getStorageBuilder() - .setStorageFormat(fromHiveStorageFormat(PARQUET)) - .setLocation(getTableDefaultLocation( - metastoreClient.getDatabase(tableName.getSchemaName()).orElseThrow(), - new HdfsContext(session.getIdentity()), - hdfsEnvironment, - tableName.getSchemaName(), - tableName.getTableName()).toString()); - metastoreClient.createTable(table.build(), NO_PRIVILEGES); - - try { - // Verify the table was created as a Delta Lake table - try (Transaction transaction = newTransaction()) { - ConnectorMetadata metadata = transaction.getMetadata(); - metadata.beginQuery(session); - assertThatThrownBy(() -> getTableHandle(metadata, tableName)) - .hasMessage(format("Cannot query Delta Lake table '%s'", tableName)); - } - - // Verify the hidden `$properties` and `$partitions` Delta Lake table handle can't be obtained within the hive connector - try (Transaction transaction = newTransaction()) { - ConnectorMetadata metadata = transaction.getMetadata(); - metadata.beginQuery(session); - SchemaTableName propertiesTableName = new SchemaTableName(tableName.getSchemaName(), format("%s$properties", tableName.getTableName())); - assertThat(metadata.getSystemTable(newSession(), propertiesTableName)).isEmpty(); - SchemaTableName partitionsTableName = new SchemaTableName(tableName.getSchemaName(), format("%s$partitions", tableName.getTableName())); - assertThat(metadata.getSystemTable(newSession(), partitionsTableName)).isEmpty(); - } - - // Assert that table is hidden - try (Transaction transaction = newTransaction()) { - ConnectorMetadata metadata = transaction.getMetadata(); - - // TODO (https://github.com/trinodb/trino/issues/5426) these assertions should use information_schema instead of metadata directly, - // as information_schema or MetadataManager may apply additional logic - - // list all tables - assertThat(metadata.listTables(session, Optional.empty())) - .doesNotContain(tableName); - - // list all tables in a schema - assertThat(metadata.listTables(session, Optional.of(tableName.getSchemaName()))) - .doesNotContain(tableName); - - // list all columns - assertThat(listTableColumns(metadata, session, new SchemaTablePrefix()).keySet()) - .doesNotContain(tableName); - - // list all columns in a schema - assertThat(listTableColumns(metadata, session, new SchemaTablePrefix(tableName.getSchemaName())).keySet()) - .doesNotContain(tableName); - - // list all columns in a table - assertThat(listTableColumns(metadata, session, new SchemaTablePrefix(tableName.getSchemaName(), tableName.getTableName())).keySet()) - .doesNotContain(tableName); - } - } - finally { - // Clean up - metastoreClient.dropTable(tableName.getSchemaName(), tableName.getTableName(), true); - } - } - - @Test - public void testDisallowQueryingOfIcebergTables() - { - ConnectorSession session = newSession(); - SchemaTableName tableName = temporaryTable("trino_iceberg_table"); - - Table.Builder table = Table.builder() - .setDatabaseName(tableName.getSchemaName()) - .setTableName(tableName.getTableName()) - .setOwner(Optional.of(session.getUser())) - .setTableType(MANAGED_TABLE.name()) - .setPartitionColumns(List.of(new Column("a_partition_column", HIVE_INT, Optional.empty()))) - .setDataColumns(List.of(new Column("a_column", HIVE_STRING, Optional.empty()))) - .setParameter(ICEBERG_TABLE_TYPE_NAME, ICEBERG_TABLE_TYPE_VALUE); - table.getStorageBuilder() - .setStorageFormat(fromHiveStorageFormat(PARQUET)) - .setLocation(getTableDefaultLocation( - metastoreClient.getDatabase(tableName.getSchemaName()).orElseThrow(), - new HdfsContext(session.getIdentity()), - hdfsEnvironment, - tableName.getSchemaName(), - tableName.getTableName()).toString()); - metastoreClient.createTable(table.build(), NO_PRIVILEGES); - - try { - // Verify that the table was created as a Iceberg table can't be queried in hive - try (Transaction transaction = newTransaction()) { - ConnectorMetadata metadata = transaction.getMetadata(); - metadata.beginQuery(session); - assertThatThrownBy(() -> getTableHandle(metadata, tableName)) - .hasMessage(format("Cannot query Iceberg table '%s'", tableName)); - } - - // Verify the hidden `$properties` and `$partitions` hive system tables table handle can't be obtained for the Iceberg tables - try (Transaction transaction = newTransaction()) { - ConnectorMetadata metadata = transaction.getMetadata(); - metadata.beginQuery(session); - SchemaTableName propertiesTableName = new SchemaTableName(tableName.getSchemaName(), format("%s$properties", tableName.getTableName())); - assertThat(metadata.getSystemTable(newSession(), propertiesTableName)).isEmpty(); - SchemaTableName partitionsTableName = new SchemaTableName(tableName.getSchemaName(), format("%s$partitions", tableName.getTableName())); - assertThat(metadata.getSystemTable(newSession(), partitionsTableName)).isEmpty(); - } - } - finally { - // Clean up - metastoreClient.dropTable(tableName.getSchemaName(), tableName.getTableName(), true); - } - } - - @Test - public void testUpdateBasicTableStatistics() - throws Exception - { - SchemaTableName tableName = temporaryTable("update_basic_table_statistics"); - try { - doCreateEmptyTable(tableName, ORC, STATISTICS_TABLE_COLUMNS); - testUpdateTableStatistics(tableName, ZERO_TABLE_STATISTICS, BASIC_STATISTICS_1, BASIC_STATISTICS_2); - } - finally { - dropTable(tableName); - } - } - - @Test - public void testUpdateTableColumnStatistics() - throws Exception - { - SchemaTableName tableName = temporaryTable("update_table_column_statistics"); - try { - doCreateEmptyTable(tableName, ORC, STATISTICS_TABLE_COLUMNS); - testUpdateTableStatistics(tableName, ZERO_TABLE_STATISTICS, STATISTICS_1_1, STATISTICS_1_2, STATISTICS_2); - } - finally { - dropTable(tableName); - } - } - - @Test - public void testUpdateTableColumnStatisticsEmptyOptionalFields() - throws Exception - { - SchemaTableName tableName = temporaryTable("update_table_column_statistics_empty_optional_fields"); - try { - doCreateEmptyTable(tableName, ORC, STATISTICS_TABLE_COLUMNS); - testUpdateTableStatistics(tableName, ZERO_TABLE_STATISTICS, STATISTICS_EMPTY_OPTIONAL_FIELDS); - } - finally { - dropTable(tableName); - } - } - - protected void testUpdateTableStatistics(SchemaTableName tableName, PartitionStatistics initialStatistics, PartitionStatistics... statistics) - { - HiveMetastoreClosure metastoreClient = new HiveMetastoreClosure(getMetastoreClient()); - assertThat(metastoreClient.getTableStatistics(tableName.getSchemaName(), tableName.getTableName(), Optional.empty())) - .isEqualTo(initialStatistics); - - AtomicReference expectedStatistics = new AtomicReference<>(initialStatistics); - for (PartitionStatistics partitionStatistics : statistics) { - metastoreClient.updateTableStatistics(tableName.getSchemaName(), tableName.getTableName(), NO_ACID_TRANSACTION, actualStatistics -> { - assertThat(actualStatistics).isEqualTo(expectedStatistics.get()); - return partitionStatistics; - }); - assertThat(metastoreClient.getTableStatistics(tableName.getSchemaName(), tableName.getTableName(), Optional.empty())) - .isEqualTo(partitionStatistics); - expectedStatistics.set(partitionStatistics); - } - - assertThat(metastoreClient.getTableStatistics(tableName.getSchemaName(), tableName.getTableName(), Optional.empty())) - .isEqualTo(expectedStatistics.get()); - - metastoreClient.updateTableStatistics(tableName.getSchemaName(), tableName.getTableName(), NO_ACID_TRANSACTION, actualStatistics -> { - assertThat(actualStatistics).isEqualTo(expectedStatistics.get()); - return initialStatistics; - }); - - assertThat(metastoreClient.getTableStatistics(tableName.getSchemaName(), tableName.getTableName(), Optional.empty())) - .isEqualTo(initialStatistics); - } - - @Test - public void testUpdateBasicPartitionStatistics() - throws Exception - { - SchemaTableName tableName = temporaryTable("update_basic_partition_statistics"); - try { - createDummyPartitionedTable(tableName, STATISTICS_PARTITIONED_TABLE_COLUMNS); - testUpdatePartitionStatistics( - tableName, - ZERO_TABLE_STATISTICS, - ImmutableList.of(BASIC_STATISTICS_1, BASIC_STATISTICS_2), - ImmutableList.of(BASIC_STATISTICS_2, BASIC_STATISTICS_1)); - } - finally { - dropTable(tableName); - } - } - - @Test - public void testUpdatePartitionColumnStatistics() - throws Exception - { - SchemaTableName tableName = temporaryTable("update_partition_column_statistics"); - try { - createDummyPartitionedTable(tableName, STATISTICS_PARTITIONED_TABLE_COLUMNS); - testUpdatePartitionStatistics( - tableName, - ZERO_TABLE_STATISTICS, - ImmutableList.of(STATISTICS_1_1, STATISTICS_1_2, STATISTICS_2), - ImmutableList.of(STATISTICS_1_2, STATISTICS_1_1, STATISTICS_2)); - } - finally { - dropTable(tableName); - } - } - - @Test - public void testUpdatePartitionColumnStatisticsEmptyOptionalFields() - throws Exception - { - SchemaTableName tableName = temporaryTable("update_partition_column_statistics"); - try { - createDummyPartitionedTable(tableName, STATISTICS_PARTITIONED_TABLE_COLUMNS); - testUpdatePartitionStatistics( - tableName, - ZERO_TABLE_STATISTICS, - ImmutableList.of(STATISTICS_EMPTY_OPTIONAL_FIELDS), - ImmutableList.of(STATISTICS_EMPTY_OPTIONAL_FIELDS)); - } - finally { - dropTable(tableName); - } - } - - @Test - public void testDataColumnProperties() - throws Exception - { - SchemaTableName tableName = temporaryTable("test_column_properties"); - HiveMetastoreClosure metastoreClient = new HiveMetastoreClosure(getMetastoreClient()); - try { - doCreateEmptyTable(tableName, ORC, List.of(new ColumnMetadata("id", BIGINT), new ColumnMetadata("part_key", createVarcharType(256)))); - - Table table = metastoreClient.getTable(tableName.getSchemaName(), tableName.getTableName()).orElseThrow(); - assertThat(table.getDataColumns()) - .singleElement() - .extracting(Column::getProperties, InstanceOfAssertFactories.MAP) - .isEmpty(); - assertThat(table.getPartitionColumns()) - .singleElement() - .extracting(Column::getProperties, InstanceOfAssertFactories.MAP) - .isEmpty(); - - String columnPropertyValue = "data column value ,;.!??? \" ' {} [] non-printable \000 \001 spaces \n\r\t\f hiragana だ emoji 🤷‍♂️ x"; - metastoreClient.replaceTable( - tableName.getSchemaName(), - tableName.getTableName(), - Table.builder(table) - .setDataColumns(List.of(new Column("id", HIVE_LONG, Optional.empty(), Map.of("data prop", columnPropertyValue)))) - .build(), - NO_PRIVILEGES); - - table = metastoreClient.getTable(tableName.getSchemaName(), tableName.getTableName()).orElseThrow(); - assertThat(table.getDataColumns()) - .singleElement() - .extracting(Column::getProperties, InstanceOfAssertFactories.MAP) - .isEqualTo(Map.of("data prop", columnPropertyValue)); - assertThat(table.getPartitionColumns()) - .singleElement() - .extracting(Column::getProperties, InstanceOfAssertFactories.MAP) - .isEmpty(); - } - finally { - dropTable(tableName); - } - } - - @Test - public void testPartitionColumnProperties() - throws Exception - { - SchemaTableName tableName = temporaryTable("test_column_properties"); - HiveMetastoreClosure metastoreClient = new HiveMetastoreClosure(getMetastoreClient()); - try { - doCreateEmptyTable(tableName, ORC, List.of(new ColumnMetadata("id", BIGINT), new ColumnMetadata("part_key", createVarcharType(256)))); - - Table table = metastoreClient.getTable(tableName.getSchemaName(), tableName.getTableName()).orElseThrow(); - assertThat(table.getDataColumns()) - .singleElement() - .extracting(Column::getProperties, InstanceOfAssertFactories.MAP) - .isEmpty(); - assertThat(table.getPartitionColumns()) - .singleElement() - .extracting(Column::getProperties, InstanceOfAssertFactories.MAP) - .isEmpty(); - - String columnPropertyValue = "partition column value ,;.!??? \" ' {} [] non-printable \000 \001 spaces \n\r\t\f hiragana だ emoji 🤷‍♂️ x"; - metastoreClient.replaceTable( - tableName.getSchemaName(), - tableName.getTableName(), - Table.builder(table) - .setPartitionColumns(List.of(new Column("part_key", HiveType.valueOf("varchar(256)"), Optional.empty(), Map.of("partition prop", columnPropertyValue)))) - .build(), - NO_PRIVILEGES); - - table = metastoreClient.getTable(tableName.getSchemaName(), tableName.getTableName()).orElseThrow(); - assertThat(table.getDataColumns()) - .singleElement() - .extracting(Column::getProperties, InstanceOfAssertFactories.MAP) - .isEmpty(); - assertThat(table.getPartitionColumns()) - .singleElement() - .extracting(Column::getProperties, InstanceOfAssertFactories.MAP) - .isEqualTo(Map.of("partition prop", columnPropertyValue)); - } - finally { - dropTable(tableName); - } - } - - @Test - public void testInputInfoWhenTableIsPartitioned() - throws Exception - { - SchemaTableName tableName = temporaryTable("test_input_info_with_partitioned_table"); - try { - createDummyPartitionedTable(tableName, STATISTICS_PARTITIONED_TABLE_COLUMNS); - assertInputInfo(tableName, new HiveInputInfo(ImmutableList.of(), true, Optional.of("ORC"))); - } - finally { - dropTable(tableName); - } - } - - @Test - public void testInputInfoWhenTableIsNotPartitioned() - { - SchemaTableName tableName = temporaryTable("test_input_info_without_partitioned_table"); - try { - createDummyTable(tableName); - assertInputInfo(tableName, new HiveInputInfo(ImmutableList.of(), false, Optional.of("TEXTFILE"))); - } - finally { - dropTable(tableName); - } - } - - @Test - public void testInputInfoWithParquetTableFormat() - { - SchemaTableName tableName = temporaryTable("test_input_info_with_parquet_table_format"); - try { - createDummyTable(tableName, PARQUET); - assertInputInfo(tableName, new HiveInputInfo(ImmutableList.of(), false, Optional.of("PARQUET"))); - } - finally { - dropTable(tableName); - } - } - - private void assertInputInfo(SchemaTableName tableName, HiveInputInfo expectedInputInfo) - { - try (Transaction transaction = newTransaction()) { - ConnectorSession session = newSession(); - ConnectorMetadata metadata = transaction.getMetadata(); - HiveTableHandle tableHandle = (HiveTableHandle) metadata.getTableHandle(session, tableName); - assertThat(metadata.getInfo(tableHandle)).isEqualTo(Optional.of(expectedInputInfo)); - } - } - - /** - * During table scan, the illegal storage format for some specific table should not fail the whole table scan - */ - @Test - public void testIllegalStorageFormatDuringTableScan() - { - SchemaTableName schemaTableName = temporaryTable("test_illegal_storage_format"); - try (Transaction transaction = newTransaction()) { - ConnectorSession session = newSession(); - List columns = ImmutableList.of(new Column("pk", HIVE_STRING, Optional.empty())); - String tableOwner = session.getUser(); - String schemaName = schemaTableName.getSchemaName(); - String tableName = schemaTableName.getTableName(); - Location targetPath = locationService.forNewTable(transaction.getMetastore(), session, schemaName, tableName); - //create table whose storage format is null - Table.Builder tableBuilder = Table.builder() - .setDatabaseName(schemaName) - .setTableName(tableName) - .setOwner(Optional.of(tableOwner)) - .setTableType(TableType.MANAGED_TABLE.name()) - .setParameters(ImmutableMap.of( - PRESTO_VERSION_NAME, TEST_SERVER_VERSION, - PRESTO_QUERY_ID_NAME, session.getQueryId())) - .setDataColumns(columns) - .withStorage(storage -> storage - .setLocation(targetPath.toString()) - .setStorageFormat(StorageFormat.createNullable(null, null, null)) - .setSerdeParameters(ImmutableMap.of())); - PrincipalPrivileges principalPrivileges = testingPrincipalPrivilege(tableOwner, session.getUser()); - transaction.getMetastore().createTable(session, tableBuilder.build(), principalPrivileges, Optional.empty(), Optional.empty(), true, ZERO_TABLE_STATISTICS, false); - transaction.commit(); - } - - // We retrieve the table whose storageFormat has null serde/inputFormat/outputFormat - // to make sure it can still be retrieved instead of throwing exception. - try (Transaction transaction = newTransaction()) { - ConnectorMetadata metadata = transaction.getMetadata(); - Map> allColumns = listTableColumns(metadata, newSession(), new SchemaTablePrefix(schemaTableName.getSchemaName())); - assertTrue(allColumns.containsKey(schemaTableName)); - } - finally { - dropTable(schemaTableName); - } - } - - protected static Map> listTableColumns(ConnectorMetadata metadata, ConnectorSession session, SchemaTablePrefix prefix) - { - return stream(metadata.streamTableColumns(session, prefix)) - .collect(toImmutableMap( - TableColumnsMetadata::getTable, - tableColumns -> tableColumns.getColumns().orElseThrow(() -> new IllegalStateException("Table " + tableColumns.getTable() + " reported as redirected")))); - } - - private void createDummyTable(SchemaTableName tableName) - { - createDummyTable(tableName, TEXTFILE); - } - - private void createDummyTable(SchemaTableName tableName, HiveStorageFormat storageFormat) - { - try (Transaction transaction = newTransaction()) { - ConnectorSession session = newSession(); - ConnectorMetadata metadata = transaction.getMetadata(); - - List columns = ImmutableList.of(new ColumnMetadata("dummy", createUnboundedVarcharType())); - ConnectorTableMetadata tableMetadata = new ConnectorTableMetadata(tableName, columns, createTableProperties(storageFormat)); - ConnectorOutputTableHandle handle = metadata.beginCreateTable(session, tableMetadata, Optional.empty(), NO_RETRIES); - metadata.finishCreateTable(session, handle, ImmutableList.of(), ImmutableList.of()); - - transaction.commit(); - } - } - - protected void createDummyPartitionedTable(SchemaTableName tableName, List columns) - throws Exception - { - doCreateEmptyTable(tableName, ORC, columns); - - HiveMetastoreClosure metastoreClient = new HiveMetastoreClosure(getMetastoreClient()); - Table table = metastoreClient.getTable(tableName.getSchemaName(), tableName.getTableName()) - .orElseThrow(() -> new TableNotFoundException(tableName)); - - List firstPartitionValues = ImmutableList.of("2016-01-01"); - List secondPartitionValues = ImmutableList.of("2016-01-02"); - - String firstPartitionName = makePartName(ImmutableList.of("ds"), firstPartitionValues); - String secondPartitionName = makePartName(ImmutableList.of("ds"), secondPartitionValues); - - List partitions = ImmutableList.of(firstPartitionName, secondPartitionName) - .stream() - .map(partitionName -> new PartitionWithStatistics(createDummyPartition(table, partitionName), partitionName, PartitionStatistics.empty())) - .collect(toImmutableList()); - metastoreClient.addPartitions(tableName.getSchemaName(), tableName.getTableName(), partitions); - metastoreClient.updatePartitionStatistics(tableName.getSchemaName(), tableName.getTableName(), firstPartitionName, currentStatistics -> ZERO_TABLE_STATISTICS); - metastoreClient.updatePartitionStatistics(tableName.getSchemaName(), tableName.getTableName(), secondPartitionName, currentStatistics -> ZERO_TABLE_STATISTICS); - } - - protected void testUpdatePartitionStatistics( - SchemaTableName tableName, - PartitionStatistics initialStatistics, - List firstPartitionStatistics, - List secondPartitionStatistics) - { - verify(firstPartitionStatistics.size() == secondPartitionStatistics.size()); - - String firstPartitionName = "ds=2016-01-01"; - String secondPartitionName = "ds=2016-01-02"; - - HiveMetastoreClosure metastoreClient = new HiveMetastoreClosure(getMetastoreClient()); - assertThat(metastoreClient.getPartitionStatistics(tableName.getSchemaName(), tableName.getTableName(), ImmutableSet.of(firstPartitionName, secondPartitionName))) - .isEqualTo(ImmutableMap.of(firstPartitionName, initialStatistics, secondPartitionName, initialStatistics)); - - AtomicReference expectedStatisticsPartition1 = new AtomicReference<>(initialStatistics); - AtomicReference expectedStatisticsPartition2 = new AtomicReference<>(initialStatistics); - - for (int i = 0; i < firstPartitionStatistics.size(); i++) { - PartitionStatistics statisticsPartition1 = firstPartitionStatistics.get(i); - PartitionStatistics statisticsPartition2 = secondPartitionStatistics.get(i); - metastoreClient.updatePartitionStatistics(tableName.getSchemaName(), tableName.getTableName(), firstPartitionName, actualStatistics -> { - assertThat(actualStatistics).isEqualTo(expectedStatisticsPartition1.get()); - return statisticsPartition1; - }); - metastoreClient.updatePartitionStatistics(tableName.getSchemaName(), tableName.getTableName(), secondPartitionName, actualStatistics -> { - assertThat(actualStatistics).isEqualTo(expectedStatisticsPartition2.get()); - return statisticsPartition2; - }); - assertThat(metastoreClient.getPartitionStatistics(tableName.getSchemaName(), tableName.getTableName(), ImmutableSet.of(firstPartitionName, secondPartitionName))) - .isEqualTo(ImmutableMap.of(firstPartitionName, statisticsPartition1, secondPartitionName, statisticsPartition2)); - expectedStatisticsPartition1.set(statisticsPartition1); - expectedStatisticsPartition2.set(statisticsPartition2); - } - - assertThat(metastoreClient.getPartitionStatistics(tableName.getSchemaName(), tableName.getTableName(), ImmutableSet.of(firstPartitionName, secondPartitionName))) - .isEqualTo(ImmutableMap.of(firstPartitionName, expectedStatisticsPartition1.get(), secondPartitionName, expectedStatisticsPartition2.get())); - metastoreClient.updatePartitionStatistics(tableName.getSchemaName(), tableName.getTableName(), firstPartitionName, currentStatistics -> { - assertThat(currentStatistics).isEqualTo(expectedStatisticsPartition1.get()); - return initialStatistics; - }); - metastoreClient.updatePartitionStatistics(tableName.getSchemaName(), tableName.getTableName(), secondPartitionName, currentStatistics -> { - assertThat(currentStatistics).isEqualTo(expectedStatisticsPartition2.get()); - return initialStatistics; - }); - assertThat(metastoreClient.getPartitionStatistics(tableName.getSchemaName(), tableName.getTableName(), ImmutableSet.of(firstPartitionName, secondPartitionName))) - .isEqualTo(ImmutableMap.of(firstPartitionName, initialStatistics, secondPartitionName, initialStatistics)); - } - - @Test - public void testStorePartitionWithStatistics() - throws Exception - { - testStorePartitionWithStatistics(STATISTICS_PARTITIONED_TABLE_COLUMNS, STATISTICS_1, STATISTICS_2, STATISTICS_1_1, ZERO_TABLE_STATISTICS); - } - - protected void testStorePartitionWithStatistics( - List columns, - PartitionStatistics statsForAllColumns1, - PartitionStatistics statsForAllColumns2, - PartitionStatistics statsForSubsetOfColumns, - PartitionStatistics emptyStatistics) - throws Exception - { - SchemaTableName tableName = temporaryTable("store_partition_with_statistics"); - try { - doCreateEmptyTable(tableName, ORC, columns); - - HiveMetastoreClosure metastoreClient = new HiveMetastoreClosure(getMetastoreClient()); - Table table = metastoreClient.getTable(tableName.getSchemaName(), tableName.getTableName()) - .orElseThrow(() -> new TableNotFoundException(tableName)); - - List partitionValues = ImmutableList.of("2016-01-01"); - String partitionName = makePartName(ImmutableList.of("ds"), partitionValues); - - Partition partition = createDummyPartition(table, partitionName); - - // create partition with stats for all columns - metastoreClient.addPartitions(tableName.getSchemaName(), tableName.getTableName(), ImmutableList.of(new PartitionWithStatistics(partition, partitionName, statsForAllColumns1))); - assertEquals( - metastoreClient.getPartition(tableName.getSchemaName(), tableName.getTableName(), partitionValues).get().getStorage().getStorageFormat(), - fromHiveStorageFormat(ORC)); - assertThat(metastoreClient.getPartitionStatistics(tableName.getSchemaName(), tableName.getTableName(), ImmutableSet.of(partitionName))) - .isEqualTo(ImmutableMap.of(partitionName, statsForAllColumns1)); - - // alter the partition into one with other stats - Partition modifiedPartition = Partition.builder(partition) - .withStorage(storage -> storage - .setStorageFormat(fromHiveStorageFormat(RCBINARY)) - .setLocation(partitionTargetPath(tableName, partitionName))) - .build(); - metastoreClient.alterPartition(tableName.getSchemaName(), tableName.getTableName(), new PartitionWithStatistics(modifiedPartition, partitionName, statsForAllColumns2)); - assertEquals( - metastoreClient.getPartition(tableName.getSchemaName(), tableName.getTableName(), partitionValues).get().getStorage().getStorageFormat(), - fromHiveStorageFormat(RCBINARY)); - assertThat(metastoreClient.getPartitionStatistics(tableName.getSchemaName(), tableName.getTableName(), ImmutableSet.of(partitionName))) - .isEqualTo(ImmutableMap.of(partitionName, statsForAllColumns2)); - - // alter the partition into one with stats for only subset of columns - modifiedPartition = Partition.builder(partition) - .withStorage(storage -> storage - .setStorageFormat(fromHiveStorageFormat(TEXTFILE)) - .setLocation(partitionTargetPath(tableName, partitionName))) - .build(); - metastoreClient.alterPartition(tableName.getSchemaName(), tableName.getTableName(), new PartitionWithStatistics(modifiedPartition, partitionName, statsForSubsetOfColumns)); - assertThat(metastoreClient.getPartitionStatistics(tableName.getSchemaName(), tableName.getTableName(), ImmutableSet.of(partitionName))) - .isEqualTo(ImmutableMap.of(partitionName, statsForSubsetOfColumns)); - - // alter the partition into one without stats - modifiedPartition = Partition.builder(partition) - .withStorage(storage -> storage - .setStorageFormat(fromHiveStorageFormat(TEXTFILE)) - .setLocation(partitionTargetPath(tableName, partitionName))) - .build(); - metastoreClient.alterPartition(tableName.getSchemaName(), tableName.getTableName(), new PartitionWithStatistics(modifiedPartition, partitionName, emptyStatistics)); - assertThat(metastoreClient.getPartitionStatistics(tableName.getSchemaName(), tableName.getTableName(), ImmutableSet.of(partitionName))) - .isEqualTo(ImmutableMap.of(partitionName, emptyStatistics)); - } - finally { - dropTable(tableName); - } - } - - protected Partition createDummyPartition(Table table, String partitionName) - { - return Partition.builder() - .setDatabaseName(table.getDatabaseName()) - .setTableName(table.getTableName()) - .setColumns(table.getDataColumns()) - .setValues(toPartitionValues(partitionName)) - .withStorage(storage -> storage - .setStorageFormat(fromHiveStorageFormat(ORC)) - .setLocation(partitionTargetPath(new SchemaTableName(table.getDatabaseName(), table.getTableName()), partitionName))) - .setParameters(ImmutableMap.of( - PRESTO_VERSION_NAME, "testversion", - PRESTO_QUERY_ID_NAME, "20180101_123456_00001_x1y2z")) - .build(); - } - - protected String partitionTargetPath(SchemaTableName schemaTableName, String partitionName) - { - try (Transaction transaction = newTransaction()) { - ConnectorSession session = newSession(); - SemiTransactionalHiveMetastore metastore = transaction.getMetastore(); - LocationService locationService = getLocationService(); - Table table = metastore.getTable(schemaTableName.getSchemaName(), schemaTableName.getTableName()).get(); - LocationHandle handle = locationService.forExistingTable(metastore, session, table); - return locationService.getPartitionWriteInfo(handle, Optional.empty(), partitionName).targetPath().toString(); - } - } - - /** - * This test creates 2 identical partitions and verifies that the statistics projected based on - * a single partition sample are equal to the statistics computed in a fair way - */ - @Test - public void testPartitionStatisticsSampling() - throws Exception - { - testPartitionStatisticsSampling(STATISTICS_PARTITIONED_TABLE_COLUMNS, STATISTICS_1); - } - - protected void testPartitionStatisticsSampling(List columns, PartitionStatistics statistics) - throws Exception - { - SchemaTableName tableName = temporaryTable("test_partition_statistics_sampling"); - - try { - createDummyPartitionedTable(tableName, columns); - HiveMetastoreClosure metastoreClient = new HiveMetastoreClosure(getMetastoreClient()); - metastoreClient.updatePartitionStatistics(tableName.getSchemaName(), tableName.getTableName(), "ds=2016-01-01", actualStatistics -> statistics); - metastoreClient.updatePartitionStatistics(tableName.getSchemaName(), tableName.getTableName(), "ds=2016-01-02", actualStatistics -> statistics); - - try (Transaction transaction = newTransaction()) { - ConnectorSession session = newSession(); - ConnectorMetadata metadata = transaction.getMetadata(); - - ConnectorTableHandle tableHandle = metadata.getTableHandle(session, tableName); - TableStatistics unsampledStatistics = metadata.getTableStatistics(sampleSize(2), tableHandle); - TableStatistics sampledStatistics = metadata.getTableStatistics(sampleSize(1), tableHandle); - assertEquals(sampledStatistics, unsampledStatistics); - } - } - finally { - dropTable(tableName); - } - } - - @Test - public void testApplyProjection() - throws Exception - { - ColumnMetadata bigIntColumn0 = new ColumnMetadata("int0", BIGINT); - ColumnMetadata bigIntColumn1 = new ColumnMetadata("int1", BIGINT); - - RowType oneLevelRowType = toRowType(ImmutableList.of(bigIntColumn0, bigIntColumn1)); - ColumnMetadata oneLevelRow0 = new ColumnMetadata("onelevelrow0", oneLevelRowType); - - RowType twoLevelRowType = toRowType(ImmutableList.of(oneLevelRow0, bigIntColumn0, bigIntColumn1)); - ColumnMetadata twoLevelRow0 = new ColumnMetadata("twolevelrow0", twoLevelRowType); - - List columnsForApplyProjectionTest = ImmutableList.of(bigIntColumn0, bigIntColumn1, oneLevelRow0, twoLevelRow0); - - SchemaTableName tableName = temporaryTable("apply_projection_tester"); - doCreateEmptyTable(tableName, ORC, columnsForApplyProjectionTest); - - try (Transaction transaction = newTransaction()) { - ConnectorSession session = newSession(); - ConnectorMetadata metadata = transaction.getMetadata(); - ConnectorTableHandle tableHandle = getTableHandle(metadata, tableName); - - List columnHandles = metadata.getColumnHandles(session, tableHandle).values().stream() - .filter(columnHandle -> !((HiveColumnHandle) columnHandle).isHidden()) - .collect(toList()); - assertEquals(columnHandles.size(), columnsForApplyProjectionTest.size()); - - Map columnHandleMap = columnHandles.stream() - .collect(toImmutableMap(handle -> ((HiveColumnHandle) handle).getBaseColumnName(), Function.identity())); - - // Emulate symbols coming from the query plan and map them to column handles - Map columnHandlesWithSymbols = ImmutableMap.of( - "symbol_0", columnHandleMap.get("int0"), - "symbol_1", columnHandleMap.get("int1"), - "symbol_2", columnHandleMap.get("onelevelrow0"), - "symbol_3", columnHandleMap.get("twolevelrow0")); - - // Create variables for the emulated symbols - Map symbolVariableMapping = columnHandlesWithSymbols.entrySet().stream() - .collect(toImmutableMap( - Map.Entry::getKey, - e -> new Variable( - e.getKey(), - ((HiveColumnHandle) e.getValue()).getBaseType()))); - - // Create dereference expressions for testing - FieldDereference symbol2Field0 = new FieldDereference(BIGINT, symbolVariableMapping.get("symbol_2"), 0); - FieldDereference symbol3Field0 = new FieldDereference(oneLevelRowType, symbolVariableMapping.get("symbol_3"), 0); - FieldDereference symbol3Field0Field0 = new FieldDereference(BIGINT, symbol3Field0, 0); - FieldDereference symbol3Field1 = new FieldDereference(BIGINT, symbolVariableMapping.get("symbol_3"), 1); - - Map inputAssignments; - List inputProjections; - Optional> projectionResult; - List expectedProjections; - Map expectedAssignments; - - // Test projected columns pushdown to HiveTableHandle in case of all variable references - inputAssignments = getColumnHandlesFor(columnHandlesWithSymbols, ImmutableList.of("symbol_0", "symbol_1")); - inputProjections = ImmutableList.of(symbolVariableMapping.get("symbol_0"), symbolVariableMapping.get("symbol_1")); - expectedAssignments = ImmutableMap.of( - "symbol_0", BIGINT, - "symbol_1", BIGINT); - projectionResult = metadata.applyProjection(session, tableHandle, inputProjections, inputAssignments); - assertProjectionResult(projectionResult, false, inputProjections, expectedAssignments); - - // Empty result when projected column handles are same as those present in table handle - projectionResult = metadata.applyProjection(session, projectionResult.get().getHandle(), inputProjections, inputAssignments); - assertProjectionResult(projectionResult, true, ImmutableList.of(), ImmutableMap.of()); - - // Extra columns handles in HiveTableHandle should get pruned - projectionResult = metadata.applyProjection( - session, - ((HiveTableHandle) tableHandle).withProjectedColumns(ImmutableSet.copyOf(columnHandles)), - inputProjections, - inputAssignments); - assertProjectionResult(projectionResult, false, inputProjections, expectedAssignments); - - // Test projection pushdown for dereferences - inputAssignments = getColumnHandlesFor(columnHandlesWithSymbols, ImmutableList.of("symbol_2", "symbol_3")); - inputProjections = ImmutableList.of(symbol2Field0, symbol3Field0Field0, symbol3Field1); - expectedAssignments = ImmutableMap.of( - "onelevelrow0#f_int0", BIGINT, - "twolevelrow0#f_onelevelrow0#f_int0", BIGINT, - "twolevelrow0#f_int0", BIGINT); - expectedProjections = ImmutableList.of( - new Variable("onelevelrow0#f_int0", BIGINT), - new Variable("twolevelrow0#f_onelevelrow0#f_int0", BIGINT), - new Variable("twolevelrow0#f_int0", BIGINT)); - projectionResult = metadata.applyProjection(session, tableHandle, inputProjections, inputAssignments); - assertProjectionResult(projectionResult, false, expectedProjections, expectedAssignments); - - // Test reuse of virtual column handles - // Round-1: input projections [symbol_2, symbol_2.int0]. virtual handle is created for symbol_2.int0. - inputAssignments = getColumnHandlesFor(columnHandlesWithSymbols, ImmutableList.of("symbol_2")); - inputProjections = ImmutableList.of(symbol2Field0, symbolVariableMapping.get("symbol_2")); - projectionResult = metadata.applyProjection(session, tableHandle, inputProjections, inputAssignments); - expectedProjections = ImmutableList.of(new Variable("onelevelrow0#f_int0", BIGINT), symbolVariableMapping.get("symbol_2")); - expectedAssignments = ImmutableMap.of("onelevelrow0#f_int0", BIGINT, "symbol_2", oneLevelRowType); - assertProjectionResult(projectionResult, false, expectedProjections, expectedAssignments); - - // Round-2: input projections [symbol_2.int0 and onelevelrow0#f_int0]. Virtual handle is reused. - Assignment newlyCreatedColumn = getOnlyElement(projectionResult.get().getAssignments().stream() - .filter(handle -> handle.getVariable().equals("onelevelrow0#f_int0")) - .collect(toList())); - inputAssignments = ImmutableMap.builder() - .putAll(getColumnHandlesFor(columnHandlesWithSymbols, ImmutableList.of("symbol_2"))) - .put(newlyCreatedColumn.getVariable(), newlyCreatedColumn.getColumn()) - .buildOrThrow(); - inputProjections = ImmutableList.of(symbol2Field0, new Variable("onelevelrow0#f_int0", BIGINT)); - projectionResult = metadata.applyProjection(session, tableHandle, inputProjections, inputAssignments); - expectedProjections = ImmutableList.of(new Variable("onelevelrow0#f_int0", BIGINT), new Variable("onelevelrow0#f_int0", BIGINT)); - expectedAssignments = ImmutableMap.of("onelevelrow0#f_int0", BIGINT); - assertProjectionResult(projectionResult, false, expectedProjections, expectedAssignments); - } - finally { - dropTable(tableName); - } - } - - private static Map getColumnHandlesFor(Map columnHandles, List symbols) - { - return columnHandles.entrySet().stream() - .filter(e -> symbols.contains(e.getKey())) - .collect(toImmutableMap(Map.Entry::getKey, Map.Entry::getValue)); - } - - private static void assertProjectionResult(Optional> projectionResult, boolean shouldBeEmpty, List expectedProjections, Map expectedAssignments) - { - if (shouldBeEmpty) { - assertTrue(projectionResult.isEmpty(), "expected projectionResult to be empty"); - return; - } - - assertTrue(projectionResult.isPresent(), "expected non-empty projection result"); - - ProjectionApplicationResult result = projectionResult.get(); - - // Verify projections - assertEquals(expectedProjections, result.getProjections()); - - // Verify assignments - List assignments = result.getAssignments(); - Map actualAssignments = uniqueIndex(assignments, Assignment::getVariable); - - for (String variable : expectedAssignments.keySet()) { - Type expectedType = expectedAssignments.get(variable); - assertTrue(actualAssignments.containsKey(variable)); - assertEquals(actualAssignments.get(variable).getType(), expectedType); - assertEquals(((HiveColumnHandle) actualAssignments.get(variable).getColumn()).getType(), expectedType); - } - - assertEquals(actualAssignments.size(), expectedAssignments.size()); - assertEquals( - actualAssignments.values().stream().map(Assignment::getColumn).collect(toImmutableSet()), - ((HiveTableHandle) result.getHandle()).getProjectedColumns()); - } - - @Test - public void testApplyRedirection() - throws Exception - { - SchemaTableName sourceTableName = temporaryTable("apply_redirection_tester"); - doCreateEmptyTable(sourceTableName, ORC, CREATE_TABLE_COLUMNS); - SchemaTableName tableName = temporaryTable("apply_no_redirection_tester"); - doCreateEmptyTable(tableName, ORC, CREATE_TABLE_COLUMNS); - try (Transaction transaction = newTransaction()) { - ConnectorSession session = newSession(); - ConnectorMetadata metadata = transaction.getMetadata(); - assertThat(metadata.applyTableScanRedirect(session, getTableHandle(metadata, tableName))).isEmpty(); - Optional result = metadata.applyTableScanRedirect(session, getTableHandle(metadata, sourceTableName)); - assertThat(result).isPresent(); - assertThat(result.get().getDestinationTable()) - .isEqualTo(new CatalogSchemaTableName("hive", database, "mock_redirection_target")); - } - finally { - dropTable(sourceTableName); - dropTable(tableName); - } - } - - @Test - public void testMaterializedViewMetadata() - throws Exception - { - SchemaTableName sourceTableName = temporaryTable("materialized_view_tester"); - doCreateEmptyTable(sourceTableName, ORC, CREATE_TABLE_COLUMNS); - SchemaTableName tableName = temporaryTable("mock_table"); - doCreateEmptyTable(tableName, ORC, CREATE_TABLE_COLUMNS); - try (Transaction transaction = newTransaction()) { - ConnectorSession session = newSession(); - ConnectorMetadata metadata = transaction.getMetadata(); - assertThat(metadata.getMaterializedView(session, tableName)).isEmpty(); - Optional result = metadata.getMaterializedView(session, sourceTableName); - assertThat(result).isPresent(); - assertThat(result.get().getOriginalSql()).isEqualTo("dummy_view_sql"); - } - finally { - dropTable(sourceTableName); - dropTable(tableName); - } - } - - @Test - public void testOrcPageSourceMetrics() - throws Exception - { - SchemaTableName tableName = temporaryTable("orc_page_source_metrics"); - try { - assertPageSourceMetrics(tableName, ORC, new Metrics(ImmutableMap.of(ORC_CODEC_METRIC_PREFIX + "SNAPPY", new LongCount(209)))); - } - finally { - dropTable(tableName); - } - } - - @Test - public void testParquetPageSourceMetrics() - throws Exception - { - SchemaTableName tableName = temporaryTable("parquet_page_source_metrics"); - try { - assertPageSourceMetrics(tableName, PARQUET, new Metrics(ImmutableMap.of(PARQUET_CODEC_METRIC_PREFIX + "SNAPPY", new LongCount(1157)))); - } - finally { - dropTable(tableName); - } - } - - private void assertPageSourceMetrics(SchemaTableName tableName, HiveStorageFormat storageFormat, Metrics expectedMetrics) - throws Exception - { - createEmptyTable( - tableName, - storageFormat, - ImmutableList.of( - new Column("id", HIVE_LONG, Optional.empty()), - new Column("name", HIVE_STRING, Optional.empty())), - ImmutableList.of()); - MaterializedResult.Builder inputDataBuilder = MaterializedResult.resultBuilder(SESSION, BIGINT, VARCHAR); - IntStream.range(0, 100).forEach(i -> inputDataBuilder.row((long) i, String.valueOf(i))); - insertData(tableName, inputDataBuilder.build(), ImmutableMap.of("compression_codec", "SNAPPY")); - - try (Transaction transaction = newTransaction()) { - ConnectorMetadata metadata = transaction.getMetadata(); - ConnectorSession session = newSession(); - metadata.beginQuery(session); - - ConnectorTableHandle tableHandle = getTableHandle(metadata, tableName); - - // read entire table - List columnHandles = ImmutableList.copyOf(metadata.getColumnHandles(session, tableHandle).values()); - - List splits = getAllSplits(getSplits(splitManager, transaction, session, tableHandle)); - for (ConnectorSplit split : splits) { - try (ConnectorPageSource pageSource = pageSourceProvider.createPageSource(transaction.getTransactionHandle(), session, split, tableHandle, columnHandles, DynamicFilter.EMPTY)) { - materializeSourceDataStream(session, pageSource, getTypes(columnHandles)); - assertThat(pageSource.getMetrics()).isEqualTo(expectedMetrics); - } - } - } - } - - private ConnectorSession sampleSize(int sampleSize) - { - return getHiveSession(getHiveConfig() - .setPartitionStatisticsSampleSize(sampleSize)); - } - - private void verifyViewCreation(SchemaTableName temporaryCreateView) - { - // replace works for new view - doCreateView(temporaryCreateView, true); - - // replace works for existing view - doCreateView(temporaryCreateView, true); - - // create fails for existing view - try { - doCreateView(temporaryCreateView, false); - fail("create existing should fail"); - } - catch (ViewAlreadyExistsException e) { - assertEquals(e.getViewName(), temporaryCreateView); - } - - try (Transaction transaction = newTransaction()) { - ConnectorMetadata metadata = transaction.getMetadata(); - // drop works when view exists - metadata.dropView(newSession(), temporaryCreateView); - transaction.commit(); - } - - try (Transaction transaction = newTransaction()) { - ConnectorMetadata metadata = transaction.getMetadata(); - assertThat(metadata.getView(newSession(), temporaryCreateView)) - .isEmpty(); - assertThat(metadata.getViews(newSession(), Optional.of(temporaryCreateView.getSchemaName()))) - .doesNotContainKey(temporaryCreateView); - assertThat(metadata.listViews(newSession(), Optional.of(temporaryCreateView.getSchemaName()))) - .doesNotContain(temporaryCreateView); - } - - // drop fails when view does not exist - try (Transaction transaction = newTransaction()) { - ConnectorMetadata metadata = transaction.getMetadata(); - metadata.dropView(newSession(), temporaryCreateView); - fail("drop non-existing should fail"); - } - catch (ViewNotFoundException e) { - assertEquals(e.getViewName(), temporaryCreateView); - } - - // create works for new view - doCreateView(temporaryCreateView, false); - } - - private void doCreateView(SchemaTableName viewName, boolean replace) - { - String viewData = "test data"; - ConnectorViewDefinition definition = new ConnectorViewDefinition( - viewData, - Optional.empty(), - Optional.empty(), - ImmutableList.of(new ViewColumn("test", BIGINT.getTypeId(), Optional.empty())), - Optional.empty(), - Optional.empty(), - true); - - try (Transaction transaction = newTransaction()) { - transaction.getMetadata().createView(newSession(), viewName, definition, replace); - transaction.commit(); - } - - try (Transaction transaction = newTransaction()) { - ConnectorMetadata metadata = transaction.getMetadata(); - - assertThat(metadata.getView(newSession(), viewName)) - .map(ConnectorViewDefinition::getOriginalSql) - .contains(viewData); - - Map views = metadata.getViews(newSession(), Optional.of(viewName.getSchemaName())); - assertEquals(views.size(), 1); - assertEquals(views.get(viewName).getOriginalSql(), definition.getOriginalSql()); - - assertTrue(metadata.listViews(newSession(), Optional.of(viewName.getSchemaName())).contains(viewName)); - } - } - - protected void doCreateTable(SchemaTableName tableName, HiveStorageFormat storageFormat) - throws Exception - { - String queryId; - try (Transaction transaction = newTransaction()) { - ConnectorSession session = newSession(); - ConnectorMetadata metadata = transaction.getMetadata(); - queryId = session.getQueryId(); - - // begin creating the table - ConnectorTableMetadata tableMetadata = new ConnectorTableMetadata(tableName, CREATE_TABLE_COLUMNS, createTableProperties(storageFormat)); - - ConnectorOutputTableHandle outputHandle = metadata.beginCreateTable(session, tableMetadata, Optional.empty(), NO_RETRIES); - - // write the data - ConnectorPageSink sink = pageSinkProvider.createPageSink(transaction.getTransactionHandle(), session, outputHandle, TESTING_PAGE_SINK_ID); - sink.appendPage(CREATE_TABLE_DATA.toPage()); - Collection fragments = getFutureValue(sink.finish()); - - // verify all new files start with the unique prefix - HdfsContext context = new HdfsContext(session); - for (String filePath : listAllDataFiles(context, getStagingPathRoot(outputHandle))) { - assertThat(new Path(filePath).getName()).startsWith(session.getQueryId()); - } - - // commit the table - metadata.finishCreateTable(session, outputHandle, fragments, ImmutableList.of()); - - transaction.commit(); - } - - try (Transaction transaction = newTransaction()) { - ConnectorSession session = newSession(); - ConnectorMetadata metadata = transaction.getMetadata(); - metadata.beginQuery(session); - // load the new table - ConnectorTableHandle tableHandle = getTableHandle(metadata, tableName); - List columnHandles = filterNonHiddenColumnHandles(metadata.getColumnHandles(session, tableHandle).values()); - - // verify the metadata - ConnectorTableMetadata tableMetadata = metadata.getTableMetadata(session, getTableHandle(metadata, tableName)); - assertEquals(filterNonHiddenColumnMetadata(tableMetadata.getColumns()), CREATE_TABLE_COLUMNS); - - // verify the data - MaterializedResult result = readTable(transaction, tableHandle, columnHandles, session, TupleDomain.all(), OptionalInt.empty(), Optional.of(storageFormat)); - assertEqualsIgnoreOrder(result.getMaterializedRows(), CREATE_TABLE_DATA.getMaterializedRows()); - - // verify the node version and query ID in table - Table table = getMetastoreClient().getTable(tableName.getSchemaName(), tableName.getTableName()).get(); - assertEquals(table.getParameters().get(PRESTO_VERSION_NAME), TEST_SERVER_VERSION); - assertEquals(table.getParameters().get(PRESTO_QUERY_ID_NAME), queryId); - - // verify basic statistics - HiveBasicStatistics statistics = getBasicStatisticsForTable(transaction, tableName); - assertEquals(statistics.getRowCount().getAsLong(), CREATE_TABLE_DATA.getRowCount()); - assertEquals(statistics.getFileCount().getAsLong(), 1L); - assertGreaterThan(statistics.getInMemoryDataSizeInBytes().getAsLong(), 0L); - assertGreaterThan(statistics.getOnDiskDataSizeInBytes().getAsLong(), 0L); - } - } - - protected void doCreateEmptyTable(SchemaTableName tableName, HiveStorageFormat storageFormat, List createTableColumns) - throws Exception - { - List partitionedBy = createTableColumns.stream() - .map(ColumnMetadata::getName) - .filter(PARTITION_COLUMN_FILTER) - .collect(toList()); - - doCreateEmptyTable(tableName, storageFormat, createTableColumns, partitionedBy); - } - - protected void doCreateEmptyTable(SchemaTableName tableName, HiveStorageFormat storageFormat, List createTableColumns, List partitionedBy) - throws Exception - { - String queryId; - try (Transaction transaction = newTransaction()) { - ConnectorSession session = newSession(); - ConnectorMetadata metadata = transaction.getMetadata(); - metadata.beginQuery(session); - queryId = session.getQueryId(); - - ConnectorTableMetadata tableMetadata = new ConnectorTableMetadata(tableName, createTableColumns, createTableProperties(storageFormat, partitionedBy)); - metadata.createTable(session, tableMetadata, false); - transaction.commit(); - } - - try (Transaction transaction = newTransaction()) { - ConnectorSession session = newSession(); - ConnectorMetadata metadata = transaction.getMetadata(); - metadata.beginQuery(session); - - // load the new table - ConnectorTableHandle tableHandle = getTableHandle(metadata, tableName); - - // verify the metadata - ConnectorTableMetadata tableMetadata = metadata.getTableMetadata(session, getTableHandle(metadata, tableName)); - - List expectedColumns = createTableColumns.stream() - .map(column -> ColumnMetadata.builder() - .setName(column.getName()) - .setType(column.getType()) - .setComment(Optional.ofNullable(column.getComment())) - .setExtraInfo(Optional.ofNullable(columnExtraInfo(partitionedBy.contains(column.getName())))) - .build()) - .collect(toList()); - assertEquals(filterNonHiddenColumnMetadata(tableMetadata.getColumns()), expectedColumns); - - // verify table format - Table table = transaction.getMetastore().getTable(tableName.getSchemaName(), tableName.getTableName()).get(); - assertEquals(table.getStorage().getStorageFormat().getInputFormat(), storageFormat.getInputFormat()); - - // verify the node version and query ID - assertEquals(table.getParameters().get(PRESTO_VERSION_NAME), TEST_SERVER_VERSION); - assertEquals(table.getParameters().get(PRESTO_QUERY_ID_NAME), queryId); - - // verify the table is empty - List columnHandles = filterNonHiddenColumnHandles(metadata.getColumnHandles(session, tableHandle).values()); - MaterializedResult result = readTable(transaction, tableHandle, columnHandles, session, TupleDomain.all(), OptionalInt.empty(), Optional.of(storageFormat)); - assertEquals(result.getRowCount(), 0); - - // verify basic statistics - if (partitionedBy.isEmpty()) { - HiveBasicStatistics statistics = getBasicStatisticsForTable(transaction, tableName); - assertEquals(statistics.getRowCount().getAsLong(), 0L); - assertEquals(statistics.getFileCount().getAsLong(), 0L); - assertEquals(statistics.getInMemoryDataSizeInBytes().getAsLong(), 0L); - assertEquals(statistics.getOnDiskDataSizeInBytes().getAsLong(), 0L); - } - } - } - - private void doInsert(HiveStorageFormat storageFormat, SchemaTableName tableName) - throws Exception - { - // creating the table - doCreateEmptyTable(tableName, storageFormat, CREATE_TABLE_COLUMNS); - - MaterializedResult.Builder resultBuilder = MaterializedResult.resultBuilder(SESSION, CREATE_TABLE_DATA.getTypes()); - for (int i = 0; i < 3; i++) { - insertData(tableName, CREATE_TABLE_DATA); - - try (Transaction transaction = newTransaction()) { - ConnectorSession session = newSession(); - ConnectorMetadata metadata = transaction.getMetadata(); - metadata.beginQuery(session); - - // load the new table - ConnectorTableHandle tableHandle = getTableHandle(metadata, tableName); - List columnHandles = filterNonHiddenColumnHandles(metadata.getColumnHandles(session, tableHandle).values()); - - // verify the metadata - ConnectorTableMetadata tableMetadata = metadata.getTableMetadata(session, getTableHandle(metadata, tableName)); - assertEquals(filterNonHiddenColumnMetadata(tableMetadata.getColumns()), CREATE_TABLE_COLUMNS); - - // verify the data - resultBuilder.rows(CREATE_TABLE_DATA.getMaterializedRows()); - MaterializedResult result = readTable(transaction, tableHandle, columnHandles, session, TupleDomain.all(), OptionalInt.empty(), Optional.empty()); - assertEqualsIgnoreOrder(result.getMaterializedRows(), resultBuilder.build().getMaterializedRows()); - - // statistics - HiveBasicStatistics tableStatistics = getBasicStatisticsForTable(transaction, tableName); - assertEquals(tableStatistics.getRowCount().orElse(0), CREATE_TABLE_DATA.getRowCount() * (i + 1L)); - assertEquals(tableStatistics.getFileCount().getAsLong(), i + 1L); - assertGreaterThan(tableStatistics.getInMemoryDataSizeInBytes().getAsLong(), 0L); - assertGreaterThan(tableStatistics.getOnDiskDataSizeInBytes().getAsLong(), 0L); - } - } - - // test rollback - Set existingFiles; - try (Transaction transaction = newTransaction()) { - existingFiles = listAllDataFiles(transaction, tableName.getSchemaName(), tableName.getTableName()); - assertFalse(existingFiles.isEmpty()); - } - - Location stagingPathRoot; - try (Transaction transaction = newTransaction()) { - ConnectorSession session = newSession(); - ConnectorMetadata metadata = transaction.getMetadata(); - - ConnectorTableHandle tableHandle = getTableHandle(metadata, tableName); - - // "stage" insert data - ConnectorInsertTableHandle insertTableHandle = metadata.beginInsert(session, tableHandle, ImmutableList.of(), NO_RETRIES); - ConnectorPageSink sink = pageSinkProvider.createPageSink(transaction.getTransactionHandle(), session, insertTableHandle, TESTING_PAGE_SINK_ID); - sink.appendPage(CREATE_TABLE_DATA.toPage()); - sink.appendPage(CREATE_TABLE_DATA.toPage()); - Collection fragments = getFutureValue(sink.finish()); - metadata.finishInsert(session, insertTableHandle, fragments, ImmutableList.of()); - - // statistics, visible from within transaction - HiveBasicStatistics tableStatistics = getBasicStatisticsForTable(transaction, tableName); - assertEquals(tableStatistics.getRowCount().getAsLong(), CREATE_TABLE_DATA.getRowCount() * 5L); - - try (Transaction otherTransaction = newTransaction()) { - // statistics, not visible from outside transaction - HiveBasicStatistics otherTableStatistics = getBasicStatisticsForTable(otherTransaction, tableName); - assertEquals(otherTableStatistics.getRowCount().getAsLong(), CREATE_TABLE_DATA.getRowCount() * 3L); - } - - // verify we did not modify the table directory - assertEquals(listAllDataFiles(transaction, tableName.getSchemaName(), tableName.getTableName()), existingFiles); - - // verify all temp files start with the unique prefix - stagingPathRoot = getStagingPathRoot(insertTableHandle); - HdfsContext context = new HdfsContext(session); - Set tempFiles = listAllDataFiles(context, stagingPathRoot); - assertTrue(!tempFiles.isEmpty()); - for (String filePath : tempFiles) { - assertThat(new Path(filePath).getName()).startsWith(session.getQueryId()); - } - - // rollback insert - transaction.rollback(); - } - - // verify temp directory is empty - HdfsContext context = new HdfsContext(newSession()); - assertTrue(listAllDataFiles(context, stagingPathRoot).isEmpty()); - - // verify the data is unchanged - try (Transaction transaction = newTransaction()) { - ConnectorSession session = newSession(); - ConnectorMetadata metadata = transaction.getMetadata(); - metadata.beginQuery(session); - - ConnectorTableHandle tableHandle = getTableHandle(metadata, tableName); - List columnHandles = filterNonHiddenColumnHandles(metadata.getColumnHandles(session, tableHandle).values()); - MaterializedResult result = readTable(transaction, tableHandle, columnHandles, session, TupleDomain.all(), OptionalInt.empty(), Optional.empty()); - assertEqualsIgnoreOrder(result.getMaterializedRows(), resultBuilder.build().getMaterializedRows()); - - // verify we did not modify the table directory - assertEquals(listAllDataFiles(transaction, tableName.getSchemaName(), tableName.getTableName()), existingFiles); - } - - // verify statistics unchanged - try (Transaction transaction = newTransaction()) { - HiveBasicStatistics statistics = getBasicStatisticsForTable(transaction, tableName); - assertEquals(statistics.getRowCount().getAsLong(), CREATE_TABLE_DATA.getRowCount() * 3L); - assertEquals(statistics.getFileCount().getAsLong(), 3L); - } - } - - private void doInsertOverwriteUnpartitioned(SchemaTableName tableName) - throws Exception - { - // create table with data - doCreateEmptyTable(tableName, ORC, CREATE_TABLE_COLUMNS); - insertData(tableName, CREATE_TABLE_DATA); - - // overwrite table with new data - MaterializedResult.Builder overwriteDataBuilder = MaterializedResult.resultBuilder(SESSION, CREATE_TABLE_DATA.getTypes()); - MaterializedResult overwriteData = null; - - Map overwriteProperties = ImmutableMap.of("insert_existing_partitions_behavior", "OVERWRITE"); - - for (int i = 0; i < 3; i++) { - overwriteDataBuilder.rows(reverse(CREATE_TABLE_DATA.getMaterializedRows())); - overwriteData = overwriteDataBuilder.build(); - - insertData(tableName, overwriteData, overwriteProperties); - - // verify overwrite - try (Transaction transaction = newTransaction()) { - ConnectorSession session = newSession(); - ConnectorMetadata metadata = transaction.getMetadata(); - metadata.beginQuery(session); - - // load the new table - ConnectorTableHandle tableHandle = getTableHandle(metadata, tableName); - List columnHandles = filterNonHiddenColumnHandles(metadata.getColumnHandles(session, tableHandle).values()); - - // verify the metadata - ConnectorTableMetadata tableMetadata = metadata.getTableMetadata(session, getTableHandle(metadata, tableName)); - assertEquals(filterNonHiddenColumnMetadata(tableMetadata.getColumns()), CREATE_TABLE_COLUMNS); - - // verify the data - MaterializedResult result = readTable(transaction, tableHandle, columnHandles, session, TupleDomain.all(), OptionalInt.empty(), Optional.empty()); - assertEqualsIgnoreOrder(result.getMaterializedRows(), overwriteData.getMaterializedRows()); - - // statistics - HiveBasicStatistics tableStatistics = getBasicStatisticsForTable(transaction, tableName); - assertEquals(tableStatistics.getRowCount().getAsLong(), overwriteData.getRowCount()); - assertEquals(tableStatistics.getFileCount().getAsLong(), 1L); - assertGreaterThan(tableStatistics.getInMemoryDataSizeInBytes().getAsLong(), 0L); - assertGreaterThan(tableStatistics.getOnDiskDataSizeInBytes().getAsLong(), 0L); - } - } - - // test rollback - Set existingFiles; - try (Transaction transaction = newTransaction()) { - existingFiles = listAllDataFiles(transaction, tableName.getSchemaName(), tableName.getTableName()); - assertFalse(existingFiles.isEmpty()); - } - - Location stagingPathRoot; - try (Transaction transaction = newTransaction()) { - ConnectorSession session = newSession(overwriteProperties); - ConnectorMetadata metadata = transaction.getMetadata(); - - ConnectorTableHandle tableHandle = getTableHandle(metadata, tableName); - - // "stage" insert data - ConnectorInsertTableHandle insertTableHandle = metadata.beginInsert(session, tableHandle, ImmutableList.of(), NO_RETRIES); - ConnectorPageSink sink = pageSinkProvider.createPageSink(transaction.getTransactionHandle(), session, insertTableHandle, TESTING_PAGE_SINK_ID); - for (int i = 0; i < 4; i++) { - sink.appendPage(overwriteData.toPage()); - } - Collection fragments = getFutureValue(sink.finish()); - metadata.finishInsert(session, insertTableHandle, fragments, ImmutableList.of()); - - // statistics, visible from within transaction - HiveBasicStatistics tableStatistics = getBasicStatisticsForTable(transaction, tableName); - assertEquals(tableStatistics.getRowCount().getAsLong(), overwriteData.getRowCount() * 4L); - - try (Transaction otherTransaction = newTransaction()) { - // statistics, not visible from outside transaction - HiveBasicStatistics otherTableStatistics = getBasicStatisticsForTable(otherTransaction, tableName); - assertEquals(otherTableStatistics.getRowCount().getAsLong(), overwriteData.getRowCount()); - } - - // verify we did not modify the table directory - assertEquals(listAllDataFiles(transaction, tableName.getSchemaName(), tableName.getTableName()), existingFiles); - - // verify all temp files start with the unique prefix - stagingPathRoot = getStagingPathRoot(insertTableHandle); - HdfsContext context = new HdfsContext(session); - Set tempFiles = listAllDataFiles(context, stagingPathRoot); - assertTrue(!tempFiles.isEmpty()); - for (String filePath : tempFiles) { - assertThat(new Path(filePath).getName()).startsWith(session.getQueryId()); - } - - // rollback insert - transaction.rollback(); - } - - // verify temp directory is empty - HdfsContext context = new HdfsContext(newSession()); - assertTrue(listAllDataFiles(context, stagingPathRoot).isEmpty()); - - // verify the data is unchanged - try (Transaction transaction = newTransaction()) { - ConnectorSession session = newSession(); - ConnectorMetadata metadata = transaction.getMetadata(); - metadata.beginQuery(session); - - ConnectorTableHandle tableHandle = getTableHandle(metadata, tableName); - List columnHandles = filterNonHiddenColumnHandles(metadata.getColumnHandles(session, tableHandle).values()); - MaterializedResult result = readTable(transaction, tableHandle, columnHandles, session, TupleDomain.all(), OptionalInt.empty(), Optional.empty()); - assertEqualsIgnoreOrder(result.getMaterializedRows(), overwriteData.getMaterializedRows()); - - // verify we did not modify the table directory - assertEquals(listAllDataFiles(transaction, tableName.getSchemaName(), tableName.getTableName()), existingFiles); - } - - // verify statistics unchanged - try (Transaction transaction = newTransaction()) { - HiveBasicStatistics statistics = getBasicStatisticsForTable(transaction, tableName); - assertEquals(statistics.getRowCount().getAsLong(), overwriteData.getRowCount()); - assertEquals(statistics.getFileCount().getAsLong(), 1L); - } - } - - private Location getStagingPathRoot(ConnectorInsertTableHandle insertTableHandle) - { - HiveInsertTableHandle handle = (HiveInsertTableHandle) insertTableHandle; - WriteInfo writeInfo = getLocationService().getQueryWriteInfo(handle.getLocationHandle()); - if (writeInfo.writeMode() != STAGE_AND_MOVE_TO_TARGET_DIRECTORY) { - throw new AssertionError("writeMode is not STAGE_AND_MOVE_TO_TARGET_DIRECTORY"); - } - return writeInfo.writePath(); - } - - private Location getStagingPathRoot(ConnectorOutputTableHandle outputTableHandle) - { - HiveOutputTableHandle handle = (HiveOutputTableHandle) outputTableHandle; - return getLocationService() - .getQueryWriteInfo(handle.getLocationHandle()) - .writePath(); - } - - private Location getTargetPathRoot(ConnectorInsertTableHandle insertTableHandle) - { - HiveInsertTableHandle hiveInsertTableHandle = (HiveInsertTableHandle) insertTableHandle; - - return getLocationService() - .getQueryWriteInfo(hiveInsertTableHandle.getLocationHandle()) - .targetPath(); - } - - protected Set listAllDataFiles(Transaction transaction, String schemaName, String tableName) - throws IOException - { - HdfsContext hdfsContext = new HdfsContext(newSession()); - Set existingFiles = new HashSet<>(); - for (String location : listAllDataPaths(transaction.getMetastore(), schemaName, tableName)) { - existingFiles.addAll(listAllDataFiles(hdfsContext, Location.of(location))); - } - return existingFiles; - } - - public static List listAllDataPaths(SemiTransactionalHiveMetastore metastore, String schemaName, String tableName) - { - ImmutableList.Builder locations = ImmutableList.builder(); - Table table = metastore.getTable(schemaName, tableName).get(); - if (table.getStorage().getLocation() != null) { - // For partitioned table, there should be nothing directly under this directory. - // But including this location in the set makes the directory content assert more - // extensive, which is desirable. - locations.add(table.getStorage().getLocation()); - } - - Optional> partitionNames = metastore.getPartitionNames(schemaName, tableName); - if (partitionNames.isPresent()) { - metastore.getPartitionsByNames(schemaName, tableName, partitionNames.get()).values().stream() - .map(Optional::get) - .map(partition -> partition.getStorage().getLocation()) - .filter(location -> !location.startsWith(table.getStorage().getLocation())) - .forEach(locations::add); - } - - return locations.build(); - } - - protected Set listAllDataFiles(HdfsContext context, Location location) - throws IOException - { - Path path = new Path(location.toString()); - Set result = new HashSet<>(); - FileSystem fileSystem = hdfsEnvironment.getFileSystem(context, path); - if (fileSystem.exists(path)) { - for (FileStatus fileStatus : fileSystem.listStatus(path)) { - if (fileStatus.getPath().getName().startsWith(".trino")) { - // skip hidden files - } - else if (fileStatus.isFile()) { - result.add(fileStatus.getPath().toString()); - } - else if (fileStatus.isDirectory()) { - result.addAll(listAllDataFiles(context, Location.of(fileStatus.getPath().toString()))); - } - } - } - return result; - } - - private void doInsertIntoNewPartition(HiveStorageFormat storageFormat, SchemaTableName tableName) - throws Exception - { - // creating the table - doCreateEmptyTable(tableName, storageFormat, CREATE_TABLE_COLUMNS_PARTITIONED); - - // insert the data - String queryId = insertData(tableName, CREATE_TABLE_PARTITIONED_DATA); - - Set existingFiles; - try (Transaction transaction = newTransaction()) { - // verify partitions were created - Table table = metastoreClient.getTable(tableName.getSchemaName(), tableName.getTableName()) - .orElseThrow(() -> new TableNotFoundException(tableName)); - List partitionNames = transaction.getMetastore().getPartitionNames(tableName.getSchemaName(), tableName.getTableName()) - .orElseThrow(() -> new AssertionError("Table does not exist: " + tableName)); - assertEqualsIgnoreOrder(partitionNames, CREATE_TABLE_PARTITIONED_DATA.getMaterializedRows().stream() - .map(row -> "ds=" + row.getField(CREATE_TABLE_PARTITIONED_DATA.getTypes().size() - 1)) - .collect(toImmutableList())); - - // verify the node versions in partitions - Map> partitions = getMetastoreClient().getPartitionsByNames(table, partitionNames); - assertEquals(partitions.size(), partitionNames.size()); - for (String partitionName : partitionNames) { - Partition partition = partitions.get(partitionName).get(); - assertEquals(partition.getParameters().get(PRESTO_VERSION_NAME), TEST_SERVER_VERSION); - assertEquals(partition.getParameters().get(PRESTO_QUERY_ID_NAME), queryId); - } - - // load the new table - ConnectorSession session = newSession(); - ConnectorMetadata metadata = transaction.getMetadata(); - metadata.beginQuery(session); - ConnectorTableHandle tableHandle = getTableHandle(metadata, tableName); - List columnHandles = filterNonHiddenColumnHandles(metadata.getColumnHandles(session, tableHandle).values()); - - // verify the data - MaterializedResult result = readTable(transaction, tableHandle, columnHandles, session, TupleDomain.all(), OptionalInt.empty(), Optional.of(storageFormat)); - assertEqualsIgnoreOrder(result.getMaterializedRows(), CREATE_TABLE_PARTITIONED_DATA.getMaterializedRows()); - - // test rollback - existingFiles = listAllDataFiles(transaction, tableName.getSchemaName(), tableName.getTableName()); - assertFalse(existingFiles.isEmpty()); - - // test statistics - for (String partitionName : partitionNames) { - HiveBasicStatistics partitionStatistics = getBasicStatisticsForPartition(transaction, tableName, COLUMN_NAMES_PARTITIONED, partitionName); - assertEquals(partitionStatistics.getRowCount().getAsLong(), 1L); - assertEquals(partitionStatistics.getFileCount().getAsLong(), 1L); - assertGreaterThan(partitionStatistics.getInMemoryDataSizeInBytes().getAsLong(), 0L); - assertGreaterThan(partitionStatistics.getOnDiskDataSizeInBytes().getAsLong(), 0L); - } - } - - Location stagingPathRoot; - try (Transaction transaction = newTransaction()) { - ConnectorSession session = newSession(); - ConnectorMetadata metadata = transaction.getMetadata(); - ConnectorTableHandle tableHandle = getTableHandle(metadata, tableName); - - // "stage" insert data - ConnectorInsertTableHandle insertTableHandle = metadata.beginInsert(session, tableHandle, ImmutableList.of(), NO_RETRIES); - stagingPathRoot = getStagingPathRoot(insertTableHandle); - ConnectorPageSink sink = pageSinkProvider.createPageSink(transaction.getTransactionHandle(), session, insertTableHandle, TESTING_PAGE_SINK_ID); - sink.appendPage(CREATE_TABLE_PARTITIONED_DATA_2ND.toPage()); - Collection fragments = getFutureValue(sink.finish()); - metadata.finishInsert(session, insertTableHandle, fragments, ImmutableList.of()); - - // verify all temp files start with the unique prefix - HdfsContext context = new HdfsContext(session); - Set tempFiles = listAllDataFiles(context, getStagingPathRoot(insertTableHandle)); - assertTrue(!tempFiles.isEmpty()); - for (String filePath : tempFiles) { - assertThat(new Path(filePath).getName()).startsWith(session.getQueryId()); - } - - // rollback insert - transaction.rollback(); - } - - // verify the data is unchanged - try (Transaction transaction = newTransaction()) { - ConnectorSession session = newSession(); - ConnectorMetadata metadata = transaction.getMetadata(); - metadata.beginQuery(session); - ConnectorTableHandle tableHandle = getTableHandle(metadata, tableName); - List columnHandles = filterNonHiddenColumnHandles(metadata.getColumnHandles(session, tableHandle).values()); - - MaterializedResult result = readTable(transaction, tableHandle, columnHandles, session, TupleDomain.all(), OptionalInt.empty(), Optional.empty()); - assertEqualsIgnoreOrder(result.getMaterializedRows(), CREATE_TABLE_PARTITIONED_DATA.getMaterializedRows()); - - // verify we did not modify the table directory - assertEquals(listAllDataFiles(transaction, tableName.getSchemaName(), tableName.getTableName()), existingFiles); - - // verify temp directory is empty - HdfsContext context = new HdfsContext(session); - assertTrue(listAllDataFiles(context, stagingPathRoot).isEmpty()); - } - } - - private void doInsertUnsupportedWriteType(HiveStorageFormat storageFormat, SchemaTableName tableName) - throws Exception - { - List columns = ImmutableList.of(new Column("dummy", HiveType.valueOf("uniontype"), Optional.empty())); - List partitionColumns = ImmutableList.of(new Column("name", HIVE_STRING, Optional.empty())); - - createEmptyTable(tableName, storageFormat, columns, partitionColumns); - - try (Transaction transaction = newTransaction()) { - ConnectorMetadata metadata = transaction.getMetadata(); - ConnectorSession session = newSession(); - ConnectorTableHandle tableHandle = getTableHandle(metadata, tableName); - - metadata.beginInsert(session, tableHandle, ImmutableList.of(), NO_RETRIES); - fail("expected failure"); - } - catch (TrinoException e) { - assertThat(e).hasMessageMatching("Inserting into Hive table .* with column type uniontype not supported"); - } - } - - private void doInsertIntoExistingPartition(HiveStorageFormat storageFormat, SchemaTableName tableName) - throws Exception - { - // creating the table - doCreateEmptyTable(tableName, storageFormat, CREATE_TABLE_COLUMNS_PARTITIONED); - - MaterializedResult.Builder resultBuilder = MaterializedResult.resultBuilder(SESSION, CREATE_TABLE_PARTITIONED_DATA.getTypes()); - for (int i = 0; i < 3; i++) { - // insert the data - insertData(tableName, CREATE_TABLE_PARTITIONED_DATA); - - try (Transaction transaction = newTransaction()) { - ConnectorSession session = newSession(); - ConnectorMetadata metadata = transaction.getMetadata(); - metadata.beginQuery(session); - ConnectorTableHandle tableHandle = getTableHandle(metadata, tableName); - - // verify partitions were created - List partitionNames = transaction.getMetastore().getPartitionNames(tableName.getSchemaName(), tableName.getTableName()) - .orElseThrow(() -> new AssertionError("Table does not exist: " + tableName)); - assertEqualsIgnoreOrder(partitionNames, CREATE_TABLE_PARTITIONED_DATA.getMaterializedRows().stream() - .map(row -> "ds=" + row.getField(CREATE_TABLE_PARTITIONED_DATA.getTypes().size() - 1)) - .collect(toImmutableList())); - - // load the new table - List columnHandles = filterNonHiddenColumnHandles(metadata.getColumnHandles(session, tableHandle).values()); - - // verify the data - resultBuilder.rows(CREATE_TABLE_PARTITIONED_DATA.getMaterializedRows()); - MaterializedResult result = readTable(transaction, tableHandle, columnHandles, session, TupleDomain.all(), OptionalInt.empty(), Optional.of(storageFormat)); - assertEqualsIgnoreOrder(result.getMaterializedRows(), resultBuilder.build().getMaterializedRows()); - - // test statistics - for (String partitionName : partitionNames) { - HiveBasicStatistics statistics = getBasicStatisticsForPartition(transaction, tableName, COLUMN_NAMES_PARTITIONED, partitionName); - assertEquals(statistics.getRowCount().getAsLong(), i + 1L); - assertEquals(statistics.getFileCount().getAsLong(), i + 1L); - assertGreaterThan(statistics.getInMemoryDataSizeInBytes().getAsLong(), 0L); - assertGreaterThan(statistics.getOnDiskDataSizeInBytes().getAsLong(), 0L); - } - } - } - - // test rollback - Set existingFiles; - Location stagingPathRoot; - try (Transaction transaction = newTransaction()) { - ConnectorMetadata metadata = transaction.getMetadata(); - ConnectorSession session = newSession(); - - existingFiles = listAllDataFiles(transaction, tableName.getSchemaName(), tableName.getTableName()); - assertFalse(existingFiles.isEmpty()); - - ConnectorTableHandle tableHandle = getTableHandle(metadata, tableName); - - // "stage" insert data - ConnectorInsertTableHandle insertTableHandle = metadata.beginInsert(session, tableHandle, ImmutableList.of(), NO_RETRIES); - stagingPathRoot = getStagingPathRoot(insertTableHandle); - ConnectorPageSink sink = pageSinkProvider.createPageSink(transaction.getTransactionHandle(), session, insertTableHandle, TESTING_PAGE_SINK_ID); - sink.appendPage(CREATE_TABLE_PARTITIONED_DATA.toPage()); - sink.appendPage(CREATE_TABLE_PARTITIONED_DATA.toPage()); - Collection fragments = getFutureValue(sink.finish()); - metadata.finishInsert(session, insertTableHandle, fragments, ImmutableList.of()); - - // verify all temp files start with the unique prefix - HdfsContext context = new HdfsContext(session); - Set tempFiles = listAllDataFiles(context, getStagingPathRoot(insertTableHandle)); - assertTrue(!tempFiles.isEmpty()); - for (String filePath : tempFiles) { - assertThat(new Path(filePath).getName()).startsWith(session.getQueryId()); - } - - // verify statistics are visible from within of the current transaction - List partitionNames = transaction.getMetastore().getPartitionNames(tableName.getSchemaName(), tableName.getTableName()) - .orElseThrow(() -> new AssertionError("Table does not exist: " + tableName)); - for (String partitionName : partitionNames) { - HiveBasicStatistics partitionStatistics = getBasicStatisticsForPartition(transaction, tableName, COLUMN_NAMES_PARTITIONED, partitionName); - assertEquals(partitionStatistics.getRowCount().getAsLong(), 5L); - } - - // rollback insert - transaction.rollback(); - } - - try (Transaction transaction = newTransaction()) { - ConnectorMetadata metadata = transaction.getMetadata(); - ConnectorSession session = newSession(); - metadata.beginQuery(session); - ConnectorTableHandle tableHandle = getTableHandle(metadata, tableName); - List columnHandles = filterNonHiddenColumnHandles(metadata.getColumnHandles(session, tableHandle).values()); - - // verify the data is unchanged - MaterializedResult result = readTable(transaction, tableHandle, columnHandles, session, TupleDomain.all(), OptionalInt.empty(), Optional.empty()); - assertEqualsIgnoreOrder(result.getMaterializedRows(), resultBuilder.build().getMaterializedRows()); - - // verify we did not modify the table directory - assertEquals(listAllDataFiles(transaction, tableName.getSchemaName(), tableName.getTableName()), existingFiles); - - // verify temp directory is empty - HdfsContext hdfsContext = new HdfsContext(session); - assertTrue(listAllDataFiles(hdfsContext, stagingPathRoot).isEmpty()); - - // verify statistics have been rolled back - List partitionNames = transaction.getMetastore().getPartitionNames(tableName.getSchemaName(), tableName.getTableName()) - .orElseThrow(() -> new AssertionError("Table does not exist: " + tableName)); - for (String partitionName : partitionNames) { - HiveBasicStatistics partitionStatistics = getBasicStatisticsForPartition(transaction, tableName, COLUMN_NAMES_PARTITIONED, partitionName); - assertEquals(partitionStatistics.getRowCount().getAsLong(), 3L); - } - } - } - - private void doInsertIntoExistingPartitionEmptyStatistics(HiveStorageFormat storageFormat, SchemaTableName tableName) - throws Exception - { - doCreateEmptyTable(tableName, storageFormat, CREATE_TABLE_COLUMNS_PARTITIONED); - insertData(tableName, CREATE_TABLE_PARTITIONED_DATA); - - eraseStatistics(tableName); - - insertData(tableName, CREATE_TABLE_PARTITIONED_DATA); - - try (Transaction transaction = newTransaction()) { - List partitionNames = transaction.getMetastore().getPartitionNames(tableName.getSchemaName(), tableName.getTableName()) - .orElseThrow(() -> new AssertionError("Table does not exist: " + tableName)); - - for (String partitionName : partitionNames) { - HiveBasicStatistics statistics = getBasicStatisticsForPartition(transaction, tableName, COLUMN_NAMES_PARTITIONED, partitionName); - assertThat(statistics.getRowCount()).isNotPresent(); - assertThat(statistics.getInMemoryDataSizeInBytes()).isNotPresent(); - // fileCount and rawSize statistics are computed on the fly by the metastore, thus cannot be erased - } - } - } - - private static HiveBasicStatistics getBasicStatisticsForTable(Transaction transaction, SchemaTableName table) - { - return transaction - .getMetastore() - .getTableStatistics(table.getSchemaName(), table.getTableName(), Optional.empty()) - .getBasicStatistics(); - } - - private static HiveBasicStatistics getBasicStatisticsForPartition(Transaction transaction, SchemaTableName table, Set columns, String partitionName) - { - return transaction - .getMetastore() - .getPartitionStatistics(table.getSchemaName(), table.getTableName(), columns, ImmutableSet.of(partitionName)) - .get(partitionName) - .getBasicStatistics(); - } - - private void eraseStatistics(SchemaTableName schemaTableName) - { - HiveMetastore metastoreClient = getMetastoreClient(); - metastoreClient.updateTableStatistics(schemaTableName.getSchemaName(), schemaTableName.getTableName(), NO_ACID_TRANSACTION, statistics -> new PartitionStatistics(createEmptyStatistics(), ImmutableMap.of())); - Table table = metastoreClient.getTable(schemaTableName.getSchemaName(), schemaTableName.getTableName()) - .orElseThrow(() -> new TableNotFoundException(schemaTableName)); - List partitionColumns = table.getPartitionColumns().stream() - .map(Column::getName) - .collect(toImmutableList()); - if (!table.getPartitionColumns().isEmpty()) { - List partitionNames = metastoreClient.getPartitionNamesByFilter(schemaTableName.getSchemaName(), schemaTableName.getTableName(), partitionColumns, TupleDomain.all()) - .orElse(ImmutableList.of()); - List partitions = metastoreClient - .getPartitionsByNames(table, partitionNames) - .values() - .stream() - .filter(Optional::isPresent) - .map(Optional::get) - .collect(toImmutableList()); - for (Partition partition : partitions) { - metastoreClient.updatePartitionStatistics( - table, - makePartName(partitionColumns, partition.getValues()), - statistics -> new PartitionStatistics(createEmptyStatistics(), ImmutableMap.of())); - } - } - } - - /** - * @return query id - */ - private String insertData(SchemaTableName tableName, MaterializedResult data) - throws Exception - { - return insertData(tableName, data, ImmutableMap.of()); - } - - private String insertData(SchemaTableName tableName, MaterializedResult data, Map sessionProperties) - throws Exception - { - Location writePath; - Location targetPath; - String queryId; - try (Transaction transaction = newTransaction()) { - ConnectorMetadata metadata = transaction.getMetadata(); - ConnectorSession session = newSession(sessionProperties); - ConnectorTableHandle tableHandle = getTableHandle(metadata, tableName); - ConnectorInsertTableHandle insertTableHandle = metadata.beginInsert(session, tableHandle, ImmutableList.of(), NO_RETRIES); - queryId = session.getQueryId(); - writePath = getStagingPathRoot(insertTableHandle); - targetPath = getTargetPathRoot(insertTableHandle); - - ConnectorPageSink sink = pageSinkProvider.createPageSink(transaction.getTransactionHandle(), session, insertTableHandle, TESTING_PAGE_SINK_ID); - - // write data - sink.appendPage(data.toPage()); - Collection fragments = getFutureValue(sink.finish()); - - // commit the insert - metadata.finishInsert(session, insertTableHandle, fragments, ImmutableList.of()); - transaction.commit(); - } - - // check that temporary files are removed - if (!writePath.equals(targetPath)) { - HdfsContext context = new HdfsContext(newSession()); - FileSystem fileSystem = hdfsEnvironment.getFileSystem(context, new Path(writePath.toString())); - assertFalse(fileSystem.exists(new Path(writePath.toString()))); - } - - return queryId; - } - - private void doTestMetadataDelete(HiveStorageFormat storageFormat, SchemaTableName tableName) - throws Exception - { - // creating the table - doCreateEmptyTable(tableName, storageFormat, CREATE_TABLE_COLUMNS_PARTITIONED); - - insertData(tableName, CREATE_TABLE_PARTITIONED_DATA); - - MaterializedResult.Builder expectedResultBuilder = MaterializedResult.resultBuilder(SESSION, CREATE_TABLE_PARTITIONED_DATA.getTypes()); - expectedResultBuilder.rows(CREATE_TABLE_PARTITIONED_DATA.getMaterializedRows()); - - try (Transaction transaction = newTransaction()) { - ConnectorSession session = newSession(); - ConnectorMetadata metadata = transaction.getMetadata(); - metadata.beginQuery(session); - - // verify partitions were created - List partitionNames = transaction.getMetastore().getPartitionNames(tableName.getSchemaName(), tableName.getTableName()) - .orElseThrow(() -> new AssertionError("Table does not exist: " + tableName)); - assertEqualsIgnoreOrder(partitionNames, CREATE_TABLE_PARTITIONED_DATA.getMaterializedRows().stream() - .map(row -> "ds=" + row.getField(CREATE_TABLE_PARTITIONED_DATA.getTypes().size() - 1)) - .collect(toImmutableList())); - - // verify table directory is not empty - Set filesAfterInsert = listAllDataFiles(transaction, tableName.getSchemaName(), tableName.getTableName()); - assertFalse(filesAfterInsert.isEmpty()); - - // verify the data - ConnectorTableHandle tableHandle = getTableHandle(metadata, tableName); - List columnHandles = filterNonHiddenColumnHandles(metadata.getColumnHandles(session, tableHandle).values()); - MaterializedResult result = readTable(transaction, tableHandle, columnHandles, session, TupleDomain.all(), OptionalInt.empty(), Optional.of(storageFormat)); - assertEqualsIgnoreOrder(result.getMaterializedRows(), expectedResultBuilder.build().getMaterializedRows()); - } - - try (Transaction transaction = newTransaction()) { - ConnectorSession session = newSession(); - ConnectorMetadata metadata = transaction.getMetadata(); - - // get ds column handle - ConnectorTableHandle tableHandle = getTableHandle(metadata, tableName); - HiveColumnHandle dsColumnHandle = (HiveColumnHandle) metadata.getColumnHandles(session, tableHandle).get("ds"); - - // delete ds=2015-07-03 - session = newSession(); - TupleDomain tupleDomain = TupleDomain.fromFixedValues(ImmutableMap.of(dsColumnHandle, NullableValue.of(createUnboundedVarcharType(), utf8Slice("2015-07-03")))); - Constraint constraint = new Constraint(tupleDomain, tupleDomain.asPredicate(), tupleDomain.getDomains().orElseThrow().keySet()); - tableHandle = applyFilter(metadata, tableHandle, constraint); - tableHandle = metadata.applyDelete(session, tableHandle).get(); - metadata.executeDelete(session, tableHandle); - - transaction.commit(); - } - - try (Transaction transaction = newTransaction()) { - ConnectorSession session = newSession(); - ConnectorMetadata metadata = transaction.getMetadata(); - metadata.beginQuery(session); - ConnectorTableHandle tableHandle = getTableHandle(metadata, tableName); - List columnHandles = filterNonHiddenColumnHandles(metadata.getColumnHandles(session, tableHandle).values()); - HiveColumnHandle dsColumnHandle = (HiveColumnHandle) metadata.getColumnHandles(session, tableHandle).get("ds"); - int dsColumnOrdinalPosition = columnHandles.indexOf(dsColumnHandle); - - // verify the data - ImmutableList expectedRows = expectedResultBuilder.build().getMaterializedRows().stream() - .filter(row -> !"2015-07-03".equals(row.getField(dsColumnOrdinalPosition))) - .collect(toImmutableList()); - MaterializedResult actualAfterDelete = readTable(transaction, tableHandle, columnHandles, session, TupleDomain.all(), OptionalInt.empty(), Optional.of(storageFormat)); - assertEqualsIgnoreOrder(actualAfterDelete.getMaterializedRows(), expectedRows); - } - - try (Transaction transaction = newTransaction()) { - ConnectorSession session = newSession(); - ConnectorMetadata metadata = transaction.getMetadata(); - ConnectorTableHandle tableHandle = getTableHandle(metadata, tableName); - HiveColumnHandle dsColumnHandle = (HiveColumnHandle) metadata.getColumnHandles(session, tableHandle).get("ds"); - - // delete ds=2015-07-01 and 2015-07-02 - session = newSession(); - TupleDomain tupleDomain2 = TupleDomain.withColumnDomains( - ImmutableMap.of(dsColumnHandle, Domain.create(ValueSet.ofRanges(Range.range(createUnboundedVarcharType(), utf8Slice("2015-07-01"), true, utf8Slice("2015-07-02"), true)), false))); - Constraint constraint2 = new Constraint(tupleDomain2, tupleDomain2.asPredicate(), tupleDomain2.getDomains().orElseThrow().keySet()); - tableHandle = applyFilter(metadata, tableHandle, constraint2); - tableHandle = metadata.applyDelete(session, tableHandle).get(); - metadata.executeDelete(session, tableHandle); - - transaction.commit(); - } - - try (Transaction transaction = newTransaction()) { - ConnectorSession session = newSession(); - ConnectorMetadata metadata = transaction.getMetadata(); - ConnectorTableHandle tableHandle = getTableHandle(metadata, tableName); - List columnHandles = ImmutableList.copyOf(metadata.getColumnHandles(session, tableHandle).values()); - - // verify the data - session = newSession(); - MaterializedResult actualAfterDelete2 = readTable(transaction, tableHandle, columnHandles, session, TupleDomain.all(), OptionalInt.empty(), Optional.of(storageFormat)); - assertEqualsIgnoreOrder(actualAfterDelete2.getMaterializedRows(), ImmutableList.of()); - - // verify table directory is empty - Set filesAfterDelete = listAllDataFiles(transaction, tableName.getSchemaName(), tableName.getTableName()); - assertTrue(filesAfterDelete.isEmpty()); - } - } - - protected void assertGetRecords(String tableName, HiveStorageFormat hiveStorageFormat) - throws Exception - { - try (Transaction transaction = newTransaction()) { - ConnectorSession session = newSession(); - ConnectorMetadata metadata = transaction.getMetadata(); - metadata.beginQuery(session); - - ConnectorTableHandle tableHandle = getTableHandle(metadata, new SchemaTableName(database, tableName)); - ConnectorTableMetadata tableMetadata = metadata.getTableMetadata(session, tableHandle); - HiveSplit hiveSplit = getHiveSplit(tableHandle, transaction, session); - - List columnHandles = ImmutableList.copyOf(metadata.getColumnHandles(session, tableHandle).values()); - - ConnectorPageSource pageSource = pageSourceProvider.createPageSource(transaction.getTransactionHandle(), session, hiveSplit, tableHandle, columnHandles, DynamicFilter.EMPTY); - assertGetRecords(hiveStorageFormat, tableMetadata, hiveSplit, pageSource, columnHandles); - } - } - - protected HiveSplit getHiveSplit(ConnectorTableHandle tableHandle, Transaction transaction, ConnectorSession session) - { - List splits = getAllSplits(tableHandle, transaction, session); - assertEquals(splits.size(), 1); - return (HiveSplit) getOnlyElement(splits); - } - - protected void assertGetRecords( - HiveStorageFormat hiveStorageFormat, - ConnectorTableMetadata tableMetadata, - HiveSplit hiveSplit, - ConnectorPageSource pageSource, - List columnHandles) - throws IOException - { - try { - MaterializedResult result = materializeSourceDataStream(newSession(), pageSource, getTypes(columnHandles)); - - assertPageSourceType(pageSource, hiveStorageFormat); - - ImmutableMap columnIndex = indexColumns(tableMetadata); - - long rowNumber = 0; - long completedBytes = 0; - for (MaterializedRow row : result) { - try { - assertValueTypes(row, tableMetadata.getColumns()); - } - catch (RuntimeException e) { - throw new RuntimeException("row " + rowNumber, e); - } - - rowNumber++; - Integer index; - Object value; - - // STRING - index = columnIndex.get("t_string"); - value = row.getField(index); - if (rowNumber % 19 == 0) { - assertNull(value); - } - else if (rowNumber % 19 == 1) { - assertEquals(value, ""); - } - else { - assertEquals(value, "test"); - } - - // NUMBERS - assertEquals(row.getField(columnIndex.get("t_tinyint")), (byte) (1 + rowNumber)); - assertEquals(row.getField(columnIndex.get("t_smallint")), (short) (2 + rowNumber)); - assertEquals(row.getField(columnIndex.get("t_int")), (int) (3 + rowNumber)); - - index = columnIndex.get("t_bigint"); - if ((rowNumber % 13) == 0) { - assertNull(row.getField(index)); - } - else { - assertEquals(row.getField(index), 4 + rowNumber); - } - - assertEquals((Float) row.getField(columnIndex.get("t_float")), 5.1f + rowNumber, 0.001); - assertEquals(row.getField(columnIndex.get("t_double")), 6.2 + rowNumber); - - // BOOLEAN - index = columnIndex.get("t_boolean"); - if ((rowNumber % 3) == 2) { - assertNull(row.getField(index)); - } - else { - assertEquals(row.getField(index), (rowNumber % 3) != 0); - } - - // TIMESTAMP - index = columnIndex.get("t_timestamp"); - if (index != null) { - if ((rowNumber % 17) == 0) { - assertNull(row.getField(index)); - } - else { - SqlTimestamp expected = sqlTimestampOf(3, 2011, 5, 6, 7, 8, 9, 123); - assertEquals(row.getField(index), expected); - } - } - - // BINARY - index = columnIndex.get("t_binary"); - if (index != null) { - if ((rowNumber % 23) == 0) { - assertNull(row.getField(index)); - } - else { - assertEquals(row.getField(index), new SqlVarbinary("test binary".getBytes(UTF_8))); - } - } - - // DATE - index = columnIndex.get("t_date"); - if (index != null) { - if ((rowNumber % 37) == 0) { - assertNull(row.getField(index)); - } - else { - SqlDate expected = new SqlDate(toIntExact(MILLISECONDS.toDays(new DateTime(2013, 8, 9, 0, 0, 0, UTC).getMillis()))); - assertEquals(row.getField(index), expected); - } - } - - // VARCHAR(50) - index = columnIndex.get("t_varchar"); - if (index != null) { - value = row.getField(index); - if (rowNumber % 39 == 0) { - assertNull(value); - } - else if (rowNumber % 39 == 1) { - // https://issues.apache.org/jira/browse/HIVE-13289 - // RCBINARY reads empty VARCHAR as null - if (hiveStorageFormat == RCBINARY) { - assertNull(value); - } - else { - assertEquals(value, ""); - } - } - else { - assertEquals(value, "test varchar"); - } - } - - //CHAR(25) - index = columnIndex.get("t_char"); - if (index != null) { - value = row.getField(index); - if ((rowNumber % 41) == 0) { - assertNull(value); - } - else { - assertEquals(value, (rowNumber % 41) == 1 ? " " : "test char "); - } - } - - // MAP - index = columnIndex.get("t_map"); - if (index != null) { - if ((rowNumber % 27) == 0) { - assertNull(row.getField(index)); - } - else { - assertEquals(row.getField(index), ImmutableMap.of("test key", "test value")); - } - } - - // ARRAY - index = columnIndex.get("t_array_string"); - if (index != null) { - if ((rowNumber % 29) == 0) { - assertNull(row.getField(index)); - } - else { - assertEquals(row.getField(index), ImmutableList.of("abc", "xyz", "data")); - } - } - - // ARRAY - index = columnIndex.get("t_array_timestamp"); - if (index != null) { - if ((rowNumber % 43) == 0) { - assertNull(row.getField(index)); - } - else { - SqlTimestamp expected = sqlTimestampOf(3, LocalDateTime.of(2011, 5, 6, 7, 8, 9, 123_000_000)); - assertEquals(row.getField(index), ImmutableList.of(expected)); - } - } - - // ARRAY> - index = columnIndex.get("t_array_struct"); - if (index != null) { - if ((rowNumber % 31) == 0) { - assertNull(row.getField(index)); - } - else { - List expected1 = ImmutableList.of("test abc", 0.1); - List expected2 = ImmutableList.of("test xyz", 0.2); - assertEquals(row.getField(index), ImmutableList.of(expected1, expected2)); - } - } - - // STRUCT - index = columnIndex.get("t_struct"); - if (index != null) { - if ((rowNumber % 31) == 0) { - assertNull(row.getField(index)); - } - else { - assertTrue(row.getField(index) instanceof List); - List values = (List) row.getField(index); - assertEquals(values.size(), 2); - assertEquals(values.get(0), "test abc"); - assertEquals(values.get(1), 0.1); - } - } - - // MAP>> - index = columnIndex.get("t_complex"); - if (index != null) { - if ((rowNumber % 33) == 0) { - assertNull(row.getField(index)); - } - else { - List expected1 = ImmutableList.of("test abc", 0.1); - List expected2 = ImmutableList.of("test xyz", 0.2); - assertEquals(row.getField(index), ImmutableMap.of(1, ImmutableList.of(expected1, expected2))); - } - } - - // NEW COLUMN - assertNull(row.getField(columnIndex.get("new_column"))); - - long newCompletedBytes = pageSource.getCompletedBytes(); - assertTrue(newCompletedBytes >= completedBytes); - // some formats (e.g., parquet) over read the data by a bit - assertLessThanOrEqual(newCompletedBytes, hiveSplit.getLength() + (100 * 1024)); - completedBytes = newCompletedBytes; - } - - assertLessThanOrEqual(completedBytes, hiveSplit.getLength() + (100 * 1024)); - assertEquals(rowNumber, 100); - } - finally { - pageSource.close(); - } - } - - protected void dropTable(SchemaTableName table) - { - try (Transaction transaction = newTransaction()) { - ConnectorMetadata metadata = transaction.getMetadata(); - ConnectorSession session = newSession(); - - ConnectorTableHandle handle = metadata.getTableHandle(session, table); - if (handle == null) { - return; - } - - metadata.dropTable(session, handle); - try { - // todo I have no idea why this is needed... maybe there is a propagation delay in the metastore? - metadata.dropTable(session, handle); - fail("expected NotFoundException"); - } - catch (TableNotFoundException expected) { - } - - transaction.commit(); - } - } - - protected ConnectorTableHandle getTableHandle(ConnectorMetadata metadata, SchemaTableName tableName) - { - ConnectorTableHandle handle = metadata.getTableHandle(newSession(), tableName); - checkArgument(handle != null, "table not found: %s", tableName); - return handle; - } - - private HiveTableHandle applyFilter(ConnectorMetadata metadata, ConnectorTableHandle tableHandle, Constraint constraint) - { - return metadata.applyFilter(newSession(), tableHandle, constraint) - .map(ConstraintApplicationResult::getHandle) - .map(HiveTableHandle.class::cast) - .orElseThrow(AssertionError::new); - } - - protected MaterializedResult readTable( - Transaction transaction, - ConnectorTableHandle tableHandle, - List columnHandles, - ConnectorSession session, - TupleDomain tupleDomain, - OptionalInt expectedSplitCount, - Optional expectedStorageFormat) - throws Exception - { - tableHandle = applyFilter(transaction.getMetadata(), tableHandle, new Constraint(tupleDomain)); - List splits = getAllSplits(getSplits(splitManager, transaction, session, tableHandle)); - if (expectedSplitCount.isPresent()) { - assertEquals(splits.size(), expectedSplitCount.getAsInt()); - } - - ImmutableList.Builder allRows = ImmutableList.builder(); - for (ConnectorSplit split : splits) { - try (ConnectorPageSource pageSource = pageSourceProvider.createPageSource(transaction.getTransactionHandle(), session, split, tableHandle, columnHandles, DynamicFilter.EMPTY)) { - expectedStorageFormat.ifPresent(format -> assertPageSourceType(pageSource, format)); - MaterializedResult result = materializeSourceDataStream(session, pageSource, getTypes(columnHandles)); - allRows.addAll(result.getMaterializedRows()); - } - } - return new MaterializedResult(allRows.build(), getTypes(columnHandles)); - } - - protected HiveMetastore getMetastoreClient() - { - return metastoreClient; - } - - protected LocationService getLocationService() - { - return locationService; - } - - protected static int getSplitCount(ConnectorSplitSource splitSource) - { - int splitCount = 0; - while (!splitSource.isFinished()) { - splitCount += getFutureValue(splitSource.getNextBatch(1000)).getSplits().size(); - } - return splitCount; - } - - private List getAllSplits(ConnectorTableHandle tableHandle, Transaction transaction, ConnectorSession session) - { - return getAllSplits(getSplits(splitManager, transaction, session, tableHandle)); - } - - protected static List getAllSplits(ConnectorSplitSource splitSource) - { - ImmutableList.Builder splits = ImmutableList.builder(); - while (!splitSource.isFinished()) { - splits.addAll(getFutureValue(splitSource.getNextBatch(1000)).getSplits()); - } - return splits.build(); - } - - protected static ConnectorSplitSource getSplits(ConnectorSplitManager splitManager, Transaction transaction, ConnectorSession session, ConnectorTableHandle tableHandle) - { - return splitManager.getSplits(transaction.getTransactionHandle(), session, tableHandle, DynamicFilter.EMPTY, Constraint.alwaysTrue()); - } - - protected String getPartitionId(Object partition) - { - return ((HivePartition) partition).getPartitionId(); - } - - protected static void assertPageSourceType(ConnectorPageSource pageSource, HiveStorageFormat hiveStorageFormat) - { - if (pageSource instanceof RecordPageSource) { - RecordCursor hiveRecordCursor = ((RecordPageSource) pageSource).getCursor(); - hiveRecordCursor = ((HiveRecordCursor) hiveRecordCursor).getRegularColumnRecordCursor(); - if (hiveRecordCursor instanceof HiveBucketValidationRecordCursor) { - hiveRecordCursor = ((HiveBucketValidationRecordCursor) hiveRecordCursor).delegate(); - } - assertInstanceOf(hiveRecordCursor, recordCursorType(), hiveStorageFormat.name()); - } - else { - assertInstanceOf(((HivePageSource) pageSource).getPageSource(), pageSourceType(hiveStorageFormat), hiveStorageFormat.name()); - } - } - - private static Class recordCursorType() - { - return GenericHiveRecordCursor.class; - } - - private static Class pageSourceType(HiveStorageFormat hiveStorageFormat) - { - switch (hiveStorageFormat) { - case RCTEXT: - case RCBINARY: - return RcFilePageSource.class; - case ORC: - return OrcPageSource.class; - case PARQUET: - return ParquetPageSource.class; - case CSV: - case JSON: - case OPENX_JSON: - case TEXTFILE: - case SEQUENCEFILE: - return LinePageSource.class; - default: - throw new AssertionError("File type does not use a PageSource: " + hiveStorageFormat); - } - } - - private static void assertValueTypes(MaterializedRow row, List schema) - { - for (int columnIndex = 0; columnIndex < schema.size(); columnIndex++) { - ColumnMetadata column = schema.get(columnIndex); - Object value = row.getField(columnIndex); - if (value != null) { - if (BOOLEAN.equals(column.getType())) { - assertInstanceOf(value, Boolean.class); - } - else if (TINYINT.equals(column.getType())) { - assertInstanceOf(value, Byte.class); - } - else if (SMALLINT.equals(column.getType())) { - assertInstanceOf(value, Short.class); - } - else if (INTEGER.equals(column.getType())) { - assertInstanceOf(value, Integer.class); - } - else if (BIGINT.equals(column.getType())) { - assertInstanceOf(value, Long.class); - } - else if (DOUBLE.equals(column.getType())) { - assertInstanceOf(value, Double.class); - } - else if (REAL.equals(column.getType())) { - assertInstanceOf(value, Float.class); - } - else if (column.getType() instanceof VarcharType) { - assertInstanceOf(value, String.class); - } - else if (column.getType() instanceof CharType) { - assertInstanceOf(value, String.class); - } - else if (VARBINARY.equals(column.getType())) { - assertInstanceOf(value, SqlVarbinary.class); - } - else if (TIMESTAMP_MILLIS.equals(column.getType())) { - assertInstanceOf(value, SqlTimestamp.class); - } - else if (TIMESTAMP_TZ_MILLIS.equals(column.getType())) { - assertInstanceOf(value, SqlTimestampWithTimeZone.class); - } - else if (DATE.equals(column.getType())) { - assertInstanceOf(value, SqlDate.class); - } - else if (column.getType() instanceof ArrayType || column.getType() instanceof RowType) { - assertInstanceOf(value, List.class); - } - else if (column.getType() instanceof MapType) { - assertInstanceOf(value, Map.class); - } - else { - fail("Unknown primitive type " + columnIndex); - } - } - } - } - - private static void assertPrimitiveField(Map map, String name, Type type, boolean partitionKey) - { - assertTrue(map.containsKey(name)); - ColumnMetadata column = map.get(name); - assertEquals(column.getType(), type, name); - assertEquals(column.getExtraInfo(), columnExtraInfo(partitionKey)); - } - - protected static ImmutableMap indexColumns(List columnHandles) - { - ImmutableMap.Builder index = ImmutableMap.builder(); - int i = 0; - for (ColumnHandle columnHandle : columnHandles) { - HiveColumnHandle hiveColumnHandle = (HiveColumnHandle) columnHandle; - index.put(hiveColumnHandle.getName(), i); - i++; - } - return index.buildOrThrow(); - } - - protected static ImmutableMap indexColumns(ConnectorTableMetadata tableMetadata) - { - ImmutableMap.Builder index = ImmutableMap.builder(); - int i = 0; - for (ColumnMetadata columnMetadata : tableMetadata.getColumns()) { - index.put(columnMetadata.getName(), i); - i++; - } - return index.buildOrThrow(); - } - - protected SchemaTableName temporaryTable(String tableName) - { - return temporaryTable(database, tableName); - } - - protected static SchemaTableName temporaryTable(String database, String tableName) - { - String randomName = UUID.randomUUID().toString().toLowerCase(ENGLISH).replace("-", ""); - return new SchemaTableName(database, TEMPORARY_TABLE_PREFIX + tableName + "_" + randomName); - } - - protected static Map createTableProperties(HiveStorageFormat storageFormat) - { - return createTableProperties(storageFormat, ImmutableList.of()); - } - - protected static Map createTableProperties(HiveStorageFormat storageFormat, Iterable partitionedBy) - { - return ImmutableMap.builder() - .put(STORAGE_FORMAT_PROPERTY, storageFormat) - .put(PARTITIONED_BY_PROPERTY, ImmutableList.copyOf(partitionedBy)) - .put(BUCKETED_BY_PROPERTY, ImmutableList.of()) - .put(BUCKET_COUNT_PROPERTY, 0) - .put(SORTED_BY_PROPERTY, ImmutableList.of()) - .buildOrThrow(); - } - - protected static List filterNonHiddenColumnHandles(Collection columnHandles) - { - return columnHandles.stream() - .filter(columnHandle -> !((HiveColumnHandle) columnHandle).isHidden()) - .collect(toList()); - } - - protected static List filterNonHiddenColumnMetadata(Collection columnMetadatas) - { - return columnMetadatas.stream() - .filter(columnMetadata -> !columnMetadata.isHidden()) - .collect(toList()); - } - - private void createEmptyTable(SchemaTableName schemaTableName, HiveStorageFormat hiveStorageFormat, List columns, List partitionColumns) - throws Exception - { - createEmptyTable(schemaTableName, hiveStorageFormat, columns, partitionColumns, Optional.empty(), false); - } - - private void createEmptyTable( - SchemaTableName schemaTableName, - HiveStorageFormat hiveStorageFormat, - List columns, - List partitionColumns, - Optional bucketProperty) - throws Exception - { - createEmptyTable(schemaTableName, hiveStorageFormat, columns, partitionColumns, bucketProperty, false); - } - - protected void createEmptyTable( - SchemaTableName schemaTableName, - HiveStorageFormat hiveStorageFormat, - List columns, - List partitionColumns, - Optional bucketProperty, - boolean isTransactional) - throws Exception - { - Path targetPath; - - try (Transaction transaction = newTransaction()) { - ConnectorSession session = newSession(); - - String tableOwner = session.getUser(); - String schemaName = schemaTableName.getSchemaName(); - String tableName = schemaTableName.getTableName(); - - LocationService locationService = getLocationService(); - targetPath = new Path(locationService.forNewTable(transaction.getMetastore(), session, schemaName, tableName).toString()); - - ImmutableMap.Builder tableParamBuilder = ImmutableMap.builder() - .put(PRESTO_VERSION_NAME, TEST_SERVER_VERSION) - .put(PRESTO_QUERY_ID_NAME, session.getQueryId()); - if (isTransactional) { - tableParamBuilder.put(TRANSACTIONAL, "true"); - } - Table.Builder tableBuilder = Table.builder() - .setDatabaseName(schemaName) - .setTableName(tableName) - .setOwner(Optional.of(tableOwner)) - .setTableType(TableType.MANAGED_TABLE.name()) - .setParameters(tableParamBuilder.buildOrThrow()) - .setDataColumns(columns) - .setPartitionColumns(partitionColumns); - - tableBuilder.getStorageBuilder() - .setLocation(targetPath.toString()) - .setStorageFormat(StorageFormat.create(hiveStorageFormat.getSerde(), hiveStorageFormat.getInputFormat(), hiveStorageFormat.getOutputFormat())) - .setBucketProperty(bucketProperty) - .setSerdeParameters(ImmutableMap.of()); - - PrincipalPrivileges principalPrivileges = testingPrincipalPrivilege(tableOwner, session.getUser()); - transaction.getMetastore().createTable(session, tableBuilder.build(), principalPrivileges, Optional.empty(), Optional.empty(), true, ZERO_TABLE_STATISTICS, false); - - transaction.commit(); - } - - HdfsContext context = new HdfsContext(newSession()); - List targetDirectoryList = listDirectory(context, targetPath); - assertEquals(targetDirectoryList, ImmutableList.of()); - } - - private void alterBucketProperty(SchemaTableName schemaTableName, Optional bucketProperty) - { - try (Transaction transaction = newTransaction()) { - ConnectorSession session = newSession(); - - String tableOwner = session.getUser(); - String schemaName = schemaTableName.getSchemaName(); - String tableName = schemaTableName.getTableName(); - - Optional
table = transaction.getMetastore().getTable(schemaName, tableName); - Table.Builder tableBuilder = Table.builder(table.get()); - tableBuilder.getStorageBuilder().setBucketProperty(bucketProperty); - PrincipalPrivileges principalPrivileges = testingPrincipalPrivilege(tableOwner, session.getUser()); - transaction.getMetastore().replaceTable(schemaName, tableName, tableBuilder.build(), principalPrivileges); - - transaction.commit(); - } - } - - protected PrincipalPrivileges testingPrincipalPrivilege(ConnectorSession session) - { - return testingPrincipalPrivilege(session.getUser(), session.getUser()); - } - - protected PrincipalPrivileges testingPrincipalPrivilege(String tableOwner, String grantor) - { - return new PrincipalPrivileges( - ImmutableMultimap.builder() - .put(tableOwner, new HivePrivilegeInfo(HivePrivilege.SELECT, true, new HivePrincipal(USER, grantor), new HivePrincipal(USER, grantor))) - .put(tableOwner, new HivePrivilegeInfo(HivePrivilege.INSERT, true, new HivePrincipal(USER, grantor), new HivePrincipal(USER, grantor))) - .put(tableOwner, new HivePrivilegeInfo(HivePrivilege.UPDATE, true, new HivePrincipal(USER, grantor), new HivePrincipal(USER, grantor))) - .put(tableOwner, new HivePrivilegeInfo(HivePrivilege.DELETE, true, new HivePrincipal(USER, grantor), new HivePrincipal(USER, grantor))) - .build(), - ImmutableMultimap.of()); - } - - private List listDirectory(HdfsContext context, Path path) - throws IOException - { - FileSystem fileSystem = hdfsEnvironment.getFileSystem(context, path); - return Arrays.stream(fileSystem.listStatus(path)) - .map(FileStatus::getPath) - .map(Path::getName) - .filter(name -> !name.startsWith(".trino")) - .collect(toList()); - } - - @Test - public void testTransactionDeleteInsert() - throws Exception - { - doTestTransactionDeleteInsert( - RCBINARY, - true, - ImmutableList.builder() - .add(new TransactionDeleteInsertTestCase(false, false, ROLLBACK_RIGHT_AWAY, Optional.empty())) - .add(new TransactionDeleteInsertTestCase(false, false, ROLLBACK_AFTER_DELETE, Optional.empty())) - .add(new TransactionDeleteInsertTestCase(false, false, ROLLBACK_AFTER_BEGIN_INSERT, Optional.empty())) - .add(new TransactionDeleteInsertTestCase(false, false, ROLLBACK_AFTER_APPEND_PAGE, Optional.empty())) - .add(new TransactionDeleteInsertTestCase(false, false, ROLLBACK_AFTER_SINK_FINISH, Optional.empty())) - .add(new TransactionDeleteInsertTestCase(false, false, ROLLBACK_AFTER_FINISH_INSERT, Optional.empty())) - .add(new TransactionDeleteInsertTestCase(false, false, COMMIT, Optional.of(new AddPartitionFailure()))) - .add(new TransactionDeleteInsertTestCase(false, false, COMMIT, Optional.of(new DirectoryRenameFailure()))) - .add(new TransactionDeleteInsertTestCase(false, false, COMMIT, Optional.of(new FileRenameFailure()))) - .add(new TransactionDeleteInsertTestCase(true, false, COMMIT, Optional.of(new DropPartitionFailure()))) - .add(new TransactionDeleteInsertTestCase(true, true, COMMIT, Optional.empty())) - .build()); - } - - @Test - public void testPreferredInsertLayout() - throws Exception - { - SchemaTableName tableName = temporaryTable("empty_partitioned_table"); - - try { - Column partitioningColumn = new Column("column2", HIVE_STRING, Optional.empty()); - List columns = ImmutableList.of( - new Column("column1", HIVE_STRING, Optional.empty()), - partitioningColumn); - createEmptyTable(tableName, ORC, columns, ImmutableList.of(partitioningColumn)); - - try (Transaction transaction = newTransaction()) { - ConnectorMetadata metadata = transaction.getMetadata(); - ConnectorSession session = newSession(); - ConnectorTableHandle tableHandle = getTableHandle(metadata, tableName); - Optional insertLayout = metadata.getInsertLayout(session, tableHandle); - assertTrue(insertLayout.isPresent()); - assertFalse(insertLayout.get().getPartitioning().isPresent()); - assertEquals(insertLayout.get().getPartitionColumns(), ImmutableList.of(partitioningColumn.getName())); - } - } - finally { - dropTable(tableName); - } - } - - @Test - public void testInsertBucketedTableLayout() - throws Exception - { - insertBucketedTableLayout(false); - } - - @Test - public void testInsertBucketedTransactionalTableLayout() - throws Exception - { - insertBucketedTableLayout(true); - } - - protected void insertBucketedTableLayout(boolean transactional) - throws Exception - { - SchemaTableName tableName = temporaryTable("empty_bucketed_table"); - try { - List columns = ImmutableList.of( - new Column("column1", HIVE_STRING, Optional.empty()), - new Column("column2", HIVE_LONG, Optional.empty())); - HiveBucketProperty bucketProperty = new HiveBucketProperty(ImmutableList.of("column1"), BUCKETING_V1, 4, ImmutableList.of()); - createEmptyTable(tableName, ORC, columns, ImmutableList.of(), Optional.of(bucketProperty), transactional); - - try (Transaction transaction = newTransaction()) { - ConnectorMetadata metadata = transaction.getMetadata(); - ConnectorSession session = newSession(); - ConnectorTableHandle tableHandle = getTableHandle(metadata, tableName); - Optional insertLayout = metadata.getInsertLayout(session, tableHandle); - assertTrue(insertLayout.isPresent()); - ConnectorPartitioningHandle partitioningHandle = new HivePartitioningHandle( - bucketProperty.getBucketingVersion(), - bucketProperty.getBucketCount(), - ImmutableList.of(HIVE_STRING), - OptionalInt.empty(), - false); - assertEquals(insertLayout.get().getPartitioning(), Optional.of(partitioningHandle)); - assertEquals(insertLayout.get().getPartitionColumns(), ImmutableList.of("column1")); - ConnectorBucketNodeMap connectorBucketNodeMap = nodePartitioningProvider.getBucketNodeMapping(transaction.getTransactionHandle(), session, partitioningHandle).orElseThrow(); - assertEquals(connectorBucketNodeMap.getBucketCount(), 4); - assertFalse(connectorBucketNodeMap.hasFixedMapping()); - } - } - finally { - dropTable(tableName); - } - } - - @Test - public void testInsertPartitionedBucketedTableLayout() - throws Exception - { - insertPartitionedBucketedTableLayout(false); - } - - @Test - public void testInsertPartitionedBucketedTransactionalTableLayout() - throws Exception - { - insertPartitionedBucketedTableLayout(true); - } - - protected void insertPartitionedBucketedTableLayout(boolean transactional) - throws Exception - { - SchemaTableName tableName = temporaryTable("empty_partitioned_table"); - try { - Column partitioningColumn = new Column("column2", HIVE_LONG, Optional.empty()); - List columns = ImmutableList.of( - new Column("column1", HIVE_STRING, Optional.empty()), - partitioningColumn); - HiveBucketProperty bucketProperty = new HiveBucketProperty(ImmutableList.of("column1"), BUCKETING_V1, 4, ImmutableList.of()); - createEmptyTable(tableName, ORC, columns, ImmutableList.of(partitioningColumn), Optional.of(bucketProperty), transactional); - - try (Transaction transaction = newTransaction()) { - ConnectorMetadata metadata = transaction.getMetadata(); - ConnectorSession session = newSession(); - ConnectorTableHandle tableHandle = getTableHandle(metadata, tableName); - Optional insertLayout = metadata.getInsertLayout(session, tableHandle); - assertTrue(insertLayout.isPresent()); - ConnectorPartitioningHandle partitioningHandle = new HivePartitioningHandle( - bucketProperty.getBucketingVersion(), - bucketProperty.getBucketCount(), - ImmutableList.of(HIVE_STRING), - OptionalInt.empty(), - true); - assertEquals(insertLayout.get().getPartitioning(), Optional.of(partitioningHandle)); - assertEquals(insertLayout.get().getPartitionColumns(), ImmutableList.of("column1", "column2")); - ConnectorBucketNodeMap connectorBucketNodeMap = nodePartitioningProvider.getBucketNodeMapping(transaction.getTransactionHandle(), session, partitioningHandle).orElseThrow(); - assertEquals(connectorBucketNodeMap.getBucketCount(), 32); - assertFalse(connectorBucketNodeMap.hasFixedMapping()); - } - } - finally { - dropTable(tableName); - } - } - - @Test - public void testPreferredCreateTableLayout() - { - try (Transaction transaction = newTransaction()) { - ConnectorMetadata metadata = transaction.getMetadata(); - ConnectorSession session = newSession(); - Optional newTableLayout = metadata.getNewTableLayout( - session, - new ConnectorTableMetadata( - new SchemaTableName("schema", "table"), - ImmutableList.of( - new ColumnMetadata("column1", BIGINT), - new ColumnMetadata("column2", BIGINT)), - ImmutableMap.of( - PARTITIONED_BY_PROPERTY, ImmutableList.of("column2"), - BUCKETED_BY_PROPERTY, ImmutableList.of(), - BUCKET_COUNT_PROPERTY, 0, - SORTED_BY_PROPERTY, ImmutableList.of()))); - assertTrue(newTableLayout.isPresent()); - assertFalse(newTableLayout.get().getPartitioning().isPresent()); - assertEquals(newTableLayout.get().getPartitionColumns(), ImmutableList.of("column2")); - } - } - - @Test - public void testCreateBucketedTableLayout() - { - try (Transaction transaction = newTransaction()) { - ConnectorMetadata metadata = transaction.getMetadata(); - ConnectorSession session = newSession(); - Optional newTableLayout = metadata.getNewTableLayout( - session, - new ConnectorTableMetadata( - new SchemaTableName("schema", "table"), - ImmutableList.of( - new ColumnMetadata("column1", BIGINT), - new ColumnMetadata("column2", BIGINT)), - ImmutableMap.of( - PARTITIONED_BY_PROPERTY, ImmutableList.of(), - BUCKETED_BY_PROPERTY, ImmutableList.of("column1"), - BUCKET_COUNT_PROPERTY, 10, - SORTED_BY_PROPERTY, ImmutableList.of()))); - assertTrue(newTableLayout.isPresent()); - ConnectorPartitioningHandle partitioningHandle = new HivePartitioningHandle( - BUCKETING_V1, - 10, - ImmutableList.of(HIVE_LONG), - OptionalInt.empty(), - false); - assertEquals(newTableLayout.get().getPartitioning(), Optional.of(partitioningHandle)); - assertEquals(newTableLayout.get().getPartitionColumns(), ImmutableList.of("column1")); - ConnectorBucketNodeMap connectorBucketNodeMap = nodePartitioningProvider.getBucketNodeMapping(transaction.getTransactionHandle(), session, partitioningHandle).orElseThrow(); - assertEquals(connectorBucketNodeMap.getBucketCount(), 10); - assertFalse(connectorBucketNodeMap.hasFixedMapping()); - } - } - - @Test - public void testCreatePartitionedBucketedTableLayout() - { - try (Transaction transaction = newTransaction()) { - ConnectorMetadata metadata = transaction.getMetadata(); - ConnectorSession session = newSession(); - Optional newTableLayout = metadata.getNewTableLayout( - session, - new ConnectorTableMetadata( - new SchemaTableName("schema", "table"), - ImmutableList.of( - new ColumnMetadata("column1", BIGINT), - new ColumnMetadata("column2", BIGINT)), - ImmutableMap.of( - PARTITIONED_BY_PROPERTY, ImmutableList.of("column2"), - BUCKETED_BY_PROPERTY, ImmutableList.of("column1"), - BUCKET_COUNT_PROPERTY, 10, - SORTED_BY_PROPERTY, ImmutableList.of()))); - assertTrue(newTableLayout.isPresent()); - ConnectorPartitioningHandle partitioningHandle = new HivePartitioningHandle( - BUCKETING_V1, - 10, - ImmutableList.of(HIVE_LONG), - OptionalInt.empty(), - true); - assertEquals(newTableLayout.get().getPartitioning(), Optional.of(partitioningHandle)); - assertEquals(newTableLayout.get().getPartitionColumns(), ImmutableList.of("column1", "column2")); - ConnectorBucketNodeMap connectorBucketNodeMap = nodePartitioningProvider.getBucketNodeMapping(transaction.getTransactionHandle(), session, partitioningHandle).orElseThrow(); - assertEquals(connectorBucketNodeMap.getBucketCount(), 32); - assertFalse(connectorBucketNodeMap.hasFixedMapping()); - } - } - - @Test - public void testNewDirectoryPermissions() - throws Exception - { - SchemaTableName tableName = temporaryTable("empty_file"); - List columns = ImmutableList.of(new Column("test", HIVE_STRING, Optional.empty())); - createEmptyTable(tableName, ORC, columns, ImmutableList.of(), Optional.empty()); - try { - Transaction transaction = newTransaction(); - ConnectorSession session = newSession(); - ConnectorMetadata metadata = transaction.getMetadata(); - metadata.beginQuery(session); - - Table table = transaction.getMetastore() - .getTable(tableName.getSchemaName(), tableName.getTableName()) - .orElseThrow(); - - // create new directory and set directory permission after creation - HdfsContext context = new HdfsContext(session); - Path location = new Path(table.getStorage().getLocation()); - Path defaultPath = new Path(location + "/defaultperms"); - createDirectory(context, hdfsEnvironment, defaultPath); - FileStatus defaultFsStatus = hdfsEnvironment.getFileSystem(context, defaultPath).getFileStatus(defaultPath); - assertEquals(defaultFsStatus.getPermission().toOctal(), 777); - - // use hdfs config that skips setting directory permissions after creation - HdfsConfig configWithSkip = new HdfsConfig(); - configWithSkip.setNewDirectoryPermissions(HdfsConfig.SKIP_DIR_PERMISSIONS); - HdfsEnvironment hdfsEnvironmentWithSkip = new HdfsEnvironment( - HDFS_CONFIGURATION, - configWithSkip, - new NoHdfsAuthentication()); - - Path skipPath = new Path(location + "/skipperms"); - createDirectory(context, hdfsEnvironmentWithSkip, skipPath); - FileStatus skipFsStatus = hdfsEnvironmentWithSkip.getFileSystem(context, skipPath).getFileStatus(skipPath); - assertEquals(skipFsStatus.getPermission().toOctal(), 755); - } - finally { - dropTable(tableName); - } - } - - protected void doTestTransactionDeleteInsert(HiveStorageFormat storageFormat, boolean allowInsertExisting, List testCases) - throws Exception - { - // There are 4 types of operations on a partition: add, drop, alter (drop then add), insert existing. - // There are 12 partitions in this test, 3 for each type. - // 3 is chosen to verify that cleanups, commit aborts, rollbacks are always as complete as possible regardless of failure. - MaterializedResult beforeData = - MaterializedResult.resultBuilder(SESSION, BIGINT, createUnboundedVarcharType(), createUnboundedVarcharType()) - .row(110L, "a", "alter1") - .row(120L, "a", "insert1") - .row(140L, "a", "drop1") - .row(210L, "b", "drop2") - .row(310L, "c", "alter2") - .row(320L, "c", "alter3") - .row(510L, "e", "drop3") - .row(610L, "f", "insert2") - .row(620L, "f", "insert3") - .build(); - Domain domainToDrop = Domain.create(ValueSet.of( - createUnboundedVarcharType(), - utf8Slice("alter1"), utf8Slice("alter2"), utf8Slice("alter3"), utf8Slice("drop1"), utf8Slice("drop2"), utf8Slice("drop3")), - false); - List extraRowsForInsertExisting = ImmutableList.of(); - if (allowInsertExisting) { - extraRowsForInsertExisting = MaterializedResult.resultBuilder(SESSION, BIGINT, createUnboundedVarcharType(), createUnboundedVarcharType()) - .row(121L, "a", "insert1") - .row(611L, "f", "insert2") - .row(621L, "f", "insert3") - .build() - .getMaterializedRows(); - } - MaterializedResult insertData = - MaterializedResult.resultBuilder(SESSION, BIGINT, createUnboundedVarcharType(), createUnboundedVarcharType()) - .row(111L, "a", "alter1") - .row(131L, "a", "add1") - .row(221L, "b", "add2") - .row(311L, "c", "alter2") - .row(321L, "c", "alter3") - .row(411L, "d", "add3") - .rows(extraRowsForInsertExisting) - .build(); - MaterializedResult afterData = - MaterializedResult.resultBuilder(SESSION, BIGINT, createUnboundedVarcharType(), createUnboundedVarcharType()) - .row(120L, "a", "insert1") - .row(610L, "f", "insert2") - .row(620L, "f", "insert3") - .rows(insertData.getMaterializedRows()) - .build(); - - for (TransactionDeleteInsertTestCase testCase : testCases) { - SchemaTableName temporaryDeleteInsert = temporaryTable("delete_insert"); - try { - createEmptyTable( - temporaryDeleteInsert, - storageFormat, - ImmutableList.of(new Column("col1", HIVE_LONG, Optional.empty())), - ImmutableList.of(new Column("pk1", HIVE_STRING, Optional.empty()), new Column("pk2", HIVE_STRING, Optional.empty()))); - insertData(temporaryDeleteInsert, beforeData); - try { - doTestTransactionDeleteInsert( - storageFormat, - temporaryDeleteInsert, - domainToDrop, - insertData, - testCase.isExpectCommittedData() ? afterData : beforeData, - testCase.getTag(), - testCase.isExpectQuerySucceed(), - testCase.getConflictTrigger()); - } - catch (AssertionError e) { - throw new AssertionError(format("Test case: %s", testCase), e); - } - } - finally { - dropTable(temporaryDeleteInsert); - } - } - } - - private void doTestTransactionDeleteInsert( - HiveStorageFormat storageFormat, - SchemaTableName tableName, - Domain domainToDrop, - MaterializedResult insertData, - MaterializedResult expectedData, - TransactionDeleteInsertTestTag tag, - boolean expectQuerySucceed, - Optional conflictTrigger) - throws Exception - { - Location writePath = null; - Location targetPath = null; - - try (Transaction transaction = newTransaction()) { - try { - ConnectorMetadata metadata = transaction.getMetadata(); - ConnectorTableHandle tableHandle = getTableHandle(metadata, tableName); - ConnectorSession session; - rollbackIfEquals(tag, ROLLBACK_RIGHT_AWAY); - - // Query 1: delete - session = newSession(); - HiveColumnHandle dsColumnHandle = (HiveColumnHandle) metadata.getColumnHandles(session, tableHandle).get("pk2"); - TupleDomain tupleDomain = TupleDomain.withColumnDomains(ImmutableMap.of( - dsColumnHandle, domainToDrop)); - Constraint constraint = new Constraint(tupleDomain, tupleDomain.asPredicate(), tupleDomain.getDomains().orElseThrow().keySet()); - tableHandle = applyFilter(metadata, tableHandle, constraint); - tableHandle = metadata.applyDelete(session, tableHandle).get(); - metadata.executeDelete(session, tableHandle); - rollbackIfEquals(tag, ROLLBACK_AFTER_DELETE); - - // Query 2: insert - session = newSession(); - ConnectorInsertTableHandle insertTableHandle = metadata.beginInsert(session, tableHandle, ImmutableList.of(), NO_RETRIES); - rollbackIfEquals(tag, ROLLBACK_AFTER_BEGIN_INSERT); - writePath = getStagingPathRoot(insertTableHandle); - targetPath = getTargetPathRoot(insertTableHandle); - ConnectorPageSink sink = pageSinkProvider.createPageSink(transaction.getTransactionHandle(), session, insertTableHandle, TESTING_PAGE_SINK_ID); - sink.appendPage(insertData.toPage()); - rollbackIfEquals(tag, ROLLBACK_AFTER_APPEND_PAGE); - Collection fragments = getFutureValue(sink.finish()); - rollbackIfEquals(tag, ROLLBACK_AFTER_SINK_FINISH); - metadata.finishInsert(session, insertTableHandle, fragments, ImmutableList.of()); - rollbackIfEquals(tag, ROLLBACK_AFTER_FINISH_INSERT); - - assertEquals(tag, COMMIT); - - if (conflictTrigger.isPresent()) { - JsonCodec partitionUpdateCodec = JsonCodec.jsonCodec(PartitionUpdate.class); - List partitionUpdates = fragments.stream() - .map(Slice::getBytes) - .map(partitionUpdateCodec::fromJson) - .collect(toList()); - conflictTrigger.get().triggerConflict(session, tableName, insertTableHandle, partitionUpdates); - } - transaction.commit(); - if (conflictTrigger.isPresent()) { - assertTrue(expectQuerySucceed); - conflictTrigger.get().verifyAndCleanup(session, tableName); - } - } - catch (TestingRollbackException e) { - transaction.rollback(); - } - catch (TrinoException e) { - assertFalse(expectQuerySucceed); - if (conflictTrigger.isPresent()) { - conflictTrigger.get().verifyAndCleanup(newSession(), tableName); - } - } - } - - // check that temporary files are removed - if (writePath != null && !writePath.equals(targetPath)) { - HdfsContext context = new HdfsContext(newSession()); - FileSystem fileSystem = hdfsEnvironment.getFileSystem(context, new Path(writePath.toString())); - assertFalse(fileSystem.exists(new Path(writePath.toString()))); - } - - try (Transaction transaction = newTransaction()) { - // verify partitions - List partitionNames = transaction.getMetastore() - .getPartitionNames(tableName.getSchemaName(), tableName.getTableName()) - .orElseThrow(() -> new AssertionError("Table does not exist: " + tableName)); - assertEqualsIgnoreOrder( - partitionNames, - expectedData.getMaterializedRows().stream() - .map(row -> format("pk1=%s/pk2=%s", row.getField(1), row.getField(2))) - .distinct() - .collect(toImmutableList())); - - // load the new table - ConnectorSession session = newSession(); - ConnectorMetadata metadata = transaction.getMetadata(); - metadata.beginQuery(session); - ConnectorTableHandle tableHandle = getTableHandle(metadata, tableName); - List columnHandles = filterNonHiddenColumnHandles(metadata.getColumnHandles(session, tableHandle).values()); - - // verify the data - MaterializedResult result = readTable(transaction, tableHandle, columnHandles, session, TupleDomain.all(), OptionalInt.empty(), Optional.of(storageFormat)); - assertEqualsIgnoreOrder(result.getMaterializedRows(), expectedData.getMaterializedRows()); - } - } - - private static void rollbackIfEquals(TransactionDeleteInsertTestTag tag, TransactionDeleteInsertTestTag expectedTag) - { - if (expectedTag == tag) { - throw new TestingRollbackException(); - } - } - - private static class TestingRollbackException - extends RuntimeException - { - } - - protected static class TransactionDeleteInsertTestCase - { - private final boolean expectCommittedData; - private final boolean expectQuerySucceed; - private final TransactionDeleteInsertTestTag tag; - private final Optional conflictTrigger; - - public TransactionDeleteInsertTestCase(boolean expectCommittedData, boolean expectQuerySucceed, TransactionDeleteInsertTestTag tag, Optional conflictTrigger) - { - this.expectCommittedData = expectCommittedData; - this.expectQuerySucceed = expectQuerySucceed; - this.tag = tag; - this.conflictTrigger = conflictTrigger; - } - - public boolean isExpectCommittedData() - { - return expectCommittedData; - } - - public boolean isExpectQuerySucceed() - { - return expectQuerySucceed; - } - - public TransactionDeleteInsertTestTag getTag() - { - return tag; - } - - public Optional getConflictTrigger() - { - return conflictTrigger; - } - - @Override - public String toString() - { - return toStringHelper(this) - .add("tag", tag) - .add("conflictTrigger", conflictTrigger.map(conflictTrigger -> conflictTrigger.getClass().getName())) - .add("expectCommittedData", expectCommittedData) - .add("expectQuerySucceed", expectQuerySucceed) - .toString(); - } - } - - protected enum TransactionDeleteInsertTestTag - { - ROLLBACK_RIGHT_AWAY, - ROLLBACK_AFTER_DELETE, - ROLLBACK_AFTER_BEGIN_INSERT, - ROLLBACK_AFTER_APPEND_PAGE, - ROLLBACK_AFTER_SINK_FINISH, - ROLLBACK_AFTER_FINISH_INSERT, - COMMIT, - } - - protected interface ConflictTrigger - { - void triggerConflict(ConnectorSession session, SchemaTableName tableName, ConnectorInsertTableHandle insertTableHandle, List partitionUpdates) - throws IOException; - - void verifyAndCleanup(ConnectorSession session, SchemaTableName tableName) - throws IOException; - } - - protected class AddPartitionFailure - implements ConflictTrigger - { - private final ImmutableList copyPartitionFrom = ImmutableList.of("a", "insert1"); - private final String partitionNameToConflict = "pk1=b/pk2=add2"; - private Partition conflictPartition; - - @Override - public void triggerConflict(ConnectorSession session, SchemaTableName tableName, ConnectorInsertTableHandle insertTableHandle, List partitionUpdates) - { - // This method bypasses transaction interface because this method is inherently hacky and doesn't work well with the transaction abstraction. - // Additionally, this method is not part of a test. Its purpose is to set up an environment for another test. - HiveMetastore metastoreClient = getMetastoreClient(); - Table table = metastoreClient.getTable(tableName.getSchemaName(), tableName.getTableName()) - .orElseThrow(() -> new TableNotFoundException(tableName)); - Optional partition = metastoreClient.getPartition(table, copyPartitionFrom); - conflictPartition = Partition.builder(partition.get()) - .setValues(toPartitionValues(partitionNameToConflict)) - .build(); - metastoreClient.addPartitions( - tableName.getSchemaName(), - tableName.getTableName(), - ImmutableList.of(new PartitionWithStatistics(conflictPartition, partitionNameToConflict, PartitionStatistics.empty()))); - } - - @Override - public void verifyAndCleanup(ConnectorSession session, SchemaTableName tableName) - { - // This method bypasses transaction interface because this method is inherently hacky and doesn't work well with the transaction abstraction. - // Additionally, this method is not part of a test. Its purpose is to set up an environment for another test. - HiveMetastore metastoreClient = getMetastoreClient(); - Table table = metastoreClient.getTable(tableName.getSchemaName(), tableName.getTableName()) - .orElseThrow(() -> new TableNotFoundException(tableName)); - Optional actualPartition = metastoreClient.getPartition(table, toPartitionValues(partitionNameToConflict)); - // Make sure the partition inserted to trigger conflict was not overwritten - // Checking storage location is sufficient because implement never uses .../pk1=a/pk2=a2 as the directory for partition [b, b2]. - assertEquals(actualPartition.get().getStorage().getLocation(), conflictPartition.getStorage().getLocation()); - metastoreClient.dropPartition(tableName.getSchemaName(), tableName.getTableName(), conflictPartition.getValues(), false); - } - } - - protected class DropPartitionFailure - implements ConflictTrigger - { - private final ImmutableList partitionValueToConflict = ImmutableList.of("b", "drop2"); - - @Override - public void triggerConflict(ConnectorSession session, SchemaTableName tableName, ConnectorInsertTableHandle insertTableHandle, List partitionUpdates) - { - // This method bypasses transaction interface because this method is inherently hacky and doesn't work well with the transaction abstraction. - // Additionally, this method is not part of a test. Its purpose is to set up an environment for another test. - HiveMetastore metastoreClient = getMetastoreClient(); - metastoreClient.dropPartition(tableName.getSchemaName(), tableName.getTableName(), partitionValueToConflict, false); - } - - @Override - public void verifyAndCleanup(ConnectorSession session, SchemaTableName tableName) - { - // Do not add back the deleted partition because the implementation is expected to move forward instead of backward when delete fails - } - } - - protected class DirectoryRenameFailure - implements ConflictTrigger - { - private HdfsContext context; - private Path path; - - @Override - public void triggerConflict(ConnectorSession session, SchemaTableName tableName, ConnectorInsertTableHandle insertTableHandle, List partitionUpdates) - { - Location writePath = getStagingPathRoot(insertTableHandle); - Location targetPath = getTargetPathRoot(insertTableHandle); - if (writePath.equals(targetPath)) { - // This conflict does not apply. Trigger a rollback right away so that this test case passes. - throw new TestingRollbackException(); - } - path = new Path(targetPath.appendPath("pk1=b").appendPath("pk2=add2").toString()); - context = new HdfsContext(session); - createDirectory(context, hdfsEnvironment, new Path(path.toString())); - } - - @Override - public void verifyAndCleanup(ConnectorSession session, SchemaTableName tableName) - throws IOException - { - assertEquals(listDirectory(context, path), ImmutableList.of()); - hdfsEnvironment.getFileSystem(context, path).delete(path, false); - } - } - - protected class FileRenameFailure - implements ConflictTrigger - { - private HdfsContext context; - private Path path; - - @Override - public void triggerConflict(ConnectorSession session, SchemaTableName tableName, ConnectorInsertTableHandle insertTableHandle, List partitionUpdates) - throws IOException - { - for (PartitionUpdate partitionUpdate : partitionUpdates) { - if ("pk2=insert2".equals(partitionUpdate.getTargetPath().fileName())) { - path = new Path(partitionUpdate.getTargetPath().toString(), partitionUpdate.getFileNames().get(0)); - break; - } - } - assertNotNull(path); - - context = new HdfsContext(session); - FileSystem fileSystem = hdfsEnvironment.getFileSystem(context, path); - fileSystem.createNewFile(path); - } - - @Override - public void verifyAndCleanup(ConnectorSession session, SchemaTableName tableName) - throws IOException - { - // The file we added to trigger a conflict was cleaned up because it matches the query prefix. - // Consider this the same as a network failure that caused the successful creation of file not reported to the caller. - assertFalse(hdfsEnvironment.getFileSystem(context, path).exists(path)); - } - } - - private static class CountingDirectoryLister - implements DirectoryLister - { - private final AtomicInteger listCount = new AtomicInteger(); - - @Override - public RemoteIterator listFilesRecursively(TrinoFileSystem fs, Table table, Location location) - throws IOException - { - listCount.incrementAndGet(); - return new TrinoFileStatusRemoteIterator(fs.listFiles(location)); - } - - public int getListCount() - { - return listCount.get(); - } - - @Override - public void invalidate(Partition partition) - { - } - - @Override - public void invalidate(Table table) - { - } - } -} diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/AbstractTestHiveFileFormats.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/AbstractTestHiveFileFormats.java index 8cd5bcdeb358..3485b52fc3a0 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/AbstractTestHiveFileFormats.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/AbstractTestHiveFileFormats.java @@ -13,31 +13,41 @@ */ package io.trino.plugin.hive; +import com.google.common.base.VerifyException; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import io.airlift.slice.Slice; import io.airlift.slice.Slices; import io.trino.filesystem.Location; -import io.trino.plugin.hive.metastore.StorageFormat; +import io.trino.filesystem.TrinoFileSystem; +import io.trino.filesystem.TrinoFileSystemFactory; +import io.trino.hadoop.ConfigurationInstantiator; +import io.trino.plugin.base.type.DecodedTimestamp; +import io.trino.plugin.hive.util.HiveTypeTranslator; import io.trino.spi.Page; import io.trino.spi.PageBuilder; -import io.trino.spi.block.Block; +import io.trino.spi.block.ArrayBlockBuilder; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.MapBlockBuilder; +import io.trino.spi.block.RowBlockBuilder; import io.trino.spi.connector.ConnectorPageSource; import io.trino.spi.connector.ConnectorSession; import io.trino.spi.connector.RecordCursor; +import io.trino.spi.security.ConnectorIdentity; import io.trino.spi.type.ArrayType; import io.trino.spi.type.CharType; import io.trino.spi.type.DateType; import io.trino.spi.type.DecimalType; import io.trino.spi.type.Int128; +import io.trino.spi.type.Int128Math; +import io.trino.spi.type.MapType; import io.trino.spi.type.RowType; import io.trino.spi.type.SqlDate; -import io.trino.spi.type.SqlDecimal; import io.trino.spi.type.SqlTimestamp; import io.trino.spi.type.SqlVarbinary; import io.trino.spi.type.TimestampType; import io.trino.spi.type.Type; +import io.trino.spi.type.TypeOperators; import io.trino.spi.type.VarcharType; import io.trino.testing.MaterializedResult; import io.trino.testing.MaterializedRow; @@ -47,19 +57,20 @@ import org.apache.hadoop.hive.common.type.HiveDecimal; import org.apache.hadoop.hive.common.type.HiveVarchar; import org.apache.hadoop.hive.common.type.Timestamp; -import org.apache.hadoop.hive.ql.exec.FileSinkOperator.RecordWriter; +import org.apache.hadoop.hive.ql.exec.FileSinkOperator; import org.apache.hadoop.hive.ql.io.HiveOutputFormat; import org.apache.hadoop.hive.serde2.Serializer; +import org.apache.hadoop.hive.serde2.io.HiveDecimalWritable; import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspector; -import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspector.Category; import org.apache.hadoop.hive.serde2.objectinspector.SettableStructObjectInspector; import org.apache.hadoop.hive.serde2.objectinspector.StructField; import org.apache.hadoop.hive.serde2.objectinspector.primitive.JavaHiveCharObjectInspector; import org.apache.hadoop.hive.serde2.objectinspector.primitive.JavaHiveDecimalObjectInspector; +import org.apache.hadoop.hive.serde2.typeinfo.CharTypeInfo; import org.apache.hadoop.hive.serde2.typeinfo.DecimalTypeInfo; +import org.apache.hadoop.hive.serde2.typeinfo.VarcharTypeInfo; import org.apache.hadoop.io.Text; import org.apache.hadoop.io.Writable; -import org.apache.hadoop.mapred.FileSplit; import org.apache.hadoop.mapred.JobConf; import org.joda.time.DateTime; import org.joda.time.DateTimeZone; @@ -67,7 +78,7 @@ import java.io.File; import java.io.IOException; -import java.lang.invoke.MethodHandle; +import java.io.OutputStream; import java.math.BigDecimal; import java.math.BigInteger; import java.util.ArrayList; @@ -82,31 +93,39 @@ import java.util.stream.Collectors; import static com.google.common.base.Preconditions.checkArgument; +import static com.google.common.base.Verify.verify; import static com.google.common.collect.ImmutableList.toImmutableList; -import static com.google.common.collect.ImmutableMap.toImmutableMap; -import static io.trino.hadoop.ConfigurationInstantiator.newEmptyConfiguration; +import static io.airlift.slice.Slices.utf8Slice; +import static io.trino.plugin.base.type.TrinoTimestampEncoderFactory.createTimestampEncoder; import static io.trino.plugin.hive.HiveColumnHandle.ColumnType.PARTITION_KEY; import static io.trino.plugin.hive.HiveColumnHandle.ColumnType.REGULAR; import static io.trino.plugin.hive.HiveColumnHandle.createBaseColumn; -import static io.trino.plugin.hive.HiveColumnProjectionInfo.generatePartialName; import static io.trino.plugin.hive.HivePartitionKey.HIVE_DEFAULT_DYNAMIC_PARTITION; import static io.trino.plugin.hive.HiveTestUtils.SESSION; -import static io.trino.plugin.hive.HiveTestUtils.isDistinctFrom; import static io.trino.plugin.hive.HiveTestUtils.mapType; import static io.trino.plugin.hive.acid.AcidTransaction.NO_ACID_TRANSACTION; import static io.trino.plugin.hive.util.CompressionConfigUtil.configureCompression; +import static io.trino.plugin.hive.util.HiveTypeTranslator.toHiveType; import static io.trino.plugin.hive.util.HiveUtil.isStructuralType; -import static io.trino.plugin.hive.util.SerDeUtils.serializeObject; +import static io.trino.plugin.hive.util.SerdeConstants.LIST_COLUMNS; +import static io.trino.plugin.hive.util.SerdeConstants.LIST_COLUMN_TYPES; import static io.trino.spi.type.BigintType.BIGINT; import static io.trino.spi.type.BooleanType.BOOLEAN; import static io.trino.spi.type.CharType.createCharType; import static io.trino.spi.type.Chars.padSpaces; +import static io.trino.spi.type.Chars.truncateToLengthAndTrimSpaces; +import static io.trino.spi.type.DateType.DATE; import static io.trino.spi.type.DoubleType.DOUBLE; import static io.trino.spi.type.IntegerType.INTEGER; import static io.trino.spi.type.RealType.REAL; +import static io.trino.spi.type.RowType.field; +import static io.trino.spi.type.RowType.rowType; import static io.trino.spi.type.SmallintType.SMALLINT; +import static io.trino.spi.type.TimestampType.TIMESTAMP_MILLIS; +import static io.trino.spi.type.Timestamps.round; import static io.trino.spi.type.TinyintType.TINYINT; import static io.trino.spi.type.VarbinaryType.VARBINARY; +import static io.trino.spi.type.VarcharType.VARCHAR; import static io.trino.spi.type.VarcharType.createUnboundedVarcharType; import static io.trino.spi.type.VarcharType.createVarcharType; import static io.trino.testing.DateTimeTestingUtils.sqlTimestampOf; @@ -120,31 +139,32 @@ import static io.trino.type.InternalTypeManager.TESTING_TYPE_MANAGER; import static java.lang.Float.intBitsToFloat; import static java.lang.Math.floorDiv; +import static java.lang.Math.toIntExact; import static java.nio.charset.StandardCharsets.UTF_8; +import static java.nio.file.Files.readAllBytes; import static java.util.Arrays.fill; import static java.util.Objects.requireNonNull; -import static java.util.function.Function.identity; import static java.util.stream.Collectors.toList; +import static org.apache.hadoop.hive.common.type.HiveVarchar.MAX_VARCHAR_LENGTH; import static org.apache.hadoop.hive.serde2.objectinspector.ObjectInspectorFactory.getStandardListObjectInspector; import static org.apache.hadoop.hive.serde2.objectinspector.ObjectInspectorFactory.getStandardMapObjectInspector; import static org.apache.hadoop.hive.serde2.objectinspector.ObjectInspectorFactory.getStandardStructObjectInspector; +import static org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorFactory.getPrimitiveJavaObjectInspector; import static org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorFactory.javaBooleanObjectInspector; import static org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorFactory.javaByteArrayObjectInspector; import static org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorFactory.javaByteObjectInspector; import static org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorFactory.javaDateObjectInspector; import static org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorFactory.javaDoubleObjectInspector; import static org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorFactory.javaFloatObjectInspector; -import static org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorFactory.javaHiveVarcharObjectInspector; import static org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorFactory.javaIntObjectInspector; import static org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorFactory.javaLongObjectInspector; import static org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorFactory.javaShortObjectInspector; import static org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorFactory.javaStringObjectInspector; import static org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorFactory.javaTimestampObjectInspector; import static org.apache.hadoop.hive.serde2.typeinfo.TypeInfoFactory.getCharTypeInfo; +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.data.Offset.offset; import static org.joda.time.DateTimeZone.UTC; -import static org.testng.Assert.assertEquals; -import static org.testng.Assert.assertFalse; -import static org.testng.Assert.assertTrue; public abstract class AbstractTestHiveFileFormats { @@ -152,17 +172,43 @@ public abstract class AbstractTestHiveFileFormats private static final double EPSILON = 0.001; + private static final Type VARCHAR_100 = createVarcharType(100); + private static final Type VARCHAR_HIVE_MAX = createVarcharType(MAX_VARCHAR_LENGTH); + private static final Type CHAR_10 = createCharType(10); + private static final String VARCHAR_MAX_LENGTH_STRING; + private static final long DATE_MILLIS_UTC = new DateTime(2011, 5, 6, 0, 0, UTC).getMillis(); private static final long DATE_DAYS = TimeUnit.MILLISECONDS.toDays(DATE_MILLIS_UTC); private static final String DATE_STRING = DateTimeFormat.forPattern("yyyy-MM-dd").withZoneUTC().print(DATE_MILLIS_UTC); private static final Date HIVE_DATE = Date.ofEpochMilli(DATE_MILLIS_UTC); - private static final DateTime TIMESTAMP = new DateTime(2011, 5, 6, 7, 8, 9, 123, UTC); - private static final long TIMESTAMP_MICROS = TIMESTAMP.getMillis() * MICROSECONDS_PER_MILLISECOND; - private static final String TIMESTAMP_STRING = DateTimeFormat.forPattern("yyyy-MM-dd HH:mm:ss.SSS").withZoneUTC().print(TIMESTAMP.getMillis()); - private static final Timestamp HIVE_TIMESTAMP = Timestamp.ofEpochMilli(TIMESTAMP.getMillis()); - - private static final String VARCHAR_MAX_LENGTH_STRING; + private static final DateTime TIMESTAMP_VALUE = new DateTime(2011, 5, 6, 7, 8, 9, 123, UTC); + private static final long TIMESTAMP_MICROS_VALUE = TIMESTAMP_VALUE.getMillis() * MICROSECONDS_PER_MILLISECOND; + private static final String TIMESTAMP_STRING_VALUE = DateTimeFormat.forPattern("yyyy-MM-dd HH:mm:ss.SSS").withZoneUTC().print(TIMESTAMP_VALUE.getMillis()); + private static final Timestamp HIVE_TIMESTAMP = Timestamp.ofEpochMilli(TIMESTAMP_VALUE.getMillis()); + + private static final DecimalType DECIMAL_TYPE_2 = DecimalType.createDecimalType(2, 1); + private static final DecimalType DECIMAL_TYPE_4 = DecimalType.createDecimalType(4, 2); + private static final DecimalType DECIMAL_TYPE_8 = DecimalType.createDecimalType(8, 4); + private static final DecimalType DECIMAL_TYPE_17 = DecimalType.createDecimalType(17, 8); + private static final DecimalType DECIMAL_TYPE_18 = DecimalType.createDecimalType(18, 8); + private static final DecimalType DECIMAL_TYPE_38 = DecimalType.createDecimalType(38, 16); + + private static final HiveDecimal WRITE_DECIMAL_2 = HiveDecimal.create(new BigDecimal("-1.2")); + private static final HiveDecimal WRITE_DECIMAL_4 = HiveDecimal.create(new BigDecimal("12.3")); + private static final HiveDecimal WRITE_DECIMAL_8 = HiveDecimal.create(new BigDecimal("-1234.5678")); + private static final HiveDecimal WRITE_DECIMAL_17 = HiveDecimal.create(new BigDecimal("123456789.1234")); + private static final HiveDecimal WRITE_DECIMAL_18 = HiveDecimal.create(new BigDecimal("-1234567890.12345678")); + private static final HiveDecimal WRITE_DECIMAL_38 = HiveDecimal.create(new BigDecimal("1234567890123456789012.12345678")); + + private static final BigDecimal EXPECTED_DECIMAL_2 = new BigDecimal("-1.2"); + private static final BigDecimal EXPECTED_DECIMAL_4 = new BigDecimal("12.30"); + private static final BigDecimal EXPECTED_DECIMAL_8 = new BigDecimal("-1234.5678"); + private static final BigDecimal EXPECTED_DECIMAL_17 = new BigDecimal("123456789.12340000"); + private static final BigDecimal EXPECTED_DECIMAL_18 = new BigDecimal("-1234567890.12345678"); + private static final BigDecimal EXPECTED_DECIMAL_38 = new BigDecimal("1234567890123456789012.1234567800000000"); + + private static final TypeOperators TYPE_OPERATORS = TESTING_TYPE_MANAGER.getTypeOperators(); static { char[] varcharMaxLengthCharArray = new char[HiveVarchar.MAX_VARCHAR_LENGTH]; @@ -207,249 +253,248 @@ public abstract class AbstractTestHiveFileFormats private static final JavaHiveCharObjectInspector CHAR_INSPECTOR_LENGTH_10 = new JavaHiveCharObjectInspector(getCharTypeInfo(10)); - // TODO: support null values and determine if timestamp and binary are allowed as partition keys - public static final List TEST_COLUMNS = ImmutableList.builder() - .add(new TestColumn("p_empty_string", javaStringObjectInspector, "", Slices.EMPTY_SLICE, true)) - .add(new TestColumn("p_string", javaStringObjectInspector, "test", Slices.utf8Slice("test"), true)) - .add(new TestColumn("p_empty_varchar", javaHiveVarcharObjectInspector, "", Slices.EMPTY_SLICE, true)) - .add(new TestColumn("p_varchar", javaHiveVarcharObjectInspector, "test", Slices.utf8Slice("test"), true)) - .add(new TestColumn("p_varchar_max_length", javaHiveVarcharObjectInspector, VARCHAR_MAX_LENGTH_STRING, Slices.utf8Slice(VARCHAR_MAX_LENGTH_STRING), true)) - .add(new TestColumn("p_char_10", CHAR_INSPECTOR_LENGTH_10, "test", Slices.utf8Slice("test"), true)) - .add(new TestColumn("p_tinyint", javaByteObjectInspector, "1", (byte) 1, true)) - .add(new TestColumn("p_smallint", javaShortObjectInspector, "2", (short) 2, true)) - .add(new TestColumn("p_int", javaIntObjectInspector, "3", 3, true)) - .add(new TestColumn("p_bigint", javaLongObjectInspector, "4", 4L, true)) - .add(new TestColumn("p_float", javaFloatObjectInspector, "5.1", 5.1f, true)) - .add(new TestColumn("p_double", javaDoubleObjectInspector, "6.2", 6.2, true)) - .add(new TestColumn("p_boolean", javaBooleanObjectInspector, "true", true, true)) - .add(new TestColumn("p_date", javaDateObjectInspector, DATE_STRING, DATE_DAYS, true)) - .add(new TestColumn("p_timestamp", javaTimestampObjectInspector, TIMESTAMP_STRING, TIMESTAMP_MICROS, true)) - .add(new TestColumn("p_decimal_precision_2", DECIMAL_INSPECTOR_PRECISION_2, WRITE_DECIMAL_PRECISION_2.toString(), EXPECTED_DECIMAL_PRECISION_2, true)) - .add(new TestColumn("p_decimal_precision_4", DECIMAL_INSPECTOR_PRECISION_4, WRITE_DECIMAL_PRECISION_4.toString(), EXPECTED_DECIMAL_PRECISION_4, true)) - .add(new TestColumn("p_decimal_precision_8", DECIMAL_INSPECTOR_PRECISION_8, WRITE_DECIMAL_PRECISION_8.toString(), EXPECTED_DECIMAL_PRECISION_8, true)) - .add(new TestColumn("p_decimal_precision_17", DECIMAL_INSPECTOR_PRECISION_17, WRITE_DECIMAL_PRECISION_17.toString(), EXPECTED_DECIMAL_PRECISION_17, true)) - .add(new TestColumn("p_decimal_precision_18", DECIMAL_INSPECTOR_PRECISION_18, WRITE_DECIMAL_PRECISION_18.toString(), EXPECTED_DECIMAL_PRECISION_18, true)) - .add(new TestColumn("p_decimal_precision_38", DECIMAL_INSPECTOR_PRECISION_38, WRITE_DECIMAL_PRECISION_38.toString() + "BD", EXPECTED_DECIMAL_PRECISION_38, true)) -// .add(new TestColumn("p_binary", javaByteArrayObjectInspector, "test2", Slices.utf8Slice("test2"), true)) - .add(new TestColumn("p_null_string", javaStringObjectInspector, HIVE_DEFAULT_DYNAMIC_PARTITION, null, true)) - .add(new TestColumn("p_null_varchar", javaHiveVarcharObjectInspector, HIVE_DEFAULT_DYNAMIC_PARTITION, null, true)) - .add(new TestColumn("p_null_char", CHAR_INSPECTOR_LENGTH_10, HIVE_DEFAULT_DYNAMIC_PARTITION, null, true)) - .add(new TestColumn("p_null_tinyint", javaByteObjectInspector, HIVE_DEFAULT_DYNAMIC_PARTITION, null, true)) - .add(new TestColumn("p_null_smallint", javaShortObjectInspector, HIVE_DEFAULT_DYNAMIC_PARTITION, null, true)) - .add(new TestColumn("p_null_int", javaIntObjectInspector, HIVE_DEFAULT_DYNAMIC_PARTITION, null, true)) - .add(new TestColumn("p_null_bigint", javaLongObjectInspector, HIVE_DEFAULT_DYNAMIC_PARTITION, null, true)) - .add(new TestColumn("p_null_float", javaFloatObjectInspector, HIVE_DEFAULT_DYNAMIC_PARTITION, null, true)) - .add(new TestColumn("p_null_double", javaDoubleObjectInspector, HIVE_DEFAULT_DYNAMIC_PARTITION, null, true)) - .add(new TestColumn("p_null_boolean", javaBooleanObjectInspector, HIVE_DEFAULT_DYNAMIC_PARTITION, null, true)) - .add(new TestColumn("p_null_date", javaDateObjectInspector, HIVE_DEFAULT_DYNAMIC_PARTITION, null, true)) - .add(new TestColumn("p_null_timestamp", javaTimestampObjectInspector, HIVE_DEFAULT_DYNAMIC_PARTITION, null, true)) - .add(new TestColumn("p_null_decimal_precision_2", DECIMAL_INSPECTOR_PRECISION_2, HIVE_DEFAULT_DYNAMIC_PARTITION, null, true)) - .add(new TestColumn("p_null_decimal_precision_4", DECIMAL_INSPECTOR_PRECISION_4, HIVE_DEFAULT_DYNAMIC_PARTITION, null, true)) - .add(new TestColumn("p_null_decimal_precision_8", DECIMAL_INSPECTOR_PRECISION_8, HIVE_DEFAULT_DYNAMIC_PARTITION, null, true)) - .add(new TestColumn("p_null_decimal_precision_17", DECIMAL_INSPECTOR_PRECISION_17, HIVE_DEFAULT_DYNAMIC_PARTITION, null, true)) - .add(new TestColumn("p_null_decimal_precision_18", DECIMAL_INSPECTOR_PRECISION_18, HIVE_DEFAULT_DYNAMIC_PARTITION, null, true)) - .add(new TestColumn("p_null_decimal_precision_38", DECIMAL_INSPECTOR_PRECISION_38, HIVE_DEFAULT_DYNAMIC_PARTITION, null, true)) - -// .add(new TestColumn("p_null_binary", javaByteArrayObjectInspector, HIVE_DEFAULT_DYNAMIC_PARTITION, null, true)) - .add(new TestColumn("t_null_string", javaStringObjectInspector, null, null)) - .add(new TestColumn("t_null_varchar", javaHiveVarcharObjectInspector, null, null)) - .add(new TestColumn("t_null_char", CHAR_INSPECTOR_LENGTH_10, null, null)) - .add(new TestColumn("t_null_array_int", getStandardListObjectInspector(javaIntObjectInspector), null, null)) - .add(new TestColumn("t_null_decimal_precision_2", DECIMAL_INSPECTOR_PRECISION_2, null, null)) - .add(new TestColumn("t_null_decimal_precision_4", DECIMAL_INSPECTOR_PRECISION_4, null, null)) - .add(new TestColumn("t_null_decimal_precision_8", DECIMAL_INSPECTOR_PRECISION_8, null, null)) - .add(new TestColumn("t_null_decimal_precision_17", DECIMAL_INSPECTOR_PRECISION_17, null, null)) - .add(new TestColumn("t_null_decimal_precision_18", DECIMAL_INSPECTOR_PRECISION_18, null, null)) - .add(new TestColumn("t_null_decimal_precision_38", DECIMAL_INSPECTOR_PRECISION_38, null, null)) - .add(new TestColumn("t_empty_string", javaStringObjectInspector, "", Slices.EMPTY_SLICE)) - .add(new TestColumn("t_string", javaStringObjectInspector, "test", Slices.utf8Slice("test"))) - .add(new TestColumn("t_empty_varchar", javaHiveVarcharObjectInspector, new HiveVarchar("", HiveVarchar.MAX_VARCHAR_LENGTH), Slices.EMPTY_SLICE)) - .add(new TestColumn("t_varchar", javaHiveVarcharObjectInspector, new HiveVarchar("test", HiveVarchar.MAX_VARCHAR_LENGTH), Slices.utf8Slice("test"))) - .add(new TestColumn("t_varchar_max_length", javaHiveVarcharObjectInspector, new HiveVarchar(VARCHAR_MAX_LENGTH_STRING, HiveVarchar.MAX_VARCHAR_LENGTH), Slices.utf8Slice(VARCHAR_MAX_LENGTH_STRING))) - .add(new TestColumn("t_char", CHAR_INSPECTOR_LENGTH_10, "test", Slices.utf8Slice("test"))) - .add(new TestColumn("t_tinyint", javaByteObjectInspector, (byte) 1, (byte) 1)) - .add(new TestColumn("t_smallint", javaShortObjectInspector, (short) 2, (short) 2)) - .add(new TestColumn("t_int", javaIntObjectInspector, 3, 3)) - .add(new TestColumn("t_bigint", javaLongObjectInspector, 4L, 4L)) - .add(new TestColumn("t_float", javaFloatObjectInspector, 5.1f, 5.1f)) - .add(new TestColumn("t_double", javaDoubleObjectInspector, 6.2, 6.2)) - .add(new TestColumn("t_boolean_true", javaBooleanObjectInspector, true, true)) - .add(new TestColumn("t_boolean_false", javaBooleanObjectInspector, false, false)) - .add(new TestColumn("t_date", javaDateObjectInspector, HIVE_DATE, DATE_DAYS)) - .add(new TestColumn("t_timestamp", javaTimestampObjectInspector, HIVE_TIMESTAMP, TIMESTAMP_MICROS)) - .add(new TestColumn("t_decimal_precision_2", DECIMAL_INSPECTOR_PRECISION_2, WRITE_DECIMAL_PRECISION_2, EXPECTED_DECIMAL_PRECISION_2)) - .add(new TestColumn("t_decimal_precision_4", DECIMAL_INSPECTOR_PRECISION_4, WRITE_DECIMAL_PRECISION_4, EXPECTED_DECIMAL_PRECISION_4)) - .add(new TestColumn("t_decimal_precision_8", DECIMAL_INSPECTOR_PRECISION_8, WRITE_DECIMAL_PRECISION_8, EXPECTED_DECIMAL_PRECISION_8)) - .add(new TestColumn("t_decimal_precision_17", DECIMAL_INSPECTOR_PRECISION_17, WRITE_DECIMAL_PRECISION_17, EXPECTED_DECIMAL_PRECISION_17)) - .add(new TestColumn("t_decimal_precision_18", DECIMAL_INSPECTOR_PRECISION_18, WRITE_DECIMAL_PRECISION_18, EXPECTED_DECIMAL_PRECISION_18)) - .add(new TestColumn("t_decimal_precision_38", DECIMAL_INSPECTOR_PRECISION_38, WRITE_DECIMAL_PRECISION_38, EXPECTED_DECIMAL_PRECISION_38)) - .add(new TestColumn("t_binary", javaByteArrayObjectInspector, Slices.utf8Slice("test2").getBytes(), Slices.utf8Slice("test2"))) + protected static final List TEST_COLUMNS = ImmutableList.builder() + .add(new TestColumn("p_empty_string", VARCHAR, "", Slices.EMPTY_SLICE, true)) + .add(new TestColumn("p_string", VARCHAR, "test", utf8Slice("test"), true)) + .add(new TestColumn("p_empty_varchar", VARCHAR_100, "", Slices.EMPTY_SLICE, true)) + .add(new TestColumn("p_varchar", VARCHAR_100, "test", utf8Slice("test"), true)) + .add(new TestColumn("p_varchar_max_length", VARCHAR_HIVE_MAX, VARCHAR_MAX_LENGTH_STRING, utf8Slice(VARCHAR_MAX_LENGTH_STRING), true)) + .add(new TestColumn("p_char_10", CHAR_10, "test", utf8Slice("test"), true)) + .add(new TestColumn("p_tinyint", TINYINT, "1", (byte) 1, true)) + .add(new TestColumn("p_smallint", SMALLINT, "2", (short) 2, true)) + .add(new TestColumn("p_int", INTEGER, "3", 3, true)) + .add(new TestColumn("p_bigint", BIGINT, "4", 4L, true)) + .add(new TestColumn("p_float", REAL, "5.1", 5.1f, true)) + .add(new TestColumn("p_double", DOUBLE, "6.2", 6.2, true)) + .add(new TestColumn("p_boolean", BOOLEAN, "true", true, true)) + .add(new TestColumn("p_date", DATE, DATE_STRING, DATE_DAYS, true)) + .add(new TestColumn("p_timestamp", TIMESTAMP_MILLIS, TIMESTAMP_STRING_VALUE, TIMESTAMP_MICROS_VALUE, true)) + .add(new TestColumn("p_decimal_2", DECIMAL_TYPE_2, WRITE_DECIMAL_2.toString(), EXPECTED_DECIMAL_2, true)) + .add(new TestColumn("p_decimal_4", DECIMAL_TYPE_4, WRITE_DECIMAL_4.toString(), EXPECTED_DECIMAL_4, true)) + .add(new TestColumn("p_decimal_8", DECIMAL_TYPE_8, WRITE_DECIMAL_8.toString(), EXPECTED_DECIMAL_8, true)) + .add(new TestColumn("p_decimal_17", DECIMAL_TYPE_17, WRITE_DECIMAL_17.toString(), EXPECTED_DECIMAL_17, true)) + .add(new TestColumn("p_decimal_18", DECIMAL_TYPE_18, WRITE_DECIMAL_18.toString(), EXPECTED_DECIMAL_18, true)) + .add(new TestColumn("p_decimal_38", DECIMAL_TYPE_38, WRITE_DECIMAL_38.toString() + "BD", EXPECTED_DECIMAL_38, true)) + .add(new TestColumn("p_null_string", VARCHAR, HIVE_DEFAULT_DYNAMIC_PARTITION, null, true)) + .add(new TestColumn("p_null_varchar", VARCHAR_100, HIVE_DEFAULT_DYNAMIC_PARTITION, null, true)) + .add(new TestColumn("p_null_char", CHAR_10, HIVE_DEFAULT_DYNAMIC_PARTITION, null, true)) + .add(new TestColumn("p_null_tinyint", TINYINT, HIVE_DEFAULT_DYNAMIC_PARTITION, null, true)) + .add(new TestColumn("p_null_smallint", SMALLINT, HIVE_DEFAULT_DYNAMIC_PARTITION, null, true)) + .add(new TestColumn("p_null_int", INTEGER, HIVE_DEFAULT_DYNAMIC_PARTITION, null, true)) + .add(new TestColumn("p_null_bigint", BIGINT, HIVE_DEFAULT_DYNAMIC_PARTITION, null, true)) + .add(new TestColumn("p_null_float", REAL, HIVE_DEFAULT_DYNAMIC_PARTITION, null, true)) + .add(new TestColumn("p_null_double", DOUBLE, HIVE_DEFAULT_DYNAMIC_PARTITION, null, true)) + .add(new TestColumn("p_null_boolean", BOOLEAN, HIVE_DEFAULT_DYNAMIC_PARTITION, null, true)) + .add(new TestColumn("p_null_date", DATE, HIVE_DEFAULT_DYNAMIC_PARTITION, null, true)) + .add(new TestColumn("p_null_timestamp", TIMESTAMP_MILLIS, HIVE_DEFAULT_DYNAMIC_PARTITION, null, true)) + .add(new TestColumn("p_null_decimal_2", DECIMAL_TYPE_2, HIVE_DEFAULT_DYNAMIC_PARTITION, null, true)) + .add(new TestColumn("p_null_decimal_4", DECIMAL_TYPE_4, HIVE_DEFAULT_DYNAMIC_PARTITION, null, true)) + .add(new TestColumn("p_null_decimal_8", DECIMAL_TYPE_8, HIVE_DEFAULT_DYNAMIC_PARTITION, null, true)) + .add(new TestColumn("p_null_decimal_17", DECIMAL_TYPE_17, HIVE_DEFAULT_DYNAMIC_PARTITION, null, true)) + .add(new TestColumn("p_null_decimal_18", DECIMAL_TYPE_18, HIVE_DEFAULT_DYNAMIC_PARTITION, null, true)) + .add(new TestColumn("p_null_decimal_38", DECIMAL_TYPE_38, HIVE_DEFAULT_DYNAMIC_PARTITION, null, true)) + + .add(new TestColumn("t_null_string", VARCHAR, null, null)) + .add(new TestColumn("t_null_varchar", VARCHAR_100, null, null)) + .add(new TestColumn("t_null_char", CHAR_10, null, null)) + .add(new TestColumn("t_null_array_int", new ArrayType(INTEGER), null, null)) + .add(new TestColumn("t_null_decimal_2", DECIMAL_TYPE_2, null, null)) + .add(new TestColumn("t_null_decimal_4", DECIMAL_TYPE_4, null, null)) + .add(new TestColumn("t_null_decimal_8", DECIMAL_TYPE_8, null, null)) + .add(new TestColumn("t_null_decimal_17", DECIMAL_TYPE_17, null, null)) + .add(new TestColumn("t_null_decimal_18", DECIMAL_TYPE_18, null, null)) + .add(new TestColumn("t_null_decimal_38", DECIMAL_TYPE_38, null, null)) + .add(new TestColumn("t_empty_string", VARCHAR, "", Slices.EMPTY_SLICE)) + .add(new TestColumn("t_string", VARCHAR, "test", utf8Slice("test"))) + .add(new TestColumn("t_empty_varchar", VARCHAR_HIVE_MAX, new HiveVarchar("", MAX_VARCHAR_LENGTH), Slices.EMPTY_SLICE)) + .add(new TestColumn("t_varchar", VARCHAR_HIVE_MAX, new HiveVarchar("test", MAX_VARCHAR_LENGTH), utf8Slice("test"))) + .add(new TestColumn("t_varchar_max_length", VARCHAR_HIVE_MAX, new HiveVarchar(VARCHAR_MAX_LENGTH_STRING, MAX_VARCHAR_LENGTH), utf8Slice(VARCHAR_MAX_LENGTH_STRING))) + .add(new TestColumn("t_char", CHAR_10, "test", utf8Slice("test"))) + .add(new TestColumn("t_tinyint", TINYINT, (byte) 1, (byte) 1)) + .add(new TestColumn("t_smallint", SMALLINT, (short) 2, (short) 2)) + .add(new TestColumn("t_int", INTEGER, 3, 3)) + .add(new TestColumn("t_bigint", BIGINT, 4L, 4L)) + .add(new TestColumn("t_float", REAL, 5.1f, 5.1f)) + .add(new TestColumn("t_double", DOUBLE, 6.2, 6.2)) + .add(new TestColumn("t_boolean_true", BOOLEAN, true, true)) + .add(new TestColumn("t_boolean_false", BOOLEAN, false, false)) + .add(new TestColumn("t_date", DATE, HIVE_DATE, DATE_DAYS)) + .add(new TestColumn("t_timestamp", TIMESTAMP_MILLIS, HIVE_TIMESTAMP, TIMESTAMP_MICROS_VALUE)) + .add(new TestColumn("t_decimal_2", DECIMAL_TYPE_2, WRITE_DECIMAL_2, EXPECTED_DECIMAL_2)) + .add(new TestColumn("t_decimal_4", DECIMAL_TYPE_4, WRITE_DECIMAL_4, EXPECTED_DECIMAL_4)) + .add(new TestColumn("t_decimal_8", DECIMAL_TYPE_8, WRITE_DECIMAL_8, EXPECTED_DECIMAL_8)) + .add(new TestColumn("t_decimal_17", DECIMAL_TYPE_17, WRITE_DECIMAL_17, EXPECTED_DECIMAL_17)) + .add(new TestColumn("t_decimal_18", DECIMAL_TYPE_18, WRITE_DECIMAL_18, EXPECTED_DECIMAL_18)) + .add(new TestColumn("t_decimal_38", DECIMAL_TYPE_38, WRITE_DECIMAL_38, EXPECTED_DECIMAL_38)) + .add(new TestColumn("t_binary", VARBINARY, utf8Slice("test2").getBytes(), utf8Slice("test2"))) .add(new TestColumn("t_map_string", - getStandardMapObjectInspector(javaStringObjectInspector, javaStringObjectInspector), + new MapType(VARCHAR, VARCHAR, TYPE_OPERATORS), ImmutableMap.of("test", "test"), mapBlockOf(createUnboundedVarcharType(), createUnboundedVarcharType(), "test", "test"))) .add(new TestColumn("t_map_tinyint", - getStandardMapObjectInspector(javaByteObjectInspector, javaByteObjectInspector), + new MapType(TINYINT, TINYINT, TYPE_OPERATORS), ImmutableMap.of((byte) 1, (byte) 1), mapBlockOf(TINYINT, TINYINT, (byte) 1, (byte) 1))) .add(new TestColumn("t_map_varchar", - getStandardMapObjectInspector(javaHiveVarcharObjectInspector, javaHiveVarcharObjectInspector), - ImmutableMap.of(new HiveVarchar("test", HiveVarchar.MAX_VARCHAR_LENGTH), new HiveVarchar("test", HiveVarchar.MAX_VARCHAR_LENGTH)), - mapBlockOf(createVarcharType(HiveVarchar.MAX_VARCHAR_LENGTH), createVarcharType(HiveVarchar.MAX_VARCHAR_LENGTH), "test", "test"))) + new MapType(VARCHAR_HIVE_MAX, VARCHAR_HIVE_MAX, TYPE_OPERATORS), + ImmutableMap.of(new HiveVarchar("test", MAX_VARCHAR_LENGTH), new HiveVarchar("test", MAX_VARCHAR_LENGTH)), + mapBlockOf(createVarcharType(MAX_VARCHAR_LENGTH), createVarcharType(MAX_VARCHAR_LENGTH), "test", "test"))) .add(new TestColumn("t_map_char", - getStandardMapObjectInspector(CHAR_INSPECTOR_LENGTH_10, CHAR_INSPECTOR_LENGTH_10), + new MapType(CHAR_10, CHAR_10, TYPE_OPERATORS), ImmutableMap.of(new HiveChar("test", 10), new HiveChar("test", 10)), mapBlockOf(createCharType(10), createCharType(10), "test", "test"))) .add(new TestColumn("t_map_smallint", - getStandardMapObjectInspector(javaShortObjectInspector, javaShortObjectInspector), + new MapType(SMALLINT, SMALLINT, TYPE_OPERATORS), ImmutableMap.of((short) 2, (short) 2), mapBlockOf(SMALLINT, SMALLINT, (short) 2, (short) 2))) .add(new TestColumn("t_map_null_key", - getStandardMapObjectInspector(javaLongObjectInspector, javaLongObjectInspector), + new MapType(BIGINT, BIGINT, TYPE_OPERATORS), asMap(new Long[] {null, 2L}, new Long[] {0L, 3L}), mapBlockOf(BIGINT, BIGINT, 2, 3))) .add(new TestColumn("t_map_int", - getStandardMapObjectInspector(javaIntObjectInspector, javaIntObjectInspector), + new MapType(INTEGER, INTEGER, TYPE_OPERATORS), ImmutableMap.of(3, 3), mapBlockOf(INTEGER, INTEGER, 3, 3))) .add(new TestColumn("t_map_bigint", - getStandardMapObjectInspector(javaLongObjectInspector, javaLongObjectInspector), + new MapType(BIGINT, BIGINT, TYPE_OPERATORS), ImmutableMap.of(4L, 4L), mapBlockOf(BIGINT, BIGINT, 4L, 4L))) .add(new TestColumn("t_map_float", - getStandardMapObjectInspector(javaFloatObjectInspector, javaFloatObjectInspector), + new MapType(REAL, REAL, TYPE_OPERATORS), ImmutableMap.of(5.0f, 5.0f), mapBlockOf(REAL, REAL, 5.0f, 5.0f))) .add(new TestColumn("t_map_double", - getStandardMapObjectInspector(javaDoubleObjectInspector, javaDoubleObjectInspector), + new MapType(DOUBLE, DOUBLE, TYPE_OPERATORS), ImmutableMap.of(6.0, 6.0), mapBlockOf(DOUBLE, DOUBLE, 6.0, 6.0))) .add(new TestColumn("t_map_boolean", - getStandardMapObjectInspector(javaBooleanObjectInspector, javaBooleanObjectInspector), + new MapType(BOOLEAN, BOOLEAN, TYPE_OPERATORS), ImmutableMap.of(true, true), mapBlockOf(BOOLEAN, BOOLEAN, true, true))) .add(new TestColumn("t_map_date", - getStandardMapObjectInspector(javaDateObjectInspector, javaDateObjectInspector), + new MapType(DATE, DATE, TYPE_OPERATORS), ImmutableMap.of(HIVE_DATE, HIVE_DATE), - mapBlockOf(DateType.DATE, DateType.DATE, DATE_DAYS, DATE_DAYS))) + mapBlockOf(DATE, DATE, DATE_DAYS, DATE_DAYS))) .add(new TestColumn("t_map_timestamp", - getStandardMapObjectInspector(javaTimestampObjectInspector, javaTimestampObjectInspector), + new MapType(TIMESTAMP_MILLIS, TIMESTAMP_MILLIS, TYPE_OPERATORS), ImmutableMap.of(HIVE_TIMESTAMP, HIVE_TIMESTAMP), - mapBlockOf(TimestampType.TIMESTAMP_MILLIS, TimestampType.TIMESTAMP_MILLIS, TIMESTAMP_MICROS, TIMESTAMP_MICROS))) - .add(new TestColumn("t_map_decimal_precision_2", - getStandardMapObjectInspector(DECIMAL_INSPECTOR_PRECISION_2, DECIMAL_INSPECTOR_PRECISION_2), - ImmutableMap.of(WRITE_DECIMAL_PRECISION_2, WRITE_DECIMAL_PRECISION_2), - decimalMapBlockOf(DECIMAL_TYPE_PRECISION_2, EXPECTED_DECIMAL_PRECISION_2))) - .add(new TestColumn("t_map_decimal_precision_4", - getStandardMapObjectInspector(DECIMAL_INSPECTOR_PRECISION_4, DECIMAL_INSPECTOR_PRECISION_4), - ImmutableMap.of(WRITE_DECIMAL_PRECISION_4, WRITE_DECIMAL_PRECISION_4), - decimalMapBlockOf(DECIMAL_TYPE_PRECISION_4, EXPECTED_DECIMAL_PRECISION_4))) - .add(new TestColumn("t_map_decimal_precision_8", - getStandardMapObjectInspector(DECIMAL_INSPECTOR_PRECISION_8, DECIMAL_INSPECTOR_PRECISION_8), - ImmutableMap.of(WRITE_DECIMAL_PRECISION_8, WRITE_DECIMAL_PRECISION_8), - decimalMapBlockOf(DECIMAL_TYPE_PRECISION_8, EXPECTED_DECIMAL_PRECISION_8))) - .add(new TestColumn("t_map_decimal_precision_17", - getStandardMapObjectInspector(DECIMAL_INSPECTOR_PRECISION_17, DECIMAL_INSPECTOR_PRECISION_17), - ImmutableMap.of(WRITE_DECIMAL_PRECISION_17, WRITE_DECIMAL_PRECISION_17), - decimalMapBlockOf(DECIMAL_TYPE_PRECISION_17, EXPECTED_DECIMAL_PRECISION_17))) - .add(new TestColumn("t_map_decimal_precision_18", - getStandardMapObjectInspector(DECIMAL_INSPECTOR_PRECISION_18, DECIMAL_INSPECTOR_PRECISION_18), - ImmutableMap.of(WRITE_DECIMAL_PRECISION_18, WRITE_DECIMAL_PRECISION_18), - decimalMapBlockOf(DECIMAL_TYPE_PRECISION_18, EXPECTED_DECIMAL_PRECISION_18))) - .add(new TestColumn("t_map_decimal_precision_38", - getStandardMapObjectInspector(DECIMAL_INSPECTOR_PRECISION_38, DECIMAL_INSPECTOR_PRECISION_38), - ImmutableMap.of(WRITE_DECIMAL_PRECISION_38, WRITE_DECIMAL_PRECISION_38), - decimalMapBlockOf(DECIMAL_TYPE_PRECISION_38, EXPECTED_DECIMAL_PRECISION_38))) - .add(new TestColumn("t_array_empty", getStandardListObjectInspector(javaStringObjectInspector), ImmutableList.of(), arrayBlockOf(createUnboundedVarcharType()))) - .add(new TestColumn("t_array_string", getStandardListObjectInspector(javaStringObjectInspector), ImmutableList.of("test"), arrayBlockOf(createUnboundedVarcharType(), "test"))) - .add(new TestColumn("t_array_tinyint", getStandardListObjectInspector(javaByteObjectInspector), ImmutableList.of((byte) 1), arrayBlockOf(TINYINT, (byte) 1))) - .add(new TestColumn("t_array_smallint", getStandardListObjectInspector(javaShortObjectInspector), ImmutableList.of((short) 2), arrayBlockOf(SMALLINT, (short) 2))) - .add(new TestColumn("t_array_int", getStandardListObjectInspector(javaIntObjectInspector), ImmutableList.of(3), arrayBlockOf(INTEGER, 3))) - .add(new TestColumn("t_array_bigint", getStandardListObjectInspector(javaLongObjectInspector), ImmutableList.of(4L), arrayBlockOf(BIGINT, 4L))) - .add(new TestColumn("t_array_float", getStandardListObjectInspector(javaFloatObjectInspector), ImmutableList.of(5.0f), arrayBlockOf(REAL, 5.0f))) - .add(new TestColumn("t_array_double", getStandardListObjectInspector(javaDoubleObjectInspector), ImmutableList.of(6.0), arrayBlockOf(DOUBLE, 6.0))) - .add(new TestColumn("t_array_boolean", getStandardListObjectInspector(javaBooleanObjectInspector), ImmutableList.of(true), arrayBlockOf(BOOLEAN, true))) + mapBlockOf(TIMESTAMP_MILLIS, TIMESTAMP_MILLIS, TIMESTAMP_MICROS_VALUE, TIMESTAMP_MICROS_VALUE))) + .add(new TestColumn("t_map_decimal_2", + new MapType(DECIMAL_TYPE_2, DECIMAL_TYPE_2, TYPE_OPERATORS), + ImmutableMap.of(WRITE_DECIMAL_2, WRITE_DECIMAL_2), + decimalMapBlockOf(DECIMAL_TYPE_2, EXPECTED_DECIMAL_2))) + .add(new TestColumn("t_map_decimal_4", + new MapType(DECIMAL_TYPE_4, DECIMAL_TYPE_4, TYPE_OPERATORS), + ImmutableMap.of(WRITE_DECIMAL_4, WRITE_DECIMAL_4), + decimalMapBlockOf(DECIMAL_TYPE_4, EXPECTED_DECIMAL_4))) + .add(new TestColumn("t_map_decimal_8", + new MapType(DECIMAL_TYPE_8, DECIMAL_TYPE_8, TYPE_OPERATORS), + ImmutableMap.of(WRITE_DECIMAL_8, WRITE_DECIMAL_8), + decimalMapBlockOf(DECIMAL_TYPE_8, EXPECTED_DECIMAL_8))) + .add(new TestColumn("t_map_decimal_17", + new MapType(DECIMAL_TYPE_17, DECIMAL_TYPE_17, TYPE_OPERATORS), + ImmutableMap.of(WRITE_DECIMAL_17, WRITE_DECIMAL_17), + decimalMapBlockOf(DECIMAL_TYPE_17, EXPECTED_DECIMAL_17))) + .add(new TestColumn("t_map_decimal_18", + new MapType(DECIMAL_TYPE_18, DECIMAL_TYPE_18, TYPE_OPERATORS), + ImmutableMap.of(WRITE_DECIMAL_18, WRITE_DECIMAL_18), + decimalMapBlockOf(DECIMAL_TYPE_18, EXPECTED_DECIMAL_18))) + .add(new TestColumn("t_map_decimal_38", + new MapType(DECIMAL_TYPE_38, DECIMAL_TYPE_38, TYPE_OPERATORS), + ImmutableMap.of(WRITE_DECIMAL_38, WRITE_DECIMAL_38), + decimalMapBlockOf(DECIMAL_TYPE_38, EXPECTED_DECIMAL_38))) + .add(new TestColumn("t_array_empty", new ArrayType(VARCHAR), ImmutableList.of(), arrayBlockOf(createUnboundedVarcharType()))) + .add(new TestColumn("t_array_string", new ArrayType(VARCHAR), ImmutableList.of("test"), arrayBlockOf(createUnboundedVarcharType(), "test"))) + .add(new TestColumn("t_array_tinyint", new ArrayType(TINYINT), ImmutableList.of((byte) 1), arrayBlockOf(TINYINT, (byte) 1))) + .add(new TestColumn("t_array_smallint", new ArrayType(SMALLINT), ImmutableList.of((short) 2), arrayBlockOf(SMALLINT, (short) 2))) + .add(new TestColumn("t_array_int", new ArrayType(INTEGER), ImmutableList.of(3), arrayBlockOf(INTEGER, 3))) + .add(new TestColumn("t_array_bigint", new ArrayType(BIGINT), ImmutableList.of(4L), arrayBlockOf(BIGINT, 4L))) + .add(new TestColumn("t_array_float", new ArrayType(REAL), ImmutableList.of(5.0f), arrayBlockOf(REAL, 5.0f))) + .add(new TestColumn("t_array_double", new ArrayType(DOUBLE), ImmutableList.of(6.0), arrayBlockOf(DOUBLE, 6.0))) + .add(new TestColumn("t_array_boolean", new ArrayType(BOOLEAN), ImmutableList.of(true), arrayBlockOf(BOOLEAN, true))) .add(new TestColumn( "t_array_varchar", - getStandardListObjectInspector(javaHiveVarcharObjectInspector), - ImmutableList.of(new HiveVarchar("test", HiveVarchar.MAX_VARCHAR_LENGTH)), - arrayBlockOf(createVarcharType(HiveVarchar.MAX_VARCHAR_LENGTH), "test"))) + new ArrayType(VARCHAR_HIVE_MAX), + ImmutableList.of(new HiveVarchar("test", MAX_VARCHAR_LENGTH)), + arrayBlockOf(createVarcharType(MAX_VARCHAR_LENGTH), "test"))) .add(new TestColumn( "t_array_char", - getStandardListObjectInspector(CHAR_INSPECTOR_LENGTH_10), + new ArrayType(CHAR_10), ImmutableList.of(new HiveChar("test", 10)), arrayBlockOf(createCharType(10), "test"))) .add(new TestColumn("t_array_date", - getStandardListObjectInspector(javaDateObjectInspector), + new ArrayType(DATE), ImmutableList.of(HIVE_DATE), - arrayBlockOf(DateType.DATE, DATE_DAYS))) + arrayBlockOf(DATE, DATE_DAYS))) .add(new TestColumn("t_array_timestamp", - getStandardListObjectInspector(javaTimestampObjectInspector), + new ArrayType(TIMESTAMP_MILLIS), ImmutableList.of(HIVE_TIMESTAMP), - arrayBlockOf(TimestampType.TIMESTAMP_MILLIS, TIMESTAMP_MICROS))) - .add(new TestColumn("t_array_decimal_precision_2", - getStandardListObjectInspector(DECIMAL_INSPECTOR_PRECISION_2), - ImmutableList.of(WRITE_DECIMAL_PRECISION_2), - decimalArrayBlockOf(DECIMAL_TYPE_PRECISION_2, EXPECTED_DECIMAL_PRECISION_2))) - .add(new TestColumn("t_array_decimal_precision_4", - getStandardListObjectInspector(DECIMAL_INSPECTOR_PRECISION_4), - ImmutableList.of(WRITE_DECIMAL_PRECISION_4), - decimalArrayBlockOf(DECIMAL_TYPE_PRECISION_4, EXPECTED_DECIMAL_PRECISION_4))) - .add(new TestColumn("t_array_decimal_precision_8", - getStandardListObjectInspector(DECIMAL_INSPECTOR_PRECISION_8), - ImmutableList.of(WRITE_DECIMAL_PRECISION_8), - decimalArrayBlockOf(DECIMAL_TYPE_PRECISION_8, EXPECTED_DECIMAL_PRECISION_8))) - .add(new TestColumn("t_array_decimal_precision_17", - getStandardListObjectInspector(DECIMAL_INSPECTOR_PRECISION_17), - ImmutableList.of(WRITE_DECIMAL_PRECISION_17), - decimalArrayBlockOf(DECIMAL_TYPE_PRECISION_17, EXPECTED_DECIMAL_PRECISION_17))) - .add(new TestColumn("t_array_decimal_precision_18", - getStandardListObjectInspector(DECIMAL_INSPECTOR_PRECISION_18), - ImmutableList.of(WRITE_DECIMAL_PRECISION_18), - decimalArrayBlockOf(DECIMAL_TYPE_PRECISION_18, EXPECTED_DECIMAL_PRECISION_18))) - .add(new TestColumn("t_array_decimal_precision_38", - getStandardListObjectInspector(DECIMAL_INSPECTOR_PRECISION_38), - ImmutableList.of(WRITE_DECIMAL_PRECISION_38), - decimalArrayBlockOf(DECIMAL_TYPE_PRECISION_38, EXPECTED_DECIMAL_PRECISION_38))) + arrayBlockOf(TIMESTAMP_MILLIS, TIMESTAMP_MICROS_VALUE))) + .add(new TestColumn("t_array_decimal_2", + new ArrayType(DECIMAL_TYPE_2), + ImmutableList.of(WRITE_DECIMAL_2), + decimalArrayBlockOf(DECIMAL_TYPE_2, EXPECTED_DECIMAL_2))) + .add(new TestColumn("t_array_decimal_4", + new ArrayType(DECIMAL_TYPE_4), + ImmutableList.of(WRITE_DECIMAL_4), + decimalArrayBlockOf(DECIMAL_TYPE_4, EXPECTED_DECIMAL_4))) + .add(new TestColumn("t_array_decimal_8", + new ArrayType(DECIMAL_TYPE_8), + ImmutableList.of(WRITE_DECIMAL_8), + decimalArrayBlockOf(DECIMAL_TYPE_8, EXPECTED_DECIMAL_8))) + .add(new TestColumn("t_array_decimal_17", + new ArrayType(DECIMAL_TYPE_17), + ImmutableList.of(WRITE_DECIMAL_17), + decimalArrayBlockOf(DECIMAL_TYPE_17, EXPECTED_DECIMAL_17))) + .add(new TestColumn("t_array_decimal_18", + new ArrayType(DECIMAL_TYPE_18), + ImmutableList.of(WRITE_DECIMAL_18), + decimalArrayBlockOf(DECIMAL_TYPE_18, EXPECTED_DECIMAL_18))) + .add(new TestColumn("t_array_decimal_38", + new ArrayType(DECIMAL_TYPE_38), + ImmutableList.of(WRITE_DECIMAL_38), + decimalArrayBlockOf(DECIMAL_TYPE_38, EXPECTED_DECIMAL_38))) .add(new TestColumn("t_struct_bigint", - getStandardStructObjectInspector(ImmutableList.of("s_bigint"), ImmutableList.of(javaLongObjectInspector)), - new Long[] {1L}, + rowType(field("s_bigint", BIGINT)), + ImmutableList.of(1L), rowBlockOf(ImmutableList.of(BIGINT), 1))) .add(new TestColumn("t_complex", - getStandardMapObjectInspector( - javaStringObjectInspector, - getStandardListObjectInspector( - getStandardStructObjectInspector( - ImmutableList.of("s_int"), - ImmutableList.of(javaIntObjectInspector)))), - ImmutableMap.of("test", ImmutableList.of(new Integer[] {1})), + new MapType( + VARCHAR, + new ArrayType(rowType(field("s_int", INTEGER))), + TYPE_OPERATORS), + ImmutableMap.of("test", ImmutableList.of(ImmutableList.of(1))), mapBlockOf(createUnboundedVarcharType(), new ArrayType(RowType.anonymous(ImmutableList.of(INTEGER))), "test", arrayBlockOf(RowType.anonymous(ImmutableList.of(INTEGER)), rowBlockOf(ImmutableList.of(INTEGER), 1L))))) .add(new TestColumn("t_map_null_key_complex_value", - getStandardMapObjectInspector( - javaStringObjectInspector, - getStandardMapObjectInspector(javaLongObjectInspector, javaBooleanObjectInspector)), + new MapType( + VARCHAR, + new MapType(BIGINT, BOOLEAN, TYPE_OPERATORS), + TYPE_OPERATORS), asMap(new String[] {null, "k"}, new ImmutableMap[] {ImmutableMap.of(15L, true), ImmutableMap.of(16L, false)}), mapBlockOf(createUnboundedVarcharType(), mapType(BIGINT, BOOLEAN), "k", mapBlockOf(BIGINT, BOOLEAN, 16L, false)))) .add(new TestColumn("t_map_null_key_complex_key_value", - getStandardMapObjectInspector( - getStandardListObjectInspector(javaStringObjectInspector), - getStandardMapObjectInspector(javaLongObjectInspector, javaBooleanObjectInspector)), + new MapType( + new ArrayType(VARCHAR), + new MapType(BIGINT, BOOLEAN, TYPE_OPERATORS), + TYPE_OPERATORS), asMap(new ImmutableList[] {null, ImmutableList.of("k", "ka")}, new ImmutableMap[] {ImmutableMap.of(15L, true), ImmutableMap.of(16L, false)}), mapBlockOf(new ArrayType(createUnboundedVarcharType()), mapType(BIGINT, BOOLEAN), arrayBlockOf(createUnboundedVarcharType(), "k", "ka"), mapBlockOf(BIGINT, BOOLEAN, 16L, false)))) - .add(new TestColumn("t_struct_nested", getStandardStructObjectInspector(ImmutableList.of("struct_field"), - ImmutableList.of(getStandardListObjectInspector(javaStringObjectInspector))), ImmutableList.of(ImmutableList.of("1", "2", "3")), rowBlockOf(ImmutableList.of(new ArrayType(createUnboundedVarcharType())), arrayBlockOf(createUnboundedVarcharType(), "1", "2", "3")))) - .add(new TestColumn("t_struct_null", getStandardStructObjectInspector(ImmutableList.of("struct_field_null", "struct_field_null2"), - ImmutableList.of(javaStringObjectInspector, javaStringObjectInspector)), Arrays.asList(null, null), rowBlockOf(ImmutableList.of(createUnboundedVarcharType(), createUnboundedVarcharType()), null, null))) - .add(new TestColumn("t_struct_non_nulls_after_nulls", getStandardStructObjectInspector(ImmutableList.of("struct_non_nulls_after_nulls1", "struct_non_nulls_after_nulls2"), - ImmutableList.of(javaIntObjectInspector, javaStringObjectInspector)), Arrays.asList(null, "some string"), rowBlockOf(ImmutableList.of(INTEGER, createUnboundedVarcharType()), null, "some string"))) + .add(new TestColumn("t_struct_nested", + rowType(field("struct_field", new ArrayType(VARCHAR))), + ImmutableList.of(ImmutableList.of("1", "2", "3")), + rowBlockOf(ImmutableList.of(new ArrayType(createUnboundedVarcharType())), arrayBlockOf(createUnboundedVarcharType(), "1", "2", "3")))) + .add(new TestColumn("t_struct_null", + rowType(field("struct_field_null", VARCHAR), field("struct_field_null2", VARCHAR)), + Arrays.asList(null, null), + rowBlockOf(ImmutableList.of(createUnboundedVarcharType(), createUnboundedVarcharType()), null, null))) + .add(new TestColumn("t_struct_non_nulls_after_nulls", + rowType(field("struct_non_nulls_after_nulls1", INTEGER), field("struct_non_nulls_after_nulls2", VARCHAR)), + Arrays.asList(null, "some string"), + rowBlockOf(ImmutableList.of(INTEGER, createUnboundedVarcharType()), null, "some string"))) .add(new TestColumn("t_nested_struct_non_nulls_after_nulls", - getStandardStructObjectInspector( - ImmutableList.of("struct_field1", "struct_field2", "strict_field3"), - ImmutableList.of( - javaIntObjectInspector, - javaStringObjectInspector, - getStandardStructObjectInspector( - ImmutableList.of("nested_struct_field1", "nested_struct_field2"), - ImmutableList.of(javaIntObjectInspector, javaStringObjectInspector)))), + rowType( + field("struct_field1", INTEGER), + field("struct_field2", VARCHAR), + field("strict_field3", rowType(field("nested_struct_field1", INTEGER), field("nested_struct_field2", VARCHAR)))), Arrays.asList(null, "some string", Arrays.asList(null, "nested_string2")), rowBlockOf( ImmutableList.of( @@ -458,13 +503,13 @@ public abstract class AbstractTestHiveFileFormats RowType.anonymous(ImmutableList.of(INTEGER, createUnboundedVarcharType()))), null, "some string", rowBlockOf(ImmutableList.of(INTEGER, createUnboundedVarcharType()), null, "nested_string2")))) .add(new TestColumn("t_map_null_value", - getStandardMapObjectInspector(javaStringObjectInspector, javaStringObjectInspector), + new MapType(VARCHAR, VARCHAR, TYPE_OPERATORS), asMap(new String[] {"k1", "k2", "k3"}, new String[] {"v1", null, "v3"}), mapBlockOf(createUnboundedVarcharType(), createUnboundedVarcharType(), new String[] {"k1", "k2", "k3"}, new String[] {"v1", null, "v3"}))) - .add(new TestColumn("t_array_string_starting_with_nulls", getStandardListObjectInspector(javaStringObjectInspector), Arrays.asList(null, "test"), arrayBlockOf(createUnboundedVarcharType(), null, "test"))) - .add(new TestColumn("t_array_string_with_nulls_in_between", getStandardListObjectInspector(javaStringObjectInspector), Arrays.asList("test-1", null, "test-2"), arrayBlockOf(createUnboundedVarcharType(), "test-1", null, "test-2"))) - .add(new TestColumn("t_array_string_ending_with_nulls", getStandardListObjectInspector(javaStringObjectInspector), Arrays.asList("test", null), arrayBlockOf(createUnboundedVarcharType(), "test", null))) - .add(new TestColumn("t_array_string_all_nulls", getStandardListObjectInspector(javaStringObjectInspector), Arrays.asList(null, null, null), arrayBlockOf(createUnboundedVarcharType(), null, null, null))) + .add(new TestColumn("t_array_string_starting_with_nulls", new ArrayType(VARCHAR), Arrays.asList(null, "test"), arrayBlockOf(createUnboundedVarcharType(), null, "test"))) + .add(new TestColumn("t_array_string_with_nulls_in_between", new ArrayType(VARCHAR), Arrays.asList("test-1", null, "test-2"), arrayBlockOf(createUnboundedVarcharType(), "test-1", null, "test-2"))) + .add(new TestColumn("t_array_string_ending_with_nulls", new ArrayType(VARCHAR), Arrays.asList("test", null), arrayBlockOf(createUnboundedVarcharType(), "test", null))) + .add(new TestColumn("t_array_string_all_nulls", new ArrayType(VARCHAR), Arrays.asList(null, null, null), arrayBlockOf(createUnboundedVarcharType(), null, null, null))) .build(); private static Map asMap(K[] keys, V[] values) @@ -480,54 +525,25 @@ private static Map asMap(K[] keys, V[] values) protected List getColumnHandles(List testColumns) { - List columns = new ArrayList<>(); - Map hiveColumnIndexes = new HashMap<>(); + List columns = new ArrayList<>(testColumns.size()); int nextHiveColumnIndex = 0; - for (int i = 0; i < testColumns.size(); i++) { - TestColumn testColumn = testColumns.get(i); - + for (TestColumn testColumn : testColumns) { int columnIndex; - if (testColumn.isPartitionKey()) { + if (testColumn.partitionKey()) { columnIndex = -1; } else { - if (hiveColumnIndexes.get(testColumn.getBaseName()) != null) { - columnIndex = hiveColumnIndexes.get(testColumn.getBaseName()); - } - else { - columnIndex = nextHiveColumnIndex++; - hiveColumnIndexes.put(testColumn.getBaseName(), columnIndex); - } + columnIndex = nextHiveColumnIndex++; } - if (testColumn.getDereferenceNames().size() == 0) { - HiveType hiveType = HiveType.valueOf(testColumn.getObjectInspector().getTypeName()); - columns.add(createBaseColumn(testColumn.getName(), columnIndex, hiveType, hiveType.getType(TESTING_TYPE_MANAGER), testColumn.isPartitionKey() ? PARTITION_KEY : REGULAR, Optional.empty())); - } - else { - HiveType baseHiveType = HiveType.valueOf(testColumn.getBaseObjectInspector().getTypeName()); - HiveType partialHiveType = baseHiveType.getHiveTypeForDereferences(testColumn.getDereferenceIndices()).get(); - HiveColumnHandle hiveColumnHandle = new HiveColumnHandle( - testColumn.getBaseName(), - columnIndex, - baseHiveType, - baseHiveType.getType(TESTING_TYPE_MANAGER), - Optional.of(new HiveColumnProjectionInfo( - testColumn.getDereferenceIndices(), - testColumn.getDereferenceNames(), - partialHiveType, - partialHiveType.getType(TESTING_TYPE_MANAGER))), - testColumn.isPartitionKey() ? PARTITION_KEY : REGULAR, - Optional.empty()); - columns.add(hiveColumnHandle); - } + columns.add(testColumn.toHiveColumnHandle(columnIndex)); } return columns; } - public static FileSplit createTestFileTrino( - String filePath, + public static void createTestFileTrino( + Location location, HiveStorageFormat storageFormat, HiveCompressionCodec compressionCodec, List testColumns, @@ -537,13 +553,11 @@ public static FileSplit createTestFileTrino( { // filter out partition keys, which are not written to the file testColumns = testColumns.stream() - .filter(column -> !column.isPartitionKey()) + .filter(column -> !column.partitionKey()) .collect(toImmutableList()); List types = testColumns.stream() - .map(TestColumn::getType) - .map(HiveType::valueOf) - .map(type -> type.getType(TESTING_TYPE_MANAGER)) + .map(TestColumn::type) .collect(toList()); PageBuilder pageBuilder = new PageBuilder(types); @@ -551,35 +565,23 @@ public static FileSplit createTestFileTrino( for (int rowNumber = 0; rowNumber < numRows; rowNumber++) { pageBuilder.declarePosition(); for (int columnNumber = 0; columnNumber < testColumns.size(); columnNumber++) { - serializeObject( - types.get(columnNumber), - pageBuilder.getBlockBuilder(columnNumber), - testColumns.get(columnNumber).getWriteValue(), - testColumns.get(columnNumber).getObjectInspector(), - false); + TestColumn testColumn = testColumns.get(columnNumber); + writeValue(testColumn.type(), pageBuilder.getBlockBuilder(columnNumber), testColumn.writeValue()); } } Page page = pageBuilder.build(); - Properties tableProperties = new Properties(); - tableProperties.setProperty( - "columns", - testColumns.stream() - .map(TestColumn::getName) - .collect(Collectors.joining(","))); - - tableProperties.setProperty( - "columns.types", - testColumns.stream() - .map(TestColumn::getType) - .collect(Collectors.joining(","))); + Map tableProperties = ImmutableMap.builder() + .put(LIST_COLUMNS, testColumns.stream().map(TestColumn::name).collect(Collectors.joining(","))) + .put(LIST_COLUMN_TYPES, testColumns.stream().map(TestColumn::type).map(HiveTypeTranslator::toHiveType).map(HiveType::toString).collect(Collectors.joining(","))) + .buildOrThrow(); Optional fileWriter = fileWriterFactory.createFileWriter( - Location.of(filePath), + location, testColumns.stream() - .map(TestColumn::getName) + .map(TestColumn::name) .collect(toList()), - StorageFormat.fromHiveStorageFormat(storageFormat), + storageFormat.toStorageFormat(), compressionCodec, tableProperties, session, @@ -591,12 +593,11 @@ public static FileSplit createTestFileTrino( FileWriter hiveFileWriter = fileWriter.orElseThrow(() -> new IllegalArgumentException("fileWriterFactory")); hiveFileWriter.appendRows(page); hiveFileWriter.commit(); - - return new FileSplit(new Path(filePath), 0, new File(filePath).length(), new String[0]); } - public static FileSplit createTestFileHive( - String filePath, + public static void createTestFileHive( + TrinoFileSystemFactory fileSystemFactory, + Location location, HiveStorageFormat storageFormat, HiveCompressionCodec compressionCodec, List testColumns, @@ -608,42 +609,37 @@ public static FileSplit createTestFileHive( // filter out partition keys, which are not written to the file testColumns = testColumns.stream() - .filter(column -> !column.isPartitionKey()) + .filter(column -> !column.partitionKey()) .collect(toImmutableList()); Properties tableProperties = new Properties(); - tableProperties.setProperty( - "columns", - testColumns.stream() - .map(TestColumn::getName) - .collect(Collectors.joining(","))); - tableProperties.setProperty( - "columns.types", - testColumns.stream() - .map(TestColumn::getType) - .collect(Collectors.joining(","))); - serializer.initialize(newEmptyConfiguration(), tableProperties); + tableProperties.setProperty(LIST_COLUMNS, testColumns.stream().map(TestColumn::name).collect(Collectors.joining(","))); + tableProperties.setProperty(LIST_COLUMN_TYPES, testColumns.stream().map(testColumn -> toHiveType(testColumn.type()).toString()).collect(Collectors.joining(","))); + serializer.initialize(ConfigurationInstantiator.newEmptyConfiguration(), tableProperties); - JobConf jobConf = new JobConf(newEmptyConfiguration()); + JobConf jobConf = new JobConf(false); configureCompression(jobConf, compressionCodec); - RecordWriter recordWriter = outputFormat.getHiveRecordWriter( - jobConf, - new Path(filePath), - Text.class, - compressionCodec != HiveCompressionCodec.NONE, - tableProperties, - () -> {}); - + File file = File.createTempFile("trino_test", "data"); + verify(file.delete()); try { - serializer.initialize(newEmptyConfiguration(), tableProperties); + FileSinkOperator.RecordWriter recordWriter = outputFormat.getHiveRecordWriter( + jobConf, + new Path(file.getAbsolutePath()), + Text.class, + compressionCodec != HiveCompressionCodec.NONE, + tableProperties, + () -> {}); + + serializer.initialize(ConfigurationInstantiator.newEmptyConfiguration(), tableProperties); SettableStructObjectInspector objectInspector = getStandardStructObjectInspector( testColumns.stream() - .map(TestColumn::getName) + .map(TestColumn::name) .collect(toImmutableList()), testColumns.stream() - .map(TestColumn::getObjectInspector) + .map(TestColumn::type) + .map(AbstractTestHiveFileFormats::getJavaObjectInspector) .collect(toImmutableList())); Object row = objectInspector.create(); @@ -652,26 +648,189 @@ public static FileSplit createTestFileHive( for (int rowNumber = 0; rowNumber < numRows; rowNumber++) { for (int i = 0; i < testColumns.size(); i++) { - Object writeValue = testColumns.get(i).getWriteValue(); - if (writeValue instanceof Slice) { - writeValue = ((Slice) writeValue).getBytes(); - } - objectInspector.setStructFieldData(row, fields.get(i), writeValue); + objectInspector.setStructFieldData(row, fields.get(i), testColumns.get(i).writeValue()); } Writable record = serializer.serialize(row, objectInspector); recordWriter.write(record); } + recordWriter.close(false); + + // copy the file data to the TrinoFileSystem + TrinoFileSystem fileSystem = fileSystemFactory.create(ConnectorIdentity.ofUser("test")); + try (OutputStream outputStream = fileSystem.newOutputFile(location).create()) { + outputStream.write(readAllBytes(file.toPath())); + } } finally { - recordWriter.close(false); + verify(file.delete()); + } + } + + private static void writeValue(Type type, BlockBuilder builder, Object object) + { + requireNonNull(builder, "builder is null"); + + if (object == null) { + builder.appendNull(); + } + else if (type == BOOLEAN) { + BOOLEAN.writeBoolean(builder, (boolean) object); + } + else if (type == TINYINT) { + TINYINT.writeByte(builder, (byte) object); + } + else if (type == SMALLINT) { + SMALLINT.writeShort(builder, (short) object); + } + else if (type == INTEGER) { + INTEGER.writeInt(builder, (int) object); + } + else if (type == BIGINT) { + BIGINT.writeLong(builder, (long) object); + } + else if (type == REAL) { + REAL.writeFloat(builder, (float) object); + } + else if (type == DOUBLE) { + DOUBLE.writeDouble(builder, (double) object); + } + else if (type instanceof VarcharType varcharType) { + if (object instanceof HiveVarchar) { + object = ((HiveVarchar) object).getValue(); + } + varcharType.writeSlice(builder, utf8Slice((String) object)); + } + else if (type instanceof CharType charType) { + if (object instanceof HiveChar) { + object = ((HiveChar) object).getValue(); + } + charType.writeSlice(builder, truncateToLengthAndTrimSpaces(utf8Slice((String) object), charType.getLength())); + } + else if (type == DATE) { + long days = ((Date) object).toEpochDay(); + DATE.writeLong(builder, days); + } + else if (type instanceof TimestampType timestampType) { + Timestamp timestamp = (Timestamp) object; + long epochSecond = timestamp.toEpochSecond(); + int nanosOfSecond = (int) round(timestamp.getNanos(), 9 - timestampType.getPrecision()); + createTimestampEncoder(timestampType, UTC).write(new DecodedTimestamp(epochSecond, nanosOfSecond), builder); + } + else if (type == VARBINARY) { + VARBINARY.writeSlice(builder, Slices.wrappedBuffer((byte[]) object)); } + else if (type instanceof DecimalType decimalType) { + HiveDecimalWritable hiveDecimal = new HiveDecimalWritable((HiveDecimal) object); + Int128 value = Int128.fromBigEndian(hiveDecimal.getInternalStorage()); + value = Int128Math.rescale(value, decimalType.getScale() - hiveDecimal.getScale()); + if (decimalType.isShort()) { + type.writeLong(builder, value.toLongExact()); + } + else { + type.writeObject(builder, value); + } + } + else if (type instanceof ArrayType arrayType) { + Type elementType = arrayType.getElementType(); + List list = (List) object; + ((ArrayBlockBuilder) builder).buildEntry(elementBuilder -> { + for (Object element : list) { + writeValue(elementType, elementBuilder, element); + } + }); + } + else if (type instanceof MapType mapType) { + Type keyType = mapType.getKeyType(); + Type valueType = mapType.getValueType(); + Map map = (Map) object; + ((MapBlockBuilder) builder).buildEntry((keyBuilder, valueBuilder) -> { + for (Map.Entry entry : map.entrySet()) { + // Hive skips map entries with null keys + if (entry.getKey() != null) { + writeValue(keyType, keyBuilder, entry.getKey()); + writeValue(valueType, valueBuilder, entry.getValue()); + } + } + }); + } + else if (type instanceof RowType rowType) { + List typeParameters = rowType.getTypeParameters(); + List foo = (List) object; + ((RowBlockBuilder) builder).buildEntry(fieldBuilders -> { + for (int i = 0; i < typeParameters.size(); i++) { + writeValue(typeParameters.get(i), fieldBuilders.get(i), foo.get(i)); + } + }); + } + else { + throw new RuntimeException("Unsupported type: " + type); + } + } - // todo to test with compression, the file must be renamed with the compression extension - Path path = new Path(filePath); - path.getFileSystem(newEmptyConfiguration()).setVerifyChecksum(true); - File file = new File(filePath); - return new FileSplit(path, 0, file.length(), new String[0]); + private static ObjectInspector getJavaObjectInspector(Type type) + { + if (type.equals(BOOLEAN)) { + return javaBooleanObjectInspector; + } + if (type.equals(BIGINT)) { + return javaLongObjectInspector; + } + if (type.equals(INTEGER)) { + return javaIntObjectInspector; + } + if (type.equals(SMALLINT)) { + return javaShortObjectInspector; + } + if (type.equals(TINYINT)) { + return javaByteObjectInspector; + } + if (type.equals(REAL)) { + return javaFloatObjectInspector; + } + if (type.equals(DOUBLE)) { + return javaDoubleObjectInspector; + } + if (type instanceof VarcharType varcharType) { + return varcharType.getLength() + .map(length -> getPrimitiveJavaObjectInspector(new VarcharTypeInfo(length))) + .orElse(javaStringObjectInspector); + } + if (type instanceof CharType charType) { + return getPrimitiveJavaObjectInspector(new CharTypeInfo(charType.getLength())); + } + if (type.equals(VARBINARY)) { + return javaByteArrayObjectInspector; + } + if (type.equals(DATE)) { + return javaDateObjectInspector; + } + if (type instanceof TimestampType) { + return javaTimestampObjectInspector; + } + if (type instanceof DecimalType decimalType) { + return getPrimitiveJavaObjectInspector(new DecimalTypeInfo(decimalType.getPrecision(), decimalType.getScale())); + } + if (type instanceof ArrayType arrayType) { + return getStandardListObjectInspector(getJavaObjectInspector(arrayType.getElementType())); + } + if (type instanceof MapType mapType) { + ObjectInspector keyObjectInspector = getJavaObjectInspector(mapType.getKeyType()); + ObjectInspector valueObjectInspector = getJavaObjectInspector(mapType.getValueType()); + return getStandardMapObjectInspector(keyObjectInspector, valueObjectInspector); + } + if (type instanceof RowType rowType) { + return getStandardStructObjectInspector( + rowType.getFields().stream() + .map(RowType.Field::getName) + .map(Optional::orElseThrow) + .collect(toList()), + rowType.getFields().stream() + .map(RowType.Field::getType) + .map(AbstractTestHiveFileFormats::getJavaObjectInspector) + .collect(toList())); + } + throw new IllegalArgumentException("unsupported type: " + type); } private static T newInstance(String className, Class superType) @@ -727,241 +886,163 @@ public static Object getFieldFromCursor(RecordCursor cursor, Type type, int fiel throw new RuntimeException("unknown type"); } - protected void checkCursor(RecordCursor cursor, List testColumns, int rowCount) - { - List types = testColumns.stream() - .map(column -> column.getObjectInspector().getTypeName()) - .map(type -> HiveType.valueOf(type).getType(TESTING_TYPE_MANAGER)) - .collect(toImmutableList()); - - Map distinctFromOperators = types.stream().distinct() - .collect(toImmutableMap(identity(), HiveTestUtils::distinctFromOperator)); - - for (int row = 0; row < rowCount; row++) { - assertTrue(cursor.advanceNextPosition()); - for (int i = 0, testColumnsSize = testColumns.size(); i < testColumnsSize; i++) { - TestColumn testColumn = testColumns.get(i); - - Type type = types.get(i); - Object fieldFromCursor = getFieldFromCursor(cursor, type, i); - if (fieldFromCursor == null) { - assertEquals(null, testColumn.getExpectedValue(), "Expected null for column " + testColumn.getName()); - } - else if (type instanceof DecimalType decimalType) { - fieldFromCursor = new BigDecimal((BigInteger) fieldFromCursor, decimalType.getScale()); - assertEquals(fieldFromCursor, testColumn.getExpectedValue(), "Wrong value for column " + testColumn.getName()); - } - else if (testColumn.getObjectInspector().getTypeName().equals("float")) { - assertEquals((float) fieldFromCursor, (float) testColumn.getExpectedValue(), (float) EPSILON); - } - else if (testColumn.getObjectInspector().getTypeName().equals("double")) { - assertEquals((double) fieldFromCursor, (double) testColumn.getExpectedValue(), EPSILON); - } - else if (testColumn.getObjectInspector().getTypeName().equals("tinyint")) { - assertEquals(((Number) fieldFromCursor).byteValue(), testColumn.getExpectedValue()); - } - else if (testColumn.getObjectInspector().getTypeName().equals("smallint")) { - assertEquals(((Number) fieldFromCursor).shortValue(), testColumn.getExpectedValue()); - } - else if (testColumn.getObjectInspector().getTypeName().equals("int")) { - assertEquals(((Number) fieldFromCursor).intValue(), testColumn.getExpectedValue()); - } - else if (testColumn.getObjectInspector().getCategory() == Category.PRIMITIVE) { - assertEquals(fieldFromCursor, testColumn.getExpectedValue(), "Wrong value for column " + testColumn.getName()); - } - else { - Block expected = (Block) testColumn.getExpectedValue(); - Block actual = (Block) fieldFromCursor; - boolean distinct = isDistinctFrom(distinctFromOperators.get(type), expected, actual); - assertFalse(distinct, "Wrong value for column: " + testColumn.getName()); - } - } - } - assertFalse(cursor.advanceNextPosition()); - } - - protected void checkPageSource(ConnectorPageSource pageSource, List testColumns, List types, int rowCount) + protected void checkPageSource(ConnectorPageSource pageSource, List testColumns, int rowCount) throws IOException { - try { - MaterializedResult result = materializeSourceDataStream(SESSION, pageSource, types); - assertEquals(result.getMaterializedRows().size(), rowCount); + try (pageSource) { + MaterializedResult result = materializeSourceDataStream(SESSION, pageSource, testColumns.stream().map(TestColumn::type).collect(toImmutableList())); + assertThat(result.getMaterializedRows()).hasSize(rowCount); for (MaterializedRow row : result) { for (int i = 0, testColumnsSize = testColumns.size(); i < testColumnsSize; i++) { TestColumn testColumn = testColumns.get(i); - Type type = types.get(i); + Type type = testColumn.type(); Object actualValue = row.getField(i); - Object expectedValue = testColumn.getExpectedValue(); + Object expectedValue = testColumn.expectedValue(); if (expectedValue instanceof Slice) { expectedValue = ((Slice) expectedValue).toStringUtf8(); } if (actualValue == null || expectedValue == null) { - assertEquals(actualValue, expectedValue, "Wrong value for column " + testColumn.getName()); + assertThat(actualValue) + .describedAs("Wrong value for column " + testColumn.name()) + .isEqualTo(expectedValue); + } + else if (type == REAL) { + assertThat((float) actualValue).describedAs("Wrong value for column %s", testColumn.name()) + .isCloseTo((float) expectedValue, offset((float) EPSILON)); } - else if (testColumn.getObjectInspector().getTypeName().equals("float")) { - assertEquals((float) actualValue, (float) expectedValue, EPSILON, "Wrong value for column " + testColumn.getName()); + else if (type == DOUBLE) { + assertThat((double) actualValue).describedAs("Wrong value for column %s", testColumn.name()) + .isCloseTo((double) expectedValue, offset((double) EPSILON)); } - else if (testColumn.getObjectInspector().getTypeName().equals("double")) { - assertEquals((double) actualValue, (double) expectedValue, EPSILON, "Wrong value for column " + testColumn.getName()); + else if (type == DATE) { + SqlDate expectedDate = new SqlDate(toIntExact((long) expectedValue)); + assertThat(actualValue) + .describedAs("Wrong value for column " + testColumn.name()) + .isEqualTo(expectedDate); } - else if (testColumn.getObjectInspector().getTypeName().equals("date")) { - SqlDate expectedDate = new SqlDate(((Long) expectedValue).intValue()); - assertEquals(actualValue, expectedDate, "Wrong value for column " + testColumn.getName()); + else if (type == BIGINT || type == INTEGER || type == SMALLINT || type == TINYINT || type == BOOLEAN) { + assertThat(actualValue).isEqualTo(expectedValue); } - else if (testColumn.getObjectInspector().getTypeName().equals("int") || - testColumn.getObjectInspector().getTypeName().equals("smallint") || - testColumn.getObjectInspector().getTypeName().equals("tinyint")) { - assertEquals(actualValue, expectedValue); + else if (type instanceof TimestampType timestampType && timestampType.getPrecision() == 3) { + // the expected value is in micros to simplify the array, map, and row types + SqlTimestamp expectedTimestamp = sqlTimestampOf(3, floorDiv((long) expectedValue, MICROSECONDS_PER_MILLISECOND)); + assertThat(actualValue) + .describedAs("Wrong value for column " + testColumn.name()) + .isEqualTo(expectedTimestamp); } - else if (testColumn.getObjectInspector().getTypeName().equals("timestamp")) { - SqlTimestamp expectedTimestamp = sqlTimestampOf(floorDiv((Long) expectedValue, MICROSECONDS_PER_MILLISECOND)); - assertEquals(actualValue, expectedTimestamp, "Wrong value for column " + testColumn.getName()); + else if (type instanceof CharType) { + assertThat(actualValue) + .describedAs("Wrong value for column " + testColumn.name()) + .isEqualTo(padSpaces((String) expectedValue, (CharType) type)); } - else if (testColumn.getObjectInspector().getTypeName().startsWith("char")) { - assertEquals(actualValue, padSpaces((String) expectedValue, (CharType) type), "Wrong value for column " + testColumn.getName()); + else if (type instanceof VarcharType) { + assertThat(actualValue) + .describedAs("Wrong value for column " + testColumn.name()) + .isEqualTo(expectedValue); } - else if (testColumn.getObjectInspector().getCategory() == Category.PRIMITIVE) { - if (expectedValue instanceof Slice) { - expectedValue = ((Slice) expectedValue).toStringUtf8(); - } - - if (actualValue instanceof Slice) { - actualValue = ((Slice) actualValue).toStringUtf8(); - } - if (actualValue instanceof SqlVarbinary) { - actualValue = new String(((SqlVarbinary) actualValue).getBytes(), UTF_8); - } - - if (actualValue instanceof SqlDecimal) { - actualValue = new BigDecimal(actualValue.toString()); - } - assertEquals(actualValue, expectedValue, "Wrong value for column " + testColumn.getName()); + else if (type == VARBINARY) { + assertThat(new String(((SqlVarbinary) actualValue).getBytes(), UTF_8)) + .describedAs("Wrong value for column " + testColumn.name()) + .isEqualTo(expectedValue); + } + else if (type instanceof DecimalType) { + assertThat(new BigDecimal(actualValue.toString())) + .describedAs("Wrong value for column " + testColumn.name()) + .isEqualTo(expectedValue); } else { BlockBuilder builder = type.createBlockBuilder(null, 1); type.writeObject(builder, expectedValue); expectedValue = type.getObjectValue(SESSION, builder.build(), 0); - assertEquals(actualValue, expectedValue, "Wrong value for column " + testColumn.getName()); + assertThat(actualValue) + .describedAs("Wrong value for column " + testColumn.name()) + .isEqualTo(expectedValue); } } } } - finally { - pageSource.close(); - } } - public static final class TestColumn + public record TestColumn( + String name, + Type type, + String baseName, + Type baseType, + boolean dereference, + Object writeValue, + Object expectedValue, + boolean partitionKey) { - private final String baseName; - private final ObjectInspector baseObjectInspector; - private final List dereferenceNames; - private final List dereferenceIndices; - private final String name; - private final ObjectInspector objectInspector; - private final Object writeValue; - private final Object expectedValue; - private final boolean partitionKey; - - public TestColumn(String name, ObjectInspector objectInspector, Object writeValue, Object expectedValue) - { - this(name, objectInspector, writeValue, expectedValue, false); - } - - public TestColumn(String name, ObjectInspector objectInspector, Object writeValue, Object expectedValue, boolean partitionKey) - { - this(name, objectInspector, ImmutableList.of(), ImmutableList.of(), objectInspector, writeValue, expectedValue, partitionKey); - } - - public TestColumn( - String baseName, - ObjectInspector baseObjectInspector, - List dereferenceNames, - List dereferenceIndices, - ObjectInspector objectInspector, - Object writeValue, - Object expectedValue, - boolean partitionKey) + public TestColumn(String name, Type type, Object writeValue, Object expectedValue) { - this.baseName = requireNonNull(baseName, "baseName is null"); - this.baseObjectInspector = requireNonNull(baseObjectInspector, "baseObjectInspector is null"); - this.dereferenceNames = requireNonNull(dereferenceNames, "dereferenceNames is null"); - this.dereferenceIndices = requireNonNull(dereferenceIndices, "dereferenceIndices is null"); - checkArgument(dereferenceIndices.size() == dereferenceNames.size(), "dereferenceIndices and dereferenceNames should have the same size"); - this.name = baseName + generatePartialName(dereferenceNames); - this.objectInspector = requireNonNull(objectInspector, "objectInspector is null"); - this.writeValue = writeValue; - this.expectedValue = expectedValue; - this.partitionKey = partitionKey; - checkArgument(dereferenceNames.size() == 0 || partitionKey == false, "partial column cannot be a partition key"); - } - - public String getName() - { - return name; - } - - public String getBaseName() - { - return baseName; + this(name, type, writeValue, expectedValue, false); } - public List getDereferenceNames() + public TestColumn(String name, Type type, Object writeValue, Object expectedValue, boolean partitionKey) { - return dereferenceNames; + this(name, type, name, type, false, writeValue, expectedValue, partitionKey); } - public List getDereferenceIndices() + public TestColumn { - return dereferenceIndices; + requireNonNull(name, "name is null"); + requireNonNull(type, "type is null"); + requireNonNull(baseName, "baseName is null"); + requireNonNull(baseType, "baseType is null"); } - public String getType() + public HiveColumnHandle toHiveColumnHandle(int columnIndex) { - return objectInspector.getTypeName(); - } - - public ObjectInspector getBaseObjectInspector() - { - return baseObjectInspector; - } + checkArgument(partitionKey == (columnIndex == -1)); - public ObjectInspector getObjectInspector() - { - return objectInspector; - } + if (!dereference) { + return createBaseColumn(name, columnIndex, toHiveType(type), type, partitionKey ? PARTITION_KEY : REGULAR, Optional.empty()); + } - public Object getWriteValue() - { - return writeValue; + return new HiveColumnHandle( + baseName, + columnIndex, + toHiveType(baseType), + baseType, + Optional.of(new HiveColumnProjectionInfo(ImmutableList.of(0), ImmutableList.of(name), toHiveType(type), type)), + partitionKey ? PARTITION_KEY : REGULAR, + Optional.empty()); } - public Object getExpectedValue() + public TestColumn withDereferenceFirstField(Object writeValue, Object expectedValue) { - return expectedValue; - } + verify(!partitionKey, "dereference not supported for partition key"); + verify(!dereference, "already dereference"); + if (!(type instanceof RowType rowType)) { + throw new VerifyException("type is not a row type"); + } - public boolean isPartitionKey() - { - return partitionKey; + RowType.Field field = rowType.getFields().get(0); + return new TestColumn( + field.getName().orElseThrow(), + field.getType(), + name, + type, + true, + writeValue, + expectedValue, + false); } - @Override - public String toString() + public TestColumn withName(String newName) { - StringBuilder sb = new StringBuilder("TestColumn{"); - sb.append("baseName='").append(baseName).append("'"); - sb.append("dereferenceNames=").append("[").append(dereferenceNames.stream().collect(Collectors.joining(","))).append("]"); - sb.append("name=").append(name); - sb.append(", objectInspector=").append(objectInspector); - sb.append(", writeValue=").append(writeValue); - sb.append(", expectedValue=").append(expectedValue); - sb.append(", partitionKey=").append(partitionKey); - sb.append('}'); - return sb.toString(); + return new TestColumn( + newName, + type, + baseName, + baseType, + dereference, + writeValue, + expectedValue, + partitionKey); } } } diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/AbstractTestHiveFileSystem.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/AbstractTestHiveFileSystem.java deleted file mode 100644 index e22c4e290c3d..000000000000 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/AbstractTestHiveFileSystem.java +++ /dev/null @@ -1,848 +0,0 @@ -/* - * 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.hive; - -import com.google.common.collect.ImmutableList; -import com.google.common.collect.ImmutableMap; -import com.google.common.collect.ImmutableSet; -import com.google.common.collect.Streams; -import com.google.common.net.HostAndPort; -import io.airlift.concurrent.BoundedExecutor; -import io.airlift.json.JsonCodec; -import io.airlift.slice.Slice; -import io.airlift.stats.CounterStat; -import io.trino.filesystem.Location; -import io.trino.filesystem.TrinoFileSystem; -import io.trino.filesystem.hdfs.HdfsFileSystemFactory; -import io.trino.hdfs.HdfsConfig; -import io.trino.hdfs.HdfsConfiguration; -import io.trino.hdfs.HdfsContext; -import io.trino.hdfs.HdfsEnvironment; -import io.trino.hdfs.HdfsNamenodeStats; -import io.trino.hdfs.TrinoHdfsFileSystemStats; -import io.trino.hdfs.authentication.NoHdfsAuthentication; -import io.trino.operator.GroupByHashPageIndexerFactory; -import io.trino.plugin.base.CatalogName; -import io.trino.plugin.hive.AbstractTestHive.Transaction; -import io.trino.plugin.hive.aws.athena.PartitionProjectionService; -import io.trino.plugin.hive.fs.FileSystemDirectoryLister; -import io.trino.plugin.hive.fs.HiveFileIterator; -import io.trino.plugin.hive.fs.TransactionScopeCachingDirectoryListerFactory; -import io.trino.plugin.hive.fs.TrinoFileStatus; -import io.trino.plugin.hive.metastore.Column; -import io.trino.plugin.hive.metastore.Database; -import io.trino.plugin.hive.metastore.ForwardingHiveMetastore; -import io.trino.plugin.hive.metastore.HiveMetastore; -import io.trino.plugin.hive.metastore.HiveMetastoreConfig; -import io.trino.plugin.hive.metastore.HiveMetastoreFactory; -import io.trino.plugin.hive.metastore.PrincipalPrivileges; -import io.trino.plugin.hive.metastore.StorageFormat; -import io.trino.plugin.hive.metastore.Table; -import io.trino.plugin.hive.metastore.thrift.BridgingHiveMetastore; -import io.trino.plugin.hive.security.SqlStandardAccessControlMetadata; -import io.trino.spi.connector.ColumnHandle; -import io.trino.spi.connector.ColumnMetadata; -import io.trino.spi.connector.ConnectorMetadata; -import io.trino.spi.connector.ConnectorOutputTableHandle; -import io.trino.spi.connector.ConnectorPageSink; -import io.trino.spi.connector.ConnectorPageSinkProvider; -import io.trino.spi.connector.ConnectorPageSource; -import io.trino.spi.connector.ConnectorPageSourceProvider; -import io.trino.spi.connector.ConnectorSession; -import io.trino.spi.connector.ConnectorSplit; -import io.trino.spi.connector.ConnectorSplitManager; -import io.trino.spi.connector.ConnectorSplitSource; -import io.trino.spi.connector.ConnectorTableHandle; -import io.trino.spi.connector.ConnectorTableMetadata; -import io.trino.spi.connector.DynamicFilter; -import io.trino.spi.connector.SchemaTableName; -import io.trino.spi.connector.TableNotFoundException; -import io.trino.spi.predicate.TupleDomain; -import io.trino.spi.security.ConnectorIdentity; -import io.trino.spi.type.TestingTypeManager; -import io.trino.spi.type.TypeOperators; -import io.trino.sql.gen.JoinCompiler; -import io.trino.testing.MaterializedResult; -import io.trino.testing.TestingNodeManager; -import io.trino.type.BlockTypeOperators; -import org.apache.hadoop.fs.FileSystem; -import org.apache.hadoop.fs.Path; -import org.apache.hadoop.fs.azurebfs.AzureBlobFileSystem; -import org.testng.annotations.AfterClass; -import org.testng.annotations.BeforeClass; -import org.testng.annotations.Test; - -import java.io.BufferedReader; -import java.io.IOException; -import java.io.InputStream; -import java.io.InputStreamReader; -import java.io.OutputStream; -import java.io.UncheckedIOException; -import java.util.Collection; -import java.util.List; -import java.util.Optional; -import java.util.concurrent.ExecutorService; -import java.util.concurrent.ScheduledExecutorService; - -import static com.google.common.collect.ImmutableList.toImmutableList; -import static com.google.common.collect.Iterables.getOnlyElement; -import static com.google.common.util.concurrent.MoreExecutors.newDirectExecutorService; -import static io.airlift.concurrent.MoreFutures.getFutureValue; -import static io.airlift.concurrent.Threads.daemonThreadsNamed; -import static io.trino.hdfs.FileSystemUtils.getRawFileSystem; -import static io.trino.plugin.hive.AbstractTestHive.createTableProperties; -import static io.trino.plugin.hive.AbstractTestHive.filterNonHiddenColumnHandles; -import static io.trino.plugin.hive.AbstractTestHive.filterNonHiddenColumnMetadata; -import static io.trino.plugin.hive.AbstractTestHive.getAllSplits; -import static io.trino.plugin.hive.AbstractTestHive.getSplits; -import static io.trino.plugin.hive.HiveTestUtils.HDFS_FILE_SYSTEM_STATS; -import static io.trino.plugin.hive.HiveTestUtils.PAGE_SORTER; -import static io.trino.plugin.hive.HiveTestUtils.SESSION; -import static io.trino.plugin.hive.HiveTestUtils.getDefaultHiveFileWriterFactories; -import static io.trino.plugin.hive.HiveTestUtils.getDefaultHivePageSourceFactories; -import static io.trino.plugin.hive.HiveTestUtils.getDefaultHiveRecordCursorProviders; -import static io.trino.plugin.hive.HiveTestUtils.getHiveSessionProperties; -import static io.trino.plugin.hive.HiveTestUtils.getTypes; -import static io.trino.plugin.hive.HiveType.HIVE_LONG; -import static io.trino.plugin.hive.HiveType.HIVE_STRING; -import static io.trino.plugin.hive.TestingThriftHiveMetastoreBuilder.testingThriftHiveMetastoreBuilder; -import static io.trino.plugin.hive.metastore.PrincipalPrivileges.NO_PRIVILEGES; -import static io.trino.spi.connector.MetadataProvider.NOOP_METADATA_PROVIDER; -import static io.trino.spi.connector.RetryMode.NO_RETRIES; -import static io.trino.spi.type.BigintType.BIGINT; -import static io.trino.testing.MaterializedResult.materializeSourceDataStream; -import static io.trino.testing.QueryAssertions.assertEqualsIgnoreOrder; -import static io.trino.testing.TestingPageSinkId.TESTING_PAGE_SINK_ID; -import static io.trino.type.InternalTypeManager.TESTING_TYPE_MANAGER; -import static java.nio.charset.StandardCharsets.UTF_8; -import static java.util.Locale.ENGLISH; -import static java.util.UUID.randomUUID; -import static java.util.concurrent.Executors.newCachedThreadPool; -import static java.util.concurrent.Executors.newScheduledThreadPool; -import static org.assertj.core.api.Assertions.assertThat; -import static org.testng.Assert.assertEquals; -import static org.testng.Assert.assertFalse; -import static org.testng.Assert.assertTrue; - -public abstract class AbstractTestHiveFileSystem -{ - protected static final HdfsContext TESTING_CONTEXT = new HdfsContext(ConnectorIdentity.ofUser("test")); - - protected String database; - protected SchemaTableName table; - protected SchemaTableName tableWithHeader; - protected SchemaTableName tableWithHeaderAndFooter; - protected SchemaTableName temporaryCreateTable; - - protected HdfsEnvironment hdfsEnvironment; - protected LocationService locationService; - protected TestingHiveMetastore metastoreClient; - protected HiveMetadataFactory metadataFactory; - protected HiveTransactionManager transactionManager; - protected ConnectorSplitManager splitManager; - protected ConnectorPageSinkProvider pageSinkProvider; - protected ConnectorPageSourceProvider pageSourceProvider; - - private ExecutorService executor; - private HiveConfig config; - private ScheduledExecutorService heartbeatService; - - @BeforeClass - public void setUp() - { - executor = newCachedThreadPool(daemonThreadsNamed("hive-%s")); - heartbeatService = newScheduledThreadPool(1); - } - - @AfterClass(alwaysRun = true) - public void tearDown() - { - if (executor != null) { - executor.shutdownNow(); - executor = null; - } - if (heartbeatService != null) { - heartbeatService.shutdownNow(); - heartbeatService = null; - } - } - - protected abstract Path getBasePath(); - - protected void onSetupComplete() {} - - protected void setup(String host, int port, String databaseName, boolean s3SelectPushdownEnabled, HdfsConfiguration hdfsConfiguration) - { - database = databaseName; - table = new SchemaTableName(database, "trino_test_external_fs"); - tableWithHeader = new SchemaTableName(database, "trino_test_external_fs_with_header"); - tableWithHeaderAndFooter = new SchemaTableName(database, "trino_test_external_fs_with_header_and_footer"); - - String random = randomUUID().toString().toLowerCase(ENGLISH).replace("-", ""); - temporaryCreateTable = new SchemaTableName(database, "tmp_trino_test_create_" + random); - - config = new HiveConfig().setS3SelectPushdownEnabled(s3SelectPushdownEnabled); - - HivePartitionManager hivePartitionManager = new HivePartitionManager(config); - - hdfsEnvironment = new HdfsEnvironment(hdfsConfiguration, new HdfsConfig(), new NoHdfsAuthentication()); - metastoreClient = new TestingHiveMetastore( - new BridgingHiveMetastore( - testingThriftHiveMetastoreBuilder() - .metastoreClient(HostAndPort.fromParts(host, port)) - .hiveConfig(config) - .hdfsEnvironment(hdfsEnvironment) - .build()), - getBasePath(), - hdfsEnvironment); - locationService = new HiveLocationService(hdfsEnvironment, config); - JsonCodec partitionUpdateCodec = JsonCodec.jsonCodec(PartitionUpdate.class); - metadataFactory = new HiveMetadataFactory( - new CatalogName("hive"), - config, - new HiveMetastoreConfig(), - HiveMetastoreFactory.ofInstance(metastoreClient), - getDefaultHiveFileWriterFactories(config, hdfsEnvironment), - new HdfsFileSystemFactory(hdfsEnvironment, HDFS_FILE_SYSTEM_STATS), - hdfsEnvironment, - hivePartitionManager, - newDirectExecutorService(), - heartbeatService, - TESTING_TYPE_MANAGER, - NOOP_METADATA_PROVIDER, - locationService, - partitionUpdateCodec, - new NodeVersion("test_version"), - new NoneHiveRedirectionsProvider(), - ImmutableSet.of( - new PartitionsSystemTableProvider(hivePartitionManager, TESTING_TYPE_MANAGER), - new PropertiesSystemTableProvider()), - new DefaultHiveMaterializedViewMetadataFactory(), - SqlStandardAccessControlMetadata::new, - new FileSystemDirectoryLister(), - new TransactionScopeCachingDirectoryListerFactory(config), - new PartitionProjectionService(config, ImmutableMap.of(), new TestingTypeManager()), - true); - transactionManager = new HiveTransactionManager(metadataFactory); - splitManager = new HiveSplitManager( - transactionManager, - hivePartitionManager, - new HdfsFileSystemFactory(hdfsEnvironment, HDFS_FILE_SYSTEM_STATS), - new HdfsNamenodeStats(), - hdfsEnvironment, - new BoundedExecutor(executor, config.getMaxSplitIteratorThreads()), - new CounterStat(), - config.getMaxOutstandingSplits(), - config.getMaxOutstandingSplitsSize(), - config.getMinPartitionBatchSize(), - config.getMaxPartitionBatchSize(), - config.getMaxInitialSplits(), - config.getSplitLoaderConcurrency(), - config.getMaxSplitsPerSecond(), - config.getRecursiveDirWalkerEnabled(), - TESTING_TYPE_MANAGER, - config.getMaxPartitionsPerScan()); - TypeOperators typeOperators = new TypeOperators(); - BlockTypeOperators blockTypeOperators = new BlockTypeOperators(typeOperators); - pageSinkProvider = new HivePageSinkProvider( - getDefaultHiveFileWriterFactories(config, hdfsEnvironment), - new HdfsFileSystemFactory(hdfsEnvironment, HDFS_FILE_SYSTEM_STATS), - hdfsEnvironment, - PAGE_SORTER, - HiveMetastoreFactory.ofInstance(metastoreClient), - new GroupByHashPageIndexerFactory(new JoinCompiler(typeOperators), blockTypeOperators), - TESTING_TYPE_MANAGER, - config, - new SortingFileWriterConfig(), - locationService, - partitionUpdateCodec, - new TestingNodeManager("fake-environment"), - new HiveEventClient(), - getHiveSessionProperties(config), - new HiveWriterStats()); - pageSourceProvider = new HivePageSourceProvider( - TESTING_TYPE_MANAGER, - hdfsEnvironment, - config, - getDefaultHivePageSourceFactories(hdfsEnvironment, config), - getDefaultHiveRecordCursorProviders(config, hdfsEnvironment), - new GenericHiveRecordCursorProvider(hdfsEnvironment, config)); - - onSetupComplete(); - } - - protected ConnectorSession newSession() - { - return HiveFileSystemTestUtils.newSession(config); - } - - protected Transaction newTransaction() - { - return HiveFileSystemTestUtils.newTransaction(transactionManager); - } - - protected MaterializedResult readTable(SchemaTableName tableName) - throws IOException - { - return HiveFileSystemTestUtils.readTable(tableName, transactionManager, config, pageSourceProvider, splitManager); - } - - protected MaterializedResult filterTable(SchemaTableName tableName, List projectedColumns) - throws IOException - { - return HiveFileSystemTestUtils.filterTable(tableName, projectedColumns, transactionManager, config, pageSourceProvider, splitManager); - } - - @Test - public void testGetRecords() - throws Exception - { - assertEqualsIgnoreOrder( - readTable(table), - MaterializedResult.resultBuilder(newSession(), BIGINT) - .row(3L).row(14L).row(15L) // test_table.csv - .row(92L).row(65L).row(35L) // test_table.csv.gz - .row(89L).row(79L).row(32L) // test_table.csv.bz2 - .row(38L).row(46L).row(26L) // test_table.csv.lz4 - .build()); - } - - @Test - public void testGetRecordsWithHeader() - throws IOException - { - assertEqualsIgnoreOrder( - readTable(tableWithHeader), - MaterializedResult.resultBuilder(newSession(), BIGINT) - .row(2L).row(71L).row(82L) // test_table_with_header.csv - .row(81L).row(82L).row(84L) // test_table_with_header.csv.gz - .row(59L).row(4L).row(52L) // test_table_with_header.csv.bz2 - .row(35L).row(36L).row(2L) // test_table_with_header.csv.lz4 - .build()); - } - - @Test - public void testGetRecordsWithHeaderAndFooter() - throws IOException - { - assertEqualsIgnoreOrder( - readTable(tableWithHeaderAndFooter), - MaterializedResult.resultBuilder(newSession(), BIGINT) - .row(1L).row(41L).row(42L) // test_table_with_header_and_footer.csv - .row(13L).row(56L).row(23L) // test_table_with_header_and_footer.csv.gz - .row(73L).row(9L).row(50L) // test_table_with_header_and_footer.csv.bz2 - .row(48L).row(80L).row(16L) // test_table_with_header_and_footer.csv.lz4 - .build()); - } - - @Test - public void testGetFileStatus() - throws Exception - { - Path basePath = getBasePath(); - Path tablePath = new Path(basePath, "trino_test_external_fs"); - Path filePath = new Path(tablePath, "test_table.csv"); - FileSystem fs = hdfsEnvironment.getFileSystem(TESTING_CONTEXT, basePath); - - assertTrue(fs.getFileStatus(basePath).isDirectory(), "basePath should be considered a directory"); - assertTrue(fs.getFileStatus(tablePath).isDirectory(), "tablePath should be considered a directory"); - assertTrue(fs.getFileStatus(filePath).isFile(), "filePath should be considered a file"); - assertFalse(fs.getFileStatus(filePath).isDirectory(), "filePath should not be considered a directory"); - assertFalse(fs.exists(new Path(basePath, "foo-" + randomUUID())), "foo-random path should be found not to exist"); - assertFalse(fs.exists(new Path(basePath, "foo")), "foo path should be found not to exist"); - } - - @Test - public void testRename() - throws Exception - { - Path basePath = new Path(getBasePath(), randomUUID().toString()); - FileSystem fs = hdfsEnvironment.getFileSystem(TESTING_CONTEXT, basePath); - assertFalse(fs.exists(basePath)); - - // create file foo.txt - Path path = new Path(basePath, "foo.txt"); - assertTrue(fs.createNewFile(path)); - assertTrue(fs.exists(path)); - - // rename foo.txt to bar.txt when bar does not exist - Path newPath = new Path(basePath, "bar.txt"); - assertFalse(fs.exists(newPath)); - assertTrue(fs.rename(path, newPath)); - assertFalse(fs.exists(path)); - assertTrue(fs.exists(newPath)); - - // rename foo.txt to foo.txt when foo.txt does not exist - assertFalse(fs.rename(path, path)); - - // create file foo.txt and rename to existing bar.txt - assertTrue(fs.createNewFile(path)); - assertFalse(fs.rename(path, newPath)); - - // rename foo.txt to foo.txt when foo.txt exists - assertEquals(fs.rename(path, path), getRawFileSystem(fs) instanceof AzureBlobFileSystem); - - // delete foo.txt - assertTrue(fs.delete(path, false)); - assertFalse(fs.exists(path)); - - // create directory source with file - Path source = new Path(basePath, "source"); - assertTrue(fs.createNewFile(new Path(source, "test.txt"))); - - // rename source to non-existing target - Path target = new Path(basePath, "target"); - assertFalse(fs.exists(target)); - assertTrue(fs.rename(source, target)); - assertFalse(fs.exists(source)); - assertTrue(fs.exists(target)); - - // create directory source with file - assertTrue(fs.createNewFile(new Path(source, "test.txt"))); - - // rename source to existing target - assertTrue(fs.rename(source, target)); - assertFalse(fs.exists(source)); - target = new Path(target, "source"); - assertTrue(fs.exists(target)); - assertTrue(fs.exists(new Path(target, "test.txt"))); - - // delete target - target = new Path(basePath, "target"); - assertTrue(fs.exists(target)); - assertTrue(fs.delete(target, true)); - assertFalse(fs.exists(target)); - - // cleanup - fs.delete(basePath, true); - } - - @Test - public void testFileIteratorListing() - throws Exception - { - Table.Builder tableBuilder = Table.builder() - .setDatabaseName(table.getSchemaName()) - .setTableName(table.getTableName()) - .setDataColumns(ImmutableList.of(new Column("one", HIVE_LONG, Optional.empty()))) - .setPartitionColumns(ImmutableList.of()) - .setOwner(Optional.empty()) - .setTableType("fake"); - tableBuilder.getStorageBuilder() - .setStorageFormat(StorageFormat.fromHiveStorageFormat(HiveStorageFormat.CSV)); - Table fakeTable = tableBuilder.build(); - - // Expected file system tree: - // test-file-iterator-listing/ - // .hidden/ - // nested-file-in-hidden.txt - // parent/ - // _nested-hidden-file.txt - // nested-file.txt - // empty-directory/ - // .hidden-in-base.txt - // base-path-file.txt - Path basePath = new Path(getBasePath(), "test-file-iterator-listing"); - FileSystem fs = hdfsEnvironment.getFileSystem(TESTING_CONTEXT, basePath); - TrinoFileSystem trinoFileSystem = new HdfsFileSystemFactory(hdfsEnvironment, new TrinoHdfsFileSystemStats()).create(SESSION); - fs.mkdirs(basePath); - - // create file in hidden folder - Path fileInHiddenParent = new Path(new Path(basePath, ".hidden"), "nested-file-in-hidden.txt"); - fs.createNewFile(fileInHiddenParent); - // create hidden file in non-hidden folder - Path nestedHiddenFile = new Path(new Path(basePath, "parent"), "_nested-hidden-file.txt"); - fs.createNewFile(nestedHiddenFile); - // create file in non-hidden folder - Path nestedFile = new Path(new Path(basePath, "parent"), "nested-file.txt"); - fs.createNewFile(nestedFile); - // create file in base path - Path baseFile = new Path(basePath, "base-path-file.txt"); - fs.createNewFile(baseFile); - // create hidden file in base path - Path hiddenBase = new Path(basePath, ".hidden-in-base.txt"); - fs.createNewFile(hiddenBase); - // create empty subdirectory - Path emptyDirectory = new Path(basePath, "empty-directory"); - fs.mkdirs(emptyDirectory); - - // List recursively through hive file iterator - HiveFileIterator recursiveIterator = new HiveFileIterator( - fakeTable, - Location.of(basePath.toString()), - trinoFileSystem, - new FileSystemDirectoryLister(), - new HdfsNamenodeStats(), - HiveFileIterator.NestedDirectoryPolicy.RECURSE); - - List recursiveListing = Streams.stream(recursiveIterator) - .map(TrinoFileStatus::getPath) - .map(Path::new) - .toList(); - // Should not include directories, or files underneath hidden directories - assertEqualsIgnoreOrder(recursiveListing, ImmutableList.of(nestedFile, baseFile)); - - HiveFileIterator shallowIterator = new HiveFileIterator( - fakeTable, - Location.of(basePath.toString()), - trinoFileSystem, - new FileSystemDirectoryLister(), - new HdfsNamenodeStats(), - HiveFileIterator.NestedDirectoryPolicy.IGNORED); - List shallowListing = Streams.stream(shallowIterator) - .map(TrinoFileStatus::getPath) - .map(Path::new) - .toList(); - // Should not include any hidden files, folders, or nested files - assertEqualsIgnoreOrder(shallowListing, ImmutableList.of(baseFile)); - } - - @Test - public void testFileIteratorPartitionedListing() - throws Exception - { - Table.Builder tableBuilder = Table.builder() - .setDatabaseName(table.getSchemaName()) - .setTableName(table.getTableName()) - .setDataColumns(ImmutableList.of(new Column("data", HIVE_LONG, Optional.empty()))) - .setPartitionColumns(ImmutableList.of(new Column("part", HIVE_STRING, Optional.empty()))) - .setOwner(Optional.empty()) - .setTableType("fake"); - tableBuilder.getStorageBuilder() - .setStorageFormat(StorageFormat.fromHiveStorageFormat(HiveStorageFormat.CSV)); - Table fakeTable = tableBuilder.build(); - - // Expected file system tree: - // test-file-iterator-partitioned-listing/ - // .hidden/ - // nested-file-in-hidden.txt - // part=simple/ - // _hidden-file.txt - // plain-file.txt - // part=nested/ - // parent/ - // _nested-hidden-file.txt - // nested-file.txt - // part=plus+sign/ - // plus-file.txt - // part=percent%sign/ - // percent-file.txt - // part=url%20encoded/ - // url-encoded-file.txt - // part=level1|level2/ - // pipe-file.txt - // parent1/ - // parent2/ - // deeply-nested-file.txt - // part=level1 | level2/ - // pipe-blanks-file.txt - // empty-directory/ - // .hidden-in-base.txt - Path basePath = new Path(getBasePath(), "test-file-iterator-partitioned-listing"); - FileSystem fs = hdfsEnvironment.getFileSystem(TESTING_CONTEXT, basePath); - TrinoFileSystem trinoFileSystem = new HdfsFileSystemFactory(hdfsEnvironment, new TrinoHdfsFileSystemStats()).create(SESSION); - fs.mkdirs(basePath); - - // create file in hidden folder - Path fileInHiddenParent = new Path(new Path(basePath, ".hidden"), "nested-file-in-hidden.txt"); - fs.createNewFile(fileInHiddenParent); - // create hidden file in non-hidden folder - Path hiddenFileUnderPartitionSimple = new Path(new Path(basePath, "part=simple"), "_hidden-file.txt"); - fs.createNewFile(hiddenFileUnderPartitionSimple); - // create file in `part=simple` non-hidden folder - Path plainFilePartitionSimple = new Path(new Path(basePath, "part=simple"), "plain-file.txt"); - fs.createNewFile(plainFilePartitionSimple); - Path nestedFilePartitionNested = new Path(new Path(new Path(basePath, "part=nested"), "parent"), "nested-file.txt"); - fs.createNewFile(nestedFilePartitionNested); - // create hidden file in non-hidden folder - Path nestedHiddenFilePartitionNested = new Path(new Path(new Path(basePath, "part=nested"), "parent"), "_nested-hidden-file.txt"); - fs.createNewFile(nestedHiddenFilePartitionNested); - // create file in `part=plus+sign` non-hidden folder (which contains `+` special character) - Path plainFilePartitionPlusSign = new Path(new Path(basePath, "part=plus+sign"), "plus-file.txt"); - fs.createNewFile(plainFilePartitionPlusSign); - // create file in `part=percent%sign` non-hidden folder (which contains `%` special character) - Path plainFilePartitionPercentSign = new Path(new Path(basePath, "part=percent%sign"), "percent-file.txt"); - fs.createNewFile(plainFilePartitionPercentSign); - // create file in `part=url%20encoded` non-hidden folder (which contains `%` special character) - Path plainFilePartitionUrlEncoded = new Path(new Path(basePath, "part=url%20encoded"), "url-encoded-file.txt"); - fs.createNewFile(plainFilePartitionUrlEncoded); - // create file in `part=level1|level2` non-hidden folder (which contains `|` special character) - Path plainFilePartitionPipeSign = new Path(new Path(basePath, "part=level1|level2"), "pipe-file.txt"); - fs.createNewFile(plainFilePartitionPipeSign); - Path deeplyNestedFilePartitionPipeSign = new Path(new Path(new Path(new Path(basePath, "part=level1|level2"), "parent1"), "parent2"), "deeply-nested-file.txt"); - fs.createNewFile(deeplyNestedFilePartitionPipeSign); - // create file in `part=level1 | level2` non-hidden folder (which contains `|` and blank space special characters) - Path plainFilePartitionPipeSignBlanks = new Path(new Path(basePath, "part=level1 | level2"), "pipe-blanks-file.txt"); - fs.createNewFile(plainFilePartitionPipeSignBlanks); - - // create empty subdirectory - Path emptyDirectory = new Path(basePath, "empty-directory"); - fs.mkdirs(emptyDirectory); - // create hidden file in base path - Path hiddenBase = new Path(basePath, ".hidden-in-base.txt"); - fs.createNewFile(hiddenBase); - - // List recursively through hive file iterator - HiveFileIterator recursiveIterator = new HiveFileIterator( - fakeTable, - Location.of(basePath.toString()), - trinoFileSystem, - new FileSystemDirectoryLister(), - new HdfsNamenodeStats(), - HiveFileIterator.NestedDirectoryPolicy.RECURSE); - - List recursiveListing = Streams.stream(recursiveIterator) - .map(TrinoFileStatus::getPath) - .map(Path::new) - .toList(); - // Should not include directories, or files underneath hidden directories - assertThat(recursiveListing).containsExactlyInAnyOrder( - plainFilePartitionSimple, - nestedFilePartitionNested, - plainFilePartitionPlusSign, - plainFilePartitionPercentSign, - plainFilePartitionUrlEncoded, - plainFilePartitionPipeSign, - deeplyNestedFilePartitionPipeSign, - plainFilePartitionPipeSignBlanks); - - HiveFileIterator shallowIterator = new HiveFileIterator( - fakeTable, - Location.of(basePath.toString()), - trinoFileSystem, - new FileSystemDirectoryLister(), - new HdfsNamenodeStats(), - HiveFileIterator.NestedDirectoryPolicy.IGNORED); - List shallowListing = Streams.stream(shallowIterator) - .map(TrinoFileStatus::getPath) - .map(Path::new) - .toList(); - // Should not include any hidden files, folders, or nested files - assertThat(shallowListing).isEmpty(); - } - - @Test - public void testDirectoryWithTrailingSpace() - throws Exception - { - Path basePath = new Path(getBasePath(), randomUUID().toString()); - FileSystem fs = hdfsEnvironment.getFileSystem(TESTING_CONTEXT, basePath); - assertFalse(fs.exists(basePath)); - - Path path = new Path(new Path(basePath, "dir_with_space "), "foo.txt"); - try (OutputStream outputStream = fs.create(path)) { - outputStream.write("test".getBytes(UTF_8)); - } - assertTrue(fs.exists(path)); - - try (InputStream inputStream = fs.open(path)) { - String content = new BufferedReader(new InputStreamReader(inputStream, UTF_8)).readLine(); - assertEquals(content, "test"); - } - - fs.delete(basePath, true); - } - - @Test - public void testTableCreation() - throws Exception - { - for (HiveStorageFormat storageFormat : HiveStorageFormat.values()) { - if (storageFormat == HiveStorageFormat.CSV) { - // CSV supports only unbounded VARCHAR type - continue; - } - if (storageFormat == HiveStorageFormat.REGEX) { - // REGEX format is read-only - continue; - } - createTable(temporaryCreateTable, storageFormat); - dropTable(temporaryCreateTable); - } - } - - private void createTable(SchemaTableName tableName, HiveStorageFormat storageFormat) - throws Exception - { - List columns = ImmutableList.of(new ColumnMetadata("id", BIGINT)); - - MaterializedResult data = MaterializedResult.resultBuilder(newSession(), BIGINT) - .row(1L) - .row(3L) - .row(2L) - .build(); - - try (Transaction transaction = newTransaction()) { - ConnectorMetadata metadata = transaction.getMetadata(); - ConnectorSession session = newSession(); - - // begin creating the table - ConnectorTableMetadata tableMetadata = new ConnectorTableMetadata(tableName, columns, createTableProperties(storageFormat)); - ConnectorOutputTableHandle outputHandle = metadata.beginCreateTable(session, tableMetadata, Optional.empty(), NO_RETRIES); - - // write the records - ConnectorPageSink sink = pageSinkProvider.createPageSink(transaction.getTransactionHandle(), session, outputHandle, TESTING_PAGE_SINK_ID); - sink.appendPage(data.toPage()); - Collection fragments = getFutureValue(sink.finish()); - - // commit the table - metadata.finishCreateTable(session, outputHandle, fragments, ImmutableList.of()); - - transaction.commit(); - - // Hack to work around the metastore not being configured for S3 or other FS. - // The metastore tries to validate the location when creating the - // table, which fails without explicit configuration for file system. - // We work around that by using a dummy location when creating the - // table and update it here to the correct location. - Location location = locationService.getTableWriteInfo(((HiveOutputTableHandle) outputHandle).getLocationHandle(), false).targetPath(); - metastoreClient.updateTableLocation(database, tableName.getTableName(), location.toString()); - } - - try (Transaction transaction = newTransaction()) { - ConnectorMetadata metadata = transaction.getMetadata(); - ConnectorSession session = newSession(); - - // load the new table - ConnectorTableHandle tableHandle = getTableHandle(metadata, tableName); - List columnHandles = filterNonHiddenColumnHandles(metadata.getColumnHandles(session, tableHandle).values()); - - // verify the metadata - ConnectorTableMetadata tableMetadata = metadata.getTableMetadata(session, getTableHandle(metadata, tableName)); - assertEquals(filterNonHiddenColumnMetadata(tableMetadata.getColumns()), columns); - - // verify the data - metadata.beginQuery(session); - ConnectorSplitSource splitSource = getSplits(splitManager, transaction, session, tableHandle); - ConnectorSplit split = getOnlyElement(getAllSplits(splitSource)); - - try (ConnectorPageSource pageSource = pageSourceProvider.createPageSource(transaction.getTransactionHandle(), session, split, tableHandle, columnHandles, DynamicFilter.EMPTY)) { - MaterializedResult result = materializeSourceDataStream(session, pageSource, getTypes(columnHandles)); - assertEqualsIgnoreOrder(result.getMaterializedRows(), data.getMaterializedRows()); - } - - metadata.cleanupQuery(session); - } - } - - private void dropTable(SchemaTableName table) - { - try (Transaction transaction = newTransaction()) { - transaction.getMetastore().dropTable(newSession(), table.getSchemaName(), table.getTableName()); - transaction.commit(); - } - } - - private ConnectorTableHandle getTableHandle(ConnectorMetadata metadata, SchemaTableName tableName) - { - return HiveFileSystemTestUtils.getTableHandle(metadata, tableName, newSession()); - } - - public static class TestingHiveMetastore - extends ForwardingHiveMetastore - { - private final Path basePath; - private final HdfsEnvironment hdfsEnvironment; - - public TestingHiveMetastore(HiveMetastore delegate, Path basePath, HdfsEnvironment hdfsEnvironment) - { - super(delegate); - this.basePath = basePath; - this.hdfsEnvironment = hdfsEnvironment; - } - - @Override - public Optional getDatabase(String databaseName) - { - return super.getDatabase(databaseName) - .map(database -> Database.builder(database) - .setLocation(Optional.of(basePath.toString())) - .build()); - } - - @Override - public void createTable(Table table, PrincipalPrivileges privileges) - { - // hack to work around the metastore not being configured for S3 or other FS - Table.Builder tableBuilder = Table.builder(table); - tableBuilder.getStorageBuilder().setLocation("/"); - super.createTable(tableBuilder.build(), privileges); - } - - @Override - public void dropTable(String databaseName, String tableName, boolean deleteData) - { - try { - Table table = getTable(databaseName, tableName) - .orElseThrow(() -> new TableNotFoundException(new SchemaTableName(databaseName, tableName))); - - // hack to work around the metastore not being configured for S3 or other FS - List locations = listAllDataPaths(databaseName, tableName); - - Table.Builder tableBuilder = Table.builder(table); - tableBuilder.getStorageBuilder().setLocation("/"); - - // drop table - replaceTable(databaseName, tableName, tableBuilder.build(), NO_PRIVILEGES); - super.dropTable(databaseName, tableName, false); - - // drop data - if (deleteData) { - for (String location : locations) { - Path path = new Path(location); - hdfsEnvironment.getFileSystem(TESTING_CONTEXT, path).delete(path, true); - } - } - } - catch (IOException e) { - throw new UncheckedIOException(e); - } - } - - public void updateTableLocation(String databaseName, String tableName, String location) - { - Table table = getTable(databaseName, tableName) - .orElseThrow(() -> new TableNotFoundException(new SchemaTableName(databaseName, tableName))); - Table.Builder tableBuilder = Table.builder(table); - tableBuilder.getStorageBuilder().setLocation(location); - - // NOTE: this clears the permissions - replaceTable(databaseName, tableName, tableBuilder.build(), NO_PRIVILEGES); - } - - private List listAllDataPaths(String schemaName, String tableName) - { - ImmutableList.Builder locations = ImmutableList.builder(); - Table table = getTable(schemaName, tableName).get(); - List partitionColumnNames = table.getPartitionColumns().stream().map(Column::getName).collect(toImmutableList()); - if (table.getStorage().getLocation() != null) { - // For partitioned table, there should be nothing directly under this directory. - // But including this location in the set makes the directory content assert more - // extensive, which is desirable. - locations.add(table.getStorage().getLocation()); - } - - Optional> partitionNames = getPartitionNamesByFilter(schemaName, tableName, partitionColumnNames, TupleDomain.all()); - if (partitionNames.isPresent()) { - getPartitionsByNames(table, partitionNames.get()).values().stream() - .map(Optional::get) - .map(partition -> partition.getStorage().getLocation()) - .filter(location -> !location.startsWith(table.getStorage().getLocation())) - .forEach(locations::add); - } - - return locations.build(); - } - } -} diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/AbstractTestHiveLocal.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/AbstractTestHiveLocal.java deleted file mode 100644 index f97a28461b84..000000000000 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/AbstractTestHiveLocal.java +++ /dev/null @@ -1,285 +0,0 @@ -/* - * 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.hive; - -import com.google.common.collect.ImmutableList; -import com.google.common.collect.ImmutableMap; -import com.google.common.io.RecursiveDeleteOption; -import com.google.common.reflect.ClassPath; -import io.airlift.log.Logger; -import io.trino.plugin.hive.metastore.Column; -import io.trino.plugin.hive.metastore.Database; -import io.trino.plugin.hive.metastore.HiveMetastore; -import io.trino.plugin.hive.metastore.PrincipalPrivileges; -import io.trino.plugin.hive.metastore.SortingColumn; -import io.trino.plugin.hive.metastore.StorageFormat; -import io.trino.plugin.hive.metastore.Table; -import io.trino.spi.connector.ColumnHandle; -import io.trino.spi.connector.ConnectorMetadata; -import io.trino.spi.connector.ConnectorSession; -import io.trino.spi.connector.ConnectorTableHandle; -import io.trino.spi.connector.SchemaTableName; -import io.trino.spi.predicate.TupleDomain; -import io.trino.spi.security.PrincipalType; -import io.trino.testing.MaterializedResult; -import org.apache.hadoop.fs.Path; -import org.apache.hadoop.hive.metastore.TableType; -import org.testng.SkipException; -import org.testng.annotations.AfterClass; -import org.testng.annotations.BeforeClass; -import org.testng.annotations.Test; - -import java.io.File; -import java.io.IOException; -import java.io.InputStream; -import java.io.UncheckedIOException; -import java.net.URI; -import java.nio.file.Files; -import java.util.List; -import java.util.Optional; -import java.util.OptionalInt; - -import static com.google.common.io.MoreFiles.deleteRecursively; -import static com.google.common.io.RecursiveDeleteOption.ALLOW_INSECURE; -import static io.trino.plugin.hive.HiveMetadata.PRESTO_QUERY_ID_NAME; -import static io.trino.plugin.hive.HiveMetadata.PRESTO_VERSION_NAME; -import static io.trino.plugin.hive.HiveStorageFormat.ORC; -import static io.trino.plugin.hive.HiveTestUtils.HDFS_ENVIRONMENT; -import static io.trino.plugin.hive.HiveType.HIVE_INT; -import static io.trino.plugin.hive.HiveType.HIVE_STRING; -import static io.trino.plugin.hive.util.HiveBucketing.BucketingVersion.BUCKETING_V1; -import static io.trino.plugin.hive.util.HiveUtil.SPARK_TABLE_PROVIDER_KEY; -import static java.nio.file.Files.copy; -import static java.util.Objects.requireNonNull; -import static org.testng.Assert.assertEquals; - -public abstract class AbstractTestHiveLocal - extends AbstractTestHive -{ - private static final Logger log = Logger.get(AbstractTestHiveLocal.class); - private static final String DEFAULT_TEST_DB_NAME = "test"; - - private File tempDir; - private final String testDbName; - - protected AbstractTestHiveLocal() - { - this(DEFAULT_TEST_DB_NAME); - } - - protected AbstractTestHiveLocal(String testDbName) - { - this.testDbName = requireNonNull(testDbName, "testDbName is null"); - } - - protected abstract HiveMetastore createMetastore(File tempDir); - - @BeforeClass(alwaysRun = true) - public void initialize() - throws Exception - { - tempDir = Files.createTempDirectory(null).toFile(); - - HiveMetastore metastore = createMetastore(tempDir); - - metastore.createDatabase( - Database.builder() - .setDatabaseName(testDbName) - .setOwnerName(Optional.of("public")) - .setOwnerType(Optional.of(PrincipalType.ROLE)) - .build()); - - HiveConfig hiveConfig = new HiveConfig() - .setParquetTimeZone("America/Los_Angeles") - .setRcfileTimeZone("America/Los_Angeles"); - - setup(testDbName, hiveConfig, metastore, HDFS_ENVIRONMENT); - } - - @AfterClass(alwaysRun = true) - public void cleanup() - throws IOException - { - try { - getMetastoreClient().dropDatabase(testDbName, true); - } - finally { - deleteRecursively(tempDir.toPath(), ALLOW_INSECURE); - } - } - - @Override - protected ConnectorTableHandle getTableHandle(ConnectorMetadata metadata, SchemaTableName tableName) - { - if (tableName.getTableName().startsWith(TEMPORARY_TABLE_PREFIX)) { - return super.getTableHandle(metadata, tableName); - } - throw new SkipException("tests using existing tables are not supported"); - } - - @Override - public void testGetAllTableNames() - { - throw new SkipException("Test disabled for this subclass"); - } - - @Override - public void testGetAllTableColumns() - { - throw new SkipException("Test disabled for this subclass"); - } - - @Override - public void testGetAllTableColumnsInSchema() - { - throw new SkipException("Test disabled for this subclass"); - } - - @Override - public void testGetTableNames() - { - throw new SkipException("Test disabled for this subclass"); - } - - @Override - public void testGetTableSchemaOffline() - { - throw new SkipException("Test disabled for this subclass"); - } - - @Test - public void testSparkBucketedTableValidation() - throws Exception - { - SchemaTableName table = temporaryTable("spark_bucket_validation"); - try { - doTestSparkBucketedTableValidation(table); - } - finally { - dropTable(table); - } - } - - private void doTestSparkBucketedTableValidation(SchemaTableName tableName) - throws Exception - { - java.nio.file.Path externalLocation = copyResourceDirToTemporaryDirectory("spark_bucketed_nation"); - try { - createExternalTable( - tableName, - ORC, - ImmutableList.of( - new Column("nationkey", HIVE_INT, Optional.empty()), - new Column("name", HIVE_STRING, Optional.empty()), - new Column("regionkey", HIVE_INT, Optional.empty()), - new Column("comment", HIVE_STRING, Optional.empty())), - ImmutableList.of(), - Optional.of(new HiveBucketProperty( - ImmutableList.of("nationkey"), - BUCKETING_V1, - 3, - ImmutableList.of(new SortingColumn("name", SortingColumn.Order.ASCENDING)))), - new Path(URI.create("file://" + externalLocation.toString()))); - - assertReadFailsWithMessageMatching(ORC, tableName, "Hive table is corrupt\\. File '.*/.*' is for bucket [0-2], but contains a row for bucket [0-2]."); - markTableAsCreatedBySpark(tableName, "orc"); - assertReadReturnsRowCount(ORC, tableName, 25); - } - finally { - deleteRecursively(externalLocation, RecursiveDeleteOption.ALLOW_INSECURE); - } - } - - private void assertReadReturnsRowCount(HiveStorageFormat storageFormat, SchemaTableName tableName, int rowCount) - throws Exception - { - try (Transaction transaction = newTransaction()) { - ConnectorSession session = newSession(); - ConnectorMetadata metadata = transaction.getMetadata(); - metadata.beginQuery(session); - ConnectorTableHandle tableHandle = getTableHandle(metadata, tableName); - List columnHandles = filterNonHiddenColumnHandles(metadata.getColumnHandles(session, tableHandle).values()); - MaterializedResult result = readTable(transaction, tableHandle, columnHandles, session, TupleDomain.all(), OptionalInt.empty(), Optional.of(storageFormat)); - assertEquals(result.getRowCount(), rowCount); - } - } - - private void markTableAsCreatedBySpark(SchemaTableName tableName, String provider) - { - try (Transaction transaction = newTransaction()) { - ConnectorSession session = newSession(); - PrincipalPrivileges principalPrivileges = testingPrincipalPrivilege(session); - Table oldTable = transaction.getMetastore().getTable(tableName.getSchemaName(), tableName.getTableName()).get(); - Table.Builder newTable = Table.builder(oldTable).setParameter(SPARK_TABLE_PROVIDER_KEY, provider); - transaction.getMetastore().replaceTable(tableName.getSchemaName(), tableName.getTableName(), newTable.build(), principalPrivileges); - transaction.commit(); - } - } - - private void createExternalTable(SchemaTableName schemaTableName, HiveStorageFormat hiveStorageFormat, List columns, List partitionColumns, Optional bucketProperty, Path externalLocation) - { - try (Transaction transaction = newTransaction()) { - ConnectorSession session = newSession(); - - String tableOwner = session.getUser(); - String schemaName = schemaTableName.getSchemaName(); - String tableName = schemaTableName.getTableName(); - - Table.Builder tableBuilder = Table.builder() - .setDatabaseName(schemaName) - .setTableName(tableName) - .setOwner(Optional.of(tableOwner)) - .setTableType(TableType.EXTERNAL_TABLE.name()) - .setParameters(ImmutableMap.of( - PRESTO_VERSION_NAME, TEST_SERVER_VERSION, - PRESTO_QUERY_ID_NAME, session.getQueryId())) - .setDataColumns(columns) - .setPartitionColumns(partitionColumns); - - tableBuilder.getStorageBuilder() - .setLocation(externalLocation.toString()) - .setStorageFormat(StorageFormat.create(hiveStorageFormat.getSerde(), hiveStorageFormat.getInputFormat(), hiveStorageFormat.getOutputFormat())) - .setBucketProperty(bucketProperty) - .setSerdeParameters(ImmutableMap.of()); - - PrincipalPrivileges principalPrivileges = testingPrincipalPrivilege(tableOwner, session.getUser()); - transaction.getMetastore().createTable(session, tableBuilder.build(), principalPrivileges, Optional.of(externalLocation), Optional.empty(), true, ZERO_TABLE_STATISTICS, false); - - transaction.commit(); - } - } - - private java.nio.file.Path copyResourceDirToTemporaryDirectory(String resourceName) - throws IOException - { - java.nio.file.Path tempDir = java.nio.file.Files.createTempDirectory(getClass().getSimpleName()).normalize(); - log.info("Copying resource dir '%s' to %s", resourceName, tempDir); - ClassPath.from(getClass().getClassLoader()) - .getResources().stream() - .filter(resourceInfo -> resourceInfo.getResourceName().startsWith(resourceName)) - .forEach(resourceInfo -> { - try { - java.nio.file.Path target = tempDir.resolve(resourceInfo.getResourceName()); - java.nio.file.Files.createDirectories(target.getParent()); - try (InputStream inputStream = resourceInfo.asByteSource().openStream()) { - copy(inputStream, target); - } - } - catch (IOException e) { - throw new UncheckedIOException(e); - } - }); - return tempDir.resolve(resourceName).normalize(); - } -} diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/AbstractTestParquetPageSkipping.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/AbstractTestParquetPageSkipping.java index 89599ae8382d..029d55fab76e 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/AbstractTestParquetPageSkipping.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/AbstractTestParquetPageSkipping.java @@ -16,6 +16,8 @@ import com.google.common.io.Resources; import io.trino.Session; import io.trino.execution.QueryStats; +import io.trino.filesystem.Location; +import io.trino.filesystem.TrinoFileSystem; import io.trino.operator.OperatorStats; import io.trino.spi.QueryId; import io.trino.spi.metrics.Count; @@ -29,8 +31,11 @@ import org.testng.annotations.Test; import java.io.File; -import java.net.URISyntaxException; +import java.io.IOException; +import java.io.OutputStream; +import java.net.URL; import java.util.Map; +import java.util.UUID; import static com.google.common.collect.MoreCollectors.onlyElement; import static io.trino.parquet.reader.ParquetReader.COLUMN_INDEX_ROWS_FILTERED; @@ -42,6 +47,8 @@ public abstract class AbstractTestParquetPageSkipping extends AbstractTestQueryFramework { + protected TrinoFileSystem fileSystem; + private void buildSortedTables(String tableName, String sortByColumnName, String sortByColumnType) { String createTableTemplate = @@ -83,6 +90,7 @@ private void buildSortedTables(String tableName, String sortByColumnName, String public void testRowGroupPruningFromPageIndexes() throws Exception { + Location dataFile = copyInDataFile("parquet_page_skipping/orders_sorted_by_totalprice/data.parquet"); String tableName = "test_row_group_pruning_" + randomNameSuffix(); File parquetFile = new File(Resources.getResource("parquet_page_skipping/orders_sorted_by_totalprice").toURI()); assertUpdate( @@ -101,29 +109,29 @@ comment varchar(79), WITH ( format = 'PARQUET', external_location = '%s') - """.formatted(tableName, parquetFile.getAbsolutePath())); + """.formatted(tableName, dataFile.parentDirectory())); int rowCount = assertColumnIndexResults("SELECT * FROM " + tableName + " WHERE totalprice BETWEEN 100000 AND 131280 AND clerk = 'Clerk#000000624'"); assertThat(rowCount).isGreaterThan(0); - // `totalprice BETWEEN 51890 AND 51900` is chosen to lie between min/max values of row group - // but outside page level min/max boundaries to trigger pruning of row group using column index +// `totalprice BETWEEN 51890 AND 51900` is chosen to lie between min/max values of row group +// but outside page level min/max boundaries to trigger pruning of row group using column index assertRowGroupPruning("SELECT * FROM " + tableName + " WHERE totalprice BETWEEN 51890 AND 51900 AND orderkey > 0"); assertUpdate("DROP TABLE " + tableName); } @Test public void testPageSkippingWithNonSequentialOffsets() - throws URISyntaxException + throws IOException { + Location dataFile = copyInDataFile("parquet_page_skipping/random/data.parquet"); String tableName = "test_random_" + randomNameSuffix(); - File parquetFile = new File(Resources.getResource("parquet_page_skipping/random").toURI()); assertUpdate(format( "CREATE TABLE %s (col double) WITH (format = 'PARQUET', external_location = '%s')", tableName, - parquetFile.getAbsolutePath())); - // These queries select a subset of pages which are stored at non-sequential offsets - // This reproduces the issue identified in https://github.com/trinodb/trino/issues/9097 + dataFile.parentDirectory())); +// These queries select a subset of pages which are stored at non-sequential offsets +// This reproduces the issue identified in https://github.com/trinodb/trino/issues/9097 for (double i = 0; i < 1; i += 0.1) { assertColumnIndexResults(format("SELECT * FROM %s WHERE col BETWEEN %f AND %f", tableName, i - 0.00001, i + 0.00001)); } @@ -132,17 +140,18 @@ public void testPageSkippingWithNonSequentialOffsets() @Test public void testFilteringOnColumnNameWithDot() - throws URISyntaxException + throws IOException { + Location dataFile = copyInDataFile("parquet_page_skipping/column_name_with_dot/data.parquet"); + String nameInSql = "\"a.dot\""; String tableName = "test_column_name_with_dot_" + randomNameSuffix(); - File parquetFile = new File(Resources.getResource("parquet_page_skipping/column_name_with_dot").toURI()); assertUpdate(format( "CREATE TABLE %s (key varchar(50), %s varchar(50)) WITH (format = 'PARQUET', external_location = '%s')", tableName, nameInSql, - parquetFile.getAbsolutePath())); + dataFile.parentDirectory())); assertQuery("SELECT key FROM " + tableName + " WHERE " + nameInSql + " IS NULL", "VALUES ('null value')"); assertQuery("SELECT key FROM " + tableName + " WHERE " + nameInSql + " = 'abc'", "VALUES ('sample value')"); @@ -164,14 +173,14 @@ public void testPageSkipping(String sortByColumn, String sortByColumnType, Objec assertThat(assertColumnIndexResults(format("SELECT %s FROM %s WHERE %s < %s", sortByColumn, tableName, sortByColumn, lowValue))).isGreaterThan(0); assertThat(assertColumnIndexResults(format("SELECT %s FROM %s WHERE %s > %s", sortByColumn, tableName, sortByColumn, highValue))).isGreaterThan(0); assertThat(assertColumnIndexResults(format("SELECT %s FROM %s WHERE %s BETWEEN %s AND %s", sortByColumn, tableName, sortByColumn, middleLowValue, middleHighValue))).isGreaterThan(0); - // Tests synchronization of reading values across columns +// Tests synchronization of reading values across columns assertColumnIndexResults(format("SELECT * FROM %s WHERE %s = %s", tableName, sortByColumn, middleLowValue)); assertThat(assertColumnIndexResults(format("SELECT * FROM %s WHERE %s < %s", tableName, sortByColumn, lowValue))).isGreaterThan(0); assertThat(assertColumnIndexResults(format("SELECT * FROM %s WHERE %s > %s", tableName, sortByColumn, highValue))).isGreaterThan(0); assertThat(assertColumnIndexResults(format("SELECT * FROM %s WHERE %s BETWEEN %s AND %s", tableName, sortByColumn, middleLowValue, middleHighValue))).isGreaterThan(0); - // Nested data +// Nested data assertColumnIndexResults(format("SELECT rvalues FROM %s WHERE %s IN (%s, %s, %s, %s)", tableName, sortByColumn, lowValue, middleLowValue, middleHighValue, highValue)); - // Without nested data +// Without nested data assertColumnIndexResults(format("SELECT orderkey, orderdate FROM %s WHERE %s IN (%s, %s, %s, %s)", tableName, sortByColumn, lowValue, middleLowValue, middleHighValue, highValue)); } assertUpdate("DROP TABLE " + tableName); @@ -179,15 +188,15 @@ public void testPageSkipping(String sortByColumn, String sortByColumnType, Objec @Test public void testFilteringWithColumnIndex() - throws URISyntaxException + throws IOException { + Location dataFile = copyInDataFile("parquet_page_skipping/lineitem_sorted_by_suppkey/data.parquet"); String tableName = "test_page_filtering_" + randomNameSuffix(); - File parquetFile = new File(Resources.getResource("parquet_page_skipping/lineitem_sorted_by_suppkey").toURI()); assertUpdate(format( "CREATE TABLE %s (suppkey bigint, extendedprice decimal(12, 2), shipmode varchar(10), comment varchar(44)) " + "WITH (format = 'PARQUET', external_location = '%s')", tableName, - parquetFile.getAbsolutePath())); + dataFile.parentDirectory())); verifyFilteringWithColumnIndex("SELECT * FROM " + tableName + " WHERE suppkey = 10"); verifyFilteringWithColumnIndex("SELECT * FROM " + tableName + " WHERE suppkey BETWEEN 25 AND 35"); @@ -303,4 +312,18 @@ private OperatorStats getScanOperatorStats(QueryId queryId) .filter(summary -> summary.getOperatorType().startsWith("TableScan") || summary.getOperatorType().startsWith("Scan")) .collect(onlyElement()); } + + private Location copyInDataFile(String resourceFileName) + throws IOException + { + URL resourceLocation = Resources.getResource(resourceFileName); + + Location tempDir = Location.of("local:///temp_" + UUID.randomUUID()); + fileSystem.createDirectory(tempDir); + Location dataFile = tempDir.appendPath("data.parquet"); + try (OutputStream out = fileSystem.newOutputFile(dataFile).create()) { + Resources.copy(resourceLocation, out); + } + return dataFile; + } } diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/AvroSchemaGenerationTests.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/AvroSchemaGenerationTests.java deleted file mode 100644 index 554772fa7601..000000000000 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/AvroSchemaGenerationTests.java +++ /dev/null @@ -1,64 +0,0 @@ -/* - * 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.hive; - -import io.trino.filesystem.local.LocalFileSystem; -import io.trino.hadoop.ConfigurationInstantiator; -import io.trino.plugin.hive.avro.AvroHiveFileUtils; -import io.trino.plugin.hive.avro.TrinoAvroSerDe; -import io.trino.plugin.hive.type.TypeInfo; -import io.trino.spi.type.RowType; -import io.trino.spi.type.VarcharType; -import org.apache.avro.Schema; -import org.junit.jupiter.api.Test; - -import java.io.IOException; -import java.nio.file.Path; -import java.util.Properties; -import java.util.stream.Collectors; -import java.util.stream.Stream; - -import static io.trino.plugin.hive.avro.AvroHiveConstants.TABLE_NAME; -import static io.trino.plugin.hive.util.SerdeConstants.LIST_COLUMNS; -import static io.trino.plugin.hive.util.SerdeConstants.LIST_COLUMN_TYPES; -import static org.assertj.core.api.Assertions.assertThat; - -public class AvroSchemaGenerationTests -{ - @Test - public void testOldVsNewSchemaGeneration() - throws IOException - { - Properties properties = new Properties(); - properties.setProperty(TABLE_NAME, "testingTable"); - properties.setProperty(LIST_COLUMNS, "a,b"); - properties.setProperty(LIST_COLUMN_TYPES, Stream.of(HiveType.HIVE_INT, HiveType.HIVE_STRING).map(HiveType::getTypeInfo).map(TypeInfo::toString).collect(Collectors.joining(","))); - Schema actual = AvroHiveFileUtils.determineSchemaOrThrowException(new LocalFileSystem(Path.of("/")), properties); - Schema expected = new TrinoAvroSerDe().determineSchemaOrReturnErrorSchema(ConfigurationInstantiator.newEmptyConfiguration(), properties); - assertThat(actual).isEqualTo(expected); - } - - @Test - public void testOldVsNewSchemaGenerationWithNested() - throws IOException - { - Properties properties = new Properties(); - properties.setProperty(TABLE_NAME, "testingTable"); - properties.setProperty(LIST_COLUMNS, "a,b"); - properties.setProperty(LIST_COLUMN_TYPES, Stream.of(HiveType.toHiveType(RowType.rowType(RowType.field("a", VarcharType.VARCHAR))), HiveType.HIVE_STRING).map(HiveType::getTypeInfo).map(TypeInfo::toString).collect(Collectors.joining(","))); - Schema actual = AvroHiveFileUtils.determineSchemaOrThrowException(new LocalFileSystem(Path.of("/")), properties); - Schema expected = new TrinoAvroSerDe().determineSchemaOrReturnErrorSchema(ConfigurationInstantiator.newEmptyConfiguration(), properties); - assertThat(actual).isEqualTo(expected); - } -} diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/BaseHiveConnectorTest.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/BaseHiveConnectorTest.java index 9e58658d6186..6369c09ebb12 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/BaseHiveConnectorTest.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/BaseHiveConnectorTest.java @@ -16,6 +16,7 @@ import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; +import com.google.common.io.Resources; import io.airlift.json.JsonCodec; import io.airlift.json.JsonCodecFactory; import io.airlift.json.ObjectMapperProvider; @@ -23,6 +24,9 @@ import io.trino.Session; import io.trino.cost.StatsAndCosts; import io.trino.execution.QueryInfo; +import io.trino.filesystem.Location; +import io.trino.filesystem.TrinoFileSystem; +import io.trino.filesystem.TrinoFileSystemFactory; import io.trino.metadata.FunctionManager; import io.trino.metadata.InsertTableHandle; import io.trino.metadata.Metadata; @@ -40,6 +44,7 @@ import io.trino.spi.connector.ColumnHandle; import io.trino.spi.connector.ColumnMetadata; import io.trino.spi.connector.Constraint; +import io.trino.spi.security.ConnectorIdentity; import io.trino.spi.security.Identity; import io.trino.spi.security.SelectedRole; import io.trino.spi.type.DateType; @@ -74,9 +79,10 @@ import org.testng.annotations.DataProvider; import org.testng.annotations.Test; -import java.io.File; import java.io.IOException; +import java.io.OutputStream; import java.math.BigDecimal; +import java.net.URL; import java.time.Instant; import java.time.LocalDate; import java.time.LocalDateTime; @@ -93,6 +99,7 @@ import java.util.OptionalLong; import java.util.Set; import java.util.StringJoiner; +import java.util.UUID; import java.util.function.BiConsumer; import java.util.function.Consumer; import java.util.function.Function; @@ -144,6 +151,7 @@ import static io.trino.plugin.hive.HiveTableProperties.PARTITIONED_BY_PROPERTY; import static io.trino.plugin.hive.HiveTableProperties.STORAGE_FORMAT_PROPERTY; import static io.trino.plugin.hive.HiveType.toHiveType; +import static io.trino.plugin.hive.TestingHiveUtils.getConnectorService; import static io.trino.plugin.hive.util.HiveUtil.columnExtraInfo; import static io.trino.spi.security.Identity.ofUser; import static io.trino.spi.security.SelectedRole.Type.ROLE; @@ -173,14 +181,12 @@ import static io.trino.testing.TestingAccessControlManager.privilege; import static io.trino.testing.TestingNames.randomNameSuffix; import static io.trino.testing.TestingSession.testSessionBuilder; -import static io.trino.testing.containers.TestContainers.getPathFromClassPathResource; import static io.trino.transaction.TransactionBuilder.transaction; import static io.trino.type.InternalTypeManager.TESTING_TYPE_MANAGER; import static java.lang.String.format; import static java.lang.String.join; import static java.nio.charset.StandardCharsets.UTF_8; import static java.nio.file.Files.createTempDirectory; -import static java.nio.file.Files.writeString; import static java.util.Collections.nCopies; import static java.util.Locale.ENGLISH; import static java.util.Objects.requireNonNull; @@ -196,7 +202,6 @@ import static org.testng.Assert.assertNull; import static org.testng.Assert.assertTrue; import static org.testng.Assert.fail; -import static org.testng.FileAssert.assertFile; public abstract class BaseHiveConnectorTest extends BaseConnectorTest @@ -2053,7 +2058,7 @@ public void testPropertiesTable() assertUpdate(createTable, "SELECT count(*) FROM orders"); String queryId = (String) computeScalar("SELECT query_id FROM system.runtime.queries WHERE query LIKE 'CREATE TABLE test_show_properties%'"); String nodeVersion = (String) computeScalar("SELECT node_version FROM system.runtime.nodes WHERE coordinator"); - assertQuery("SELECT \"orc.bloom.filter.columns\", \"orc.bloom.filter.fpp\", presto_query_id, presto_version, transactional FROM \"test_show_properties$properties\"", + assertQuery("SELECT \"orc.bloom.filter.columns\", \"orc.bloom.filter.fpp\", trino_query_id, trino_version, transactional FROM \"test_show_properties$properties\"", format("SELECT 'ship_priority,order_status', '0.5', '%s', '%s', 'false'", queryId, nodeVersion)); assertUpdate("DROP TABLE test_show_properties"); } @@ -2920,7 +2925,7 @@ public void testUnregisterRegisterPartition() List paths = getQueryRunner().execute(getSession(), "SELECT \"$path\" FROM " + tableName + " ORDER BY \"$path\" ASC").toTestTypes().getMaterializedRows(); assertEquals(paths.size(), 3); - String firstPartition = new Path((String) paths.get(0).getField(0)).getParent().toString(); + String firstPartition = Location.of((String) paths.get(0).getField(0)).parentDirectory().toString(); assertAccessDenied( format("CALL system.unregister_partition('%s', '%s', ARRAY['part'], ARRAY['first'])", TPCH_SCHEMA, tableName), @@ -3799,7 +3804,7 @@ private List getPartitions(String tableName) private int getBucketCount(String tableName) { - return (int) getHiveTableProperty(tableName, table -> table.getBucketHandle().get().getTableBucketCount()); + return (int) getHiveTableProperty(tableName, table -> table.getTablePartitioning().get().tableBucketCount()); } @Test @@ -4383,14 +4388,17 @@ private void testCreateExternalTable( List tableProperties) throws Exception { - java.nio.file.Path tempDir = createTempDirectory(null); - File dataFile = tempDir.resolve("test.txt").toFile(); - writeString(dataFile.toPath(), fileContents); + TrinoFileSystem fileSystem = getTrinoFileSystem(); + Location tempDir = Location.of("local:///temp_" + UUID.randomUUID()); + fileSystem.createDirectory(tempDir); + Location dataFile = tempDir.appendPath("text.text"); + try (OutputStream out = fileSystem.newOutputFile(dataFile).create()) { + out.write(fileContents.getBytes(UTF_8)); + } // Table properties StringJoiner propertiesSql = new StringJoiner(",\n "); - propertiesSql.add( - format("external_location = '%s'", new Path(tempDir.toUri().toASCIIString()))); + propertiesSql.add(format("external_location = '%s'", tempDir)); propertiesSql.add("format = 'TEXTFILE'"); tableProperties.forEach(propertiesSql::add); @@ -4413,8 +4421,8 @@ private void testCreateExternalTable( assertQuery(format("SELECT col1, col2 from %s", tableName), expectedResults); assertUpdate(format("DROP TABLE %s", tableName)); - assertFile(dataFile); // file should still exist after drop - deleteRecursively(tempDir, ALLOW_INSECURE); + assertThat(fileSystem.newInputFile(dataFile).exists()).isTrue(); // file should still exist after drop + fileSystem.deleteDirectory(tempDir); } @Test @@ -5874,14 +5882,14 @@ public void testMismatchedBucketing() " test_mismatch_bucketingN\n" + "ON key16=keyN"; - assertUpdate(withoutMismatchOptimization, writeToTableWithMoreBuckets, 15000, assertRemoteExchangesCount(3)); + assertUpdate(withoutMismatchOptimization, writeToTableWithMoreBuckets, 15000, assertRemoteExchangesCount(5)); assertQuery("SELECT * FROM test_mismatch_bucketing_out32", "SELECT orderkey, comment, orderkey, comment, orderkey, comment FROM orders"); assertUpdate("DROP TABLE IF EXISTS test_mismatch_bucketing_out32"); - assertUpdate(withMismatchOptimization, writeToTableWithMoreBuckets, 15000, assertRemoteExchangesCount(2)); + assertUpdate(withMismatchOptimization, writeToTableWithMoreBuckets, 15000, assertRemoteExchangesCount(5)); assertQuery("SELECT * FROM test_mismatch_bucketing_out32", "SELECT orderkey, comment, orderkey, comment, orderkey, comment FROM orders"); - assertUpdate(withMismatchOptimization, writeToTableWithFewerBuckets, 15000, assertRemoteExchangesCount(2)); + assertUpdate(withMismatchOptimization, writeToTableWithFewerBuckets, 15000, assertRemoteExchangesCount(5)); assertQuery("SELECT * FROM test_mismatch_bucketing_out8", "SELECT orderkey, comment, orderkey, comment, orderkey, comment FROM orders"); } finally { @@ -5912,7 +5920,7 @@ public void testBucketedSelect() @Language("SQL") String query = "SELECT count(value1) FROM test_bucketed_select GROUP BY key1"; @Language("SQL") String expectedQuery = "SELECT count(comment) FROM orders GROUP BY orderkey"; - assertQuery(planWithTableNodePartitioning, query, expectedQuery, assertRemoteExchangesCount(0)); + assertQuery(planWithTableNodePartitioning, query, expectedQuery, assertRemoteExchangesCount(1)); assertQuery(planWithoutTableNodePartitioning, query, expectedQuery, assertRemoteExchangesCount(1)); } finally { @@ -6112,14 +6120,11 @@ private Consumer assertRemoteExchangesCount(Session session, int expectedR .findAll() .size(); if (actualRemoteExchangesCount != expectedRemoteExchangesCount) { - Metadata metadata = getDistributedQueryRunner().getCoordinator().getMetadata(); - FunctionManager functionManager = getDistributedQueryRunner().getCoordinator().getFunctionManager(); - String formattedPlan = textLogicalPlan(plan.getRoot(), plan.getTypes(), metadata, functionManager, StatsAndCosts.empty(), session, 0, false); throw new AssertionError(format( "Expected [\n%s\n] remote exchanges but found [\n%s\n] remote exchanges. Actual plan is [\n\n%s\n]", expectedRemoteExchangesCount, actualRemoteExchangesCount, - formattedPlan)); + formatPlan(session, plan))); } }; } @@ -6705,26 +6710,26 @@ public void testAnalyzePartitionedTable() "(null, null, null, null, 4.0, null, null)"); assertQuery(format("SHOW STATS FOR (SELECT * FROM %s WHERE p_varchar = 'e1' AND p_bigint = 9)", tableName), "SELECT * FROM VALUES " + - "('c_boolean', 0.0, 0.0, 1.0, null, null, null), " + - "('c_bigint', 0.0, 0.0, 1.0, null, null, null), " + - "('c_double', 0.0, 0.0, 1.0, null, null, null), " + - "('c_timestamp', 0.0, 0.0, 1.0, null, null, null), " + - "('c_varchar', 0.0, 0.0, 1.0, null, null, null), " + - "('c_varbinary', 0.0, 0.0, 1.0, null, null, null), " + - "('p_varchar', 0.0, 0.0, 1.0, null, null, null), " + - "('p_bigint', 0.0, 0.0, 1.0, null, null, null), " + - "(null, null, null, null, 0.0, null, null)"); + "('c_boolean', null, null, null, null, null, null), " + + "('c_bigint', null, null, null, null, null, null), " + + "('c_double', null, null, null, null, null, null), " + + "('c_timestamp', null, null, null, null, null, null), " + + "('c_varchar', null, null, null, null, null, null), " + + "('c_varbinary', null, null, null, null, null, null), " + + "('p_varchar', null, 1.0, 0.0, null, null, null), " + + "('p_bigint', null, 1.0, 0.0, null, 9, 9), " + + "(null, null, null, null, null, null, null)"); assertQuery(format("SHOW STATS FOR (SELECT * FROM %s WHERE p_varchar = 'e2' AND p_bigint = 9)", tableName), "SELECT * FROM VALUES " + - "('c_boolean', 0.0, 0.0, 1.0, null, null, null), " + - "('c_bigint', 0.0, 0.0, 1.0, null, null, null), " + - "('c_double', 0.0, 0.0, 1.0, null, null, null), " + - "('c_timestamp', 0.0, 0.0, 1.0, null, null, null), " + - "('c_varchar', 0.0, 0.0, 1.0, null, null, null), " + - "('c_varbinary', 0.0, 0.0, 1.0, null, null, null), " + - "('p_varchar', 0.0, 0.0, 1.0, null, null, null), " + - "('p_bigint', 0.0, 0.0, 1.0, null, null, null), " + - "(null, null, null, null, 0.0, null, null)"); + "('c_boolean', null, null, null, null, null, null), " + + "('c_bigint', null, null, null, null, null, null), " + + "('c_double', null, null, null, null, null, null), " + + "('c_timestamp', null, null, null, null, null, null), " + + "('c_varchar', null, null, null, null, null, null), " + + "('c_varbinary', null, null, null, null, null, null), " + + "('p_varchar', null, 1.0, 0.0, null, null, null), " + + "('p_bigint', null, 1.0, 0.0, null, 9, 9), " + + "(null, null, null, null, null, null, null)"); // Run analyze on the whole table assertUpdate("ANALYZE " + tableName, 16); @@ -6776,26 +6781,26 @@ public void testAnalyzePartitionedTable() "(null, null, null, null, 4.0, null, null)"); assertQuery(format("SHOW STATS FOR (SELECT * FROM %s WHERE p_varchar = 'e1' AND p_bigint = 9)", tableName), "SELECT * FROM VALUES " + - "('c_boolean', 0.0, 0.0, 1.0, null, null, null), " + - "('c_bigint', 0.0, 0.0, 1.0, null, null, null), " + - "('c_double', 0.0, 0.0, 1.0, null, null, null), " + - "('c_timestamp', 0.0, 0.0, 1.0, null, null, null), " + - "('c_varchar', 0.0, 0.0, 1.0, null, null, null), " + - "('c_varbinary', 0.0, 0.0, 1.0, null, null, null), " + - "('p_varchar', 0.0, 0.0, 1.0, null, null, null), " + - "('p_bigint', 0.0, 0.0, 1.0, null, null, null), " + - "(null, null, null, null, 0.0, null, null)"); + "('c_boolean', null, null, null, null, null, null), " + + "('c_bigint', null, null, null, null, null, null), " + + "('c_double', null, null, null, null, null, null), " + + "('c_timestamp', null, null, null, null, null, null), " + + "('c_varchar', null, null, null, null, null, null), " + + "('c_varbinary', null, null, null, null, null, null), " + + "('p_varchar', null, 1.0, 0.0, null, null, null), " + + "('p_bigint', null, 1.0, 0.0, null, 9, 9), " + + "(null, null, null, null, null, null, null)"); assertQuery(format("SHOW STATS FOR (SELECT * FROM %s WHERE p_varchar = 'e2' AND p_bigint = 9)", tableName), "SELECT * FROM VALUES " + - "('c_boolean', 0.0, 0.0, 1.0, null, null, null), " + - "('c_bigint', 0.0, 0.0, 1.0, null, null, null), " + - "('c_double', 0.0, 0.0, 1.0, null, null, null), " + - "('c_timestamp', 0.0, 0.0, 1.0, null, null, null), " + - "('c_varchar', 0.0, 0.0, 1.0, null, null, null), " + - "('c_varbinary', 0.0, 0.0, 1.0, null, null, null), " + - "('p_varchar', 0.0, 0.0, 1.0, null, null, null), " + - "('p_bigint', 0.0, 0.0, 1.0, null, null, null), " + - "(null, null, null, null, 0.0, null, null)"); + "('c_boolean', null, null, null, null, null, null), " + + "('c_bigint', null, null, null, null, null, null), " + + "('c_double', null, null, null, null, null, null), " + + "('c_timestamp', null, null, null, null, null, null), " + + "('c_varchar', null, null, null, null, null, null), " + + "('c_varbinary', null, null, null, null, null, null), " + + "('p_varchar', null, 1.0, 0.0, null, null, null), " + + "('p_bigint', null, 1.0, 0.0, null, 9, 9), " + + "(null, null, null, null, null, null, null)"); // Drop the partitioned test table assertUpdate("DROP TABLE " + tableName); @@ -6880,27 +6885,27 @@ public void testAnalyzePartitionedTableWithColumnSubset() assertQuery( format("SHOW STATS FOR (SELECT * FROM %s WHERE p_varchar = 'e1' AND p_bigint = 9)", tableName), "SELECT * FROM VALUES " + - "('c_boolean', 0.0, 0.0, 1.0, null, null, null), " + - "('c_bigint', 0.0, 0.0, 1.0, null, null, null), " + - "('c_double', 0.0, 0.0, 1.0, null, null, null), " + - "('c_timestamp', 0.0, 0.0, 1.0, null, null, null), " + - "('c_varchar', 0.0, 0.0, 1.0, null, null, null), " + - "('c_varbinary', 0.0, 0.0, 1.0, null, null, null), " + - "('p_varchar', 0.0, 0.0, 1.0, null, null, null), " + - "('p_bigint', 0.0, 0.0, 1.0, null, null, null), " + - "(null, null, null, null, 0.0, null, null)"); + "('c_boolean', null, null, null, null, null, null), " + + "('c_bigint', null, null, null, null, null, null), " + + "('c_double', null, null, null, null, null, null), " + + "('c_timestamp', null, null, null, null, null, null), " + + "('c_varchar', null, null, null, null, null, null), " + + "('c_varbinary', null, null, null, null, null, null), " + + "('p_varchar', null, 1.0, 0.0, null, null, null), " + + "('p_bigint', null, 1.0, 0.0, null, 9, 9), " + + "(null, null, null, null, null, null, null)"); assertQuery( format("SHOW STATS FOR (SELECT * FROM %s WHERE p_varchar = 'e2' AND p_bigint = 9)", tableName), "SELECT * FROM VALUES " + - "('c_boolean', 0.0, 0.0, 1.0, null, null, null), " + - "('c_bigint', 0.0, 0.0, 1.0, null, null, null), " + - "('c_double', 0.0, 0.0, 1.0, null, null, null), " + - "('c_timestamp', 0.0, 0.0, 1.0, null, null, null), " + - "('c_varchar', 0.0, 0.0, 1.0, null, null, null), " + - "('c_varbinary', 0.0, 0.0, 1.0, null, null, null), " + - "('p_varchar', 0.0, 0.0, 1.0, null, null, null), " + - "('p_bigint', 0.0, 0.0, 1.0, null, null, null), " + - "(null, null, null, null, 0.0, null, null)"); + "('c_boolean', null, null, null, null, null, null), " + + "('c_bigint', null, null, null, null, null, null), " + + "('c_double', null, null, null, null, null, null), " + + "('c_timestamp', null, null, null, null, null, null), " + + "('c_varchar', null, null, null, null, null, null), " + + "('c_varbinary', null, null, null, null, null, null), " + + "('p_varchar', null, 1.0, 0.0, null, null, null), " + + "('p_bigint', null, 1.0, 0.0, null, 9, 9), " + + "(null, null, null, null, null, null, null)"); // Run analyze again, this time on 2 new columns (for all partitions); the previously computed stats // should be preserved @@ -6958,27 +6963,27 @@ public void testAnalyzePartitionedTableWithColumnSubset() assertQuery( format("SHOW STATS FOR (SELECT * FROM %s WHERE p_varchar = 'e1' AND p_bigint = 9)", tableName), "SELECT * FROM VALUES " + - "('c_boolean', 0.0, 0.0, 1.0, null, null, null), " + - "('c_bigint', 0.0, 0.0, 1.0, null, null, null), " + - "('c_double', 0.0, 0.0, 1.0, null, null, null), " + - "('c_timestamp', 0.0, 0.0, 1.0, null, null, null), " + - "('c_varchar', 0.0, 0.0, 1.0, null, null, null), " + - "('c_varbinary', 0.0, 0.0, 1.0, null, null, null), " + - "('p_varchar', 0.0, 0.0, 1.0, null, null, null), " + - "('p_bigint', 0.0, 0.0, 1.0, null, null, null), " + - "(null, null, null, null, 0.0, null, null)"); + "('c_boolean', null, null, null, null, null, null), " + + "('c_bigint', null, null, null, null, null, null), " + + "('c_double', null, null, null, null, null, null), " + + "('c_timestamp', null, null, null, null, null, null), " + + "('c_varchar', null, null, null, null, null, null), " + + "('c_varbinary', null, null, null, null, null, null), " + + "('p_varchar', null, 1.0, 0.0, null, null, null), " + + "('p_bigint', null, 1.0, 0.0, null, 9, 9), " + + "(null, null, null, null, null, null, null)"); assertQuery( format("SHOW STATS FOR (SELECT * FROM %s WHERE p_varchar = 'e2' AND p_bigint = 9)", tableName), "SELECT * FROM VALUES " + - "('c_boolean', 0.0, 0.0, 1.0, null, null, null), " + - "('c_bigint', 0.0, 0.0, 1.0, null, null, null), " + - "('c_double', 0.0, 0.0, 1.0, null, null, null), " + - "('c_timestamp', 0.0, 0.0, 1.0, null, null, null), " + - "('c_varchar', 0.0, 0.0, 1.0, null, null, null), " + - "('c_varbinary', 0.0, 0.0, 1.0, null, null, null), " + - "('p_varchar', 0.0, 0.0, 1.0, null, null, null), " + - "('p_bigint', 0.0, 0.0, 1.0, null, null, null), " + - "(null, null, null, null, 0.0, null, null)"); + "('c_boolean', null, null, null, null, null, null), " + + "('c_bigint', null, null, null, null, null, null), " + + "('c_double', null, null, null, null, null, null), " + + "('c_timestamp', null, null, null, null, null, null), " + + "('c_varchar', null, null, null, null, null, null), " + + "('c_varbinary', null, null, null, null, null, null), " + + "('p_varchar', null, 1.0, 0.0, null, null, null), " + + "('p_bigint', null, 1.0, 0.0, null, 9, 9), " + + "(null, null, null, null, null, null, null)"); assertUpdate("DROP TABLE " + tableName); } @@ -7244,26 +7249,26 @@ public void testDropStatsPartitionedTable() "(null, null, null, null, 4.0, null, null)"); assertQuery(format("SHOW STATS FOR (SELECT * FROM %s WHERE p_varchar = 'e1' AND p_bigint = 9)", tableName), "SELECT * FROM VALUES " + - "('c_boolean', 0.0, 0.0, 1.0, null, null, null), " + - "('c_bigint', 0.0, 0.0, 1.0, null, null, null), " + - "('c_double', 0.0, 0.0, 1.0, null, null, null), " + - "('c_timestamp', 0.0, 0.0, 1.0, null, null, null), " + - "('c_varchar', 0.0, 0.0, 1.0, null, null, null), " + - "('c_varbinary', 0.0, 0.0, 1.0, null, null, null), " + - "('p_varchar', 0.0, 0.0, 1.0, null, null, null), " + - "('p_bigint', 0.0, 0.0, 1.0, null, null, null), " + - "(null, null, null, null, 0.0, null, null)"); + "('c_boolean', null, null, null, null, null, null), " + + "('c_bigint', null, null, null, null, null, null), " + + "('c_double', null, null, null, null, null, null), " + + "('c_timestamp', null, null, null, null, null, null), " + + "('c_varchar', null, null, null, null, null, null), " + + "('c_varbinary', null, null, null, null, null, null), " + + "('p_varchar', null, 1.0, 0.0, null, null, null), " + + "('p_bigint', null, 1.0, 0.0, null, 9, 9), " + + "(null, null, null, null, null, null, null)"); assertQuery(format("SHOW STATS FOR (SELECT * FROM %s WHERE p_varchar = 'e2' AND p_bigint = 9)", tableName), "SELECT * FROM VALUES " + - "('c_boolean', 0.0, 0.0, 1.0, null, null, null), " + - "('c_bigint', 0.0, 0.0, 1.0, null, null, null), " + - "('c_double', 0.0, 0.0, 1.0, null, null, null), " + - "('c_timestamp', 0.0, 0.0, 1.0, null, null, null), " + - "('c_varchar', 0.0, 0.0, 1.0, null, null, null), " + - "('c_varbinary', 0.0, 0.0, 1.0, null, null, null), " + - "('p_varchar', 0.0, 0.0, 1.0, null, null, null), " + - "('p_bigint', 0.0, 0.0, 1.0, null, null, null), " + - "(null, null, null, null, 0.0, null, null)"); + "('c_boolean', null, null, null, null, null, null), " + + "('c_bigint', null, null, null, null, null, null), " + + "('c_double', null, null, null, null, null, null), " + + "('c_timestamp', null, null, null, null, null, null), " + + "('c_varchar', null, null, null, null, null, null), " + + "('c_varbinary', null, null, null, null, null, null), " + + "('p_varchar', null, 1.0, 0.0, null, null, null), " + + "('p_bigint', null, 1.0, 0.0, null, 9, 9), " + + "(null, null, null, null, null, null, null)"); // Drop stats for 2 partitions assertUpdate(format("CALL system.drop_stats('%s', '%s', ARRAY[ARRAY['p2', '7'], ARRAY['p3', '8']])", TPCH_SCHEMA, tableName)); @@ -7323,26 +7328,26 @@ public void testDropStatsPartitionedTable() "(null, null, null, null, null, null, null)"); assertQuery(format("SHOW STATS FOR (SELECT * FROM %s WHERE p_varchar = 'e1' AND p_bigint = 9)", tableName), "SELECT * FROM VALUES " + - "('c_boolean', 0.0, 0.0, 1.0, null, null, null), " + - "('c_bigint', 0.0, 0.0, 1.0, null, null, null), " + - "('c_double', 0.0, 0.0, 1.0, null, null, null), " + - "('c_timestamp', 0.0, 0.0, 1.0, null, null, null), " + - "('c_varchar', 0.0, 0.0, 1.0, null, null, null), " + - "('c_varbinary', 0.0, 0.0, 1.0, null, null, null), " + - "('p_varchar', 0.0, 0.0, 1.0, null, null, null), " + - "('p_bigint', 0.0, 0.0, 1.0, null, null, null), " + - "(null, null, null, null, 0.0, null, null)"); + "('c_boolean', null, null, null, null, null, null), " + + "('c_bigint', null, null, null, null, null, null), " + + "('c_double', null, null, null, null, null, null), " + + "('c_timestamp', null, null, null, null, null, null), " + + "('c_varchar', null, null, null, null, null, null), " + + "('c_varbinary', null, null, null, null, null, null), " + + "('p_varchar', null, 1.0, 0.0, null, null, null), " + + "('p_bigint', null, 1.0, 0.0, null, 9, 9), " + + "(null, null, null, null, null, null, null)"); assertQuery(format("SHOW STATS FOR (SELECT * FROM %s WHERE p_varchar = 'e2' AND p_bigint = 9)", tableName), "SELECT * FROM VALUES " + - "('c_boolean', 0.0, 0.0, 1.0, null, null, null), " + - "('c_bigint', 0.0, 0.0, 1.0, null, null, null), " + - "('c_double', 0.0, 0.0, 1.0, null, null, null), " + - "('c_timestamp', 0.0, 0.0, 1.0, null, null, null), " + - "('c_varchar', 0.0, 0.0, 1.0, null, null, null), " + - "('c_varbinary', 0.0, 0.0, 1.0, null, null, null), " + - "('p_varchar', 0.0, 0.0, 1.0, null, null, null), " + - "('p_bigint', 0.0, 0.0, 1.0, null, null, null), " + - "(null, null, null, null, 0.0, null, null)"); + "('c_boolean', null, null, null, null, null, null), " + + "('c_bigint', null, null, null, null, null, null), " + + "('c_double', null, null, null, null, null, null), " + + "('c_timestamp', null, null, null, null, null, null), " + + "('c_varchar', null, null, null, null, null, null), " + + "('c_varbinary', null, null, null, null, null, null), " + + "('p_varchar', null, 1.0, 0.0, null, null, null), " + + "('p_bigint', null, 1.0, 0.0, null, 9, 9), " + + "(null, null, null, null, null, null, null)"); // Drop stats for the entire table assertUpdate(format("CALL system.drop_stats('%s', '%s')", TPCH_SCHEMA, tableName)); @@ -7400,7 +7405,7 @@ public void testDropStatsPartitionedTable() "('c_varchar', null, null, null, null, null, null), " + "('c_varbinary', null, null, null, null, null, null), " + "('p_varchar', null, 1.0, 0.0, null, null, null), " + - "('p_bigint', null, 1.0, 0.0, null, '9', '9'), " + + "('p_bigint', null, 1.0, 0.0, null, 9, 9), " + "(null, null, null, null, null, null, null)"); assertQuery(format("SHOW STATS FOR (SELECT * FROM %s WHERE p_varchar = 'e2' AND p_bigint = 9)", tableName), "SELECT * FROM VALUES " + @@ -7411,7 +7416,7 @@ public void testDropStatsPartitionedTable() "('c_varchar', null, null, null, null, null, null), " + "('c_varbinary', null, null, null, null, null, null), " + "('p_varchar', null, 1.0, 0.0, null, null, null), " + - "('p_bigint', null, 1.0, 0.0, null, '9', '9'), " + + "('p_bigint', null, 1.0, 0.0, null, 9, 9), " + "(null, null, null, null, null, null, null)"); // All table stats are gone @@ -7622,10 +7627,28 @@ public void testCreateAvroTableWithSchemaUrl() throws Exception { String tableName = "test_create_avro_table_with_schema_url"; - File schemaFile = createAvroSchemaFile(); + TrinoFileSystem fileSystem = getTrinoFileSystem(); + Location tempDir = Location.of("local:///temp_" + UUID.randomUUID()); + fileSystem.createDirectory(tempDir); + String schema = + """ + { + "namespace": "io.trino.test", + "name": "camelCase", + "type": "record", + "fields": [ + { "name":"stringCol", "type":"string" }, + { "name":"a", "type":"int" } + ] + }\ + """; + Location schemaFile = tempDir.appendPath("avro_camelCamelCase_col.avsc"); + try (OutputStream out = fileSystem.newOutputFile(schemaFile).create()) { + out.write(schema.getBytes(UTF_8)); + } - String createTableSql = getAvroCreateTableSql(tableName, schemaFile.getAbsolutePath()); - String expectedShowCreateTable = getAvroCreateTableSql(tableName, schemaFile.toURI().toString()); + String createTableSql = getAvroCreateTableSql(tableName, schemaFile.toString()); + String expectedShowCreateTable = getAvroCreateTableSql(tableName, schemaFile.toString()); assertUpdate(createTableSql); @@ -7635,7 +7658,7 @@ public void testCreateAvroTableWithSchemaUrl() } finally { assertUpdate("DROP TABLE " + tableName); - verify(schemaFile.delete(), "cannot delete temporary file: %s", schemaFile); + fileSystem.deleteDirectory(tempDir); } } @@ -7650,9 +7673,26 @@ protected void testAlterAvroTableWithSchemaUrl(boolean renameColumn, boolean add throws Exception { String tableName = "test_alter_avro_table_with_schema_url"; - File schemaFile = createAvroSchemaFile(); + TrinoFileSystem fileSystem = getTrinoFileSystem(); + Location tempDir = Location.of("local:///temp_" + UUID.randomUUID()); + fileSystem.createDirectory(tempDir); + String schema = + """ + { + "namespace": "io.trino.test", + "name": "single_column", + "type": "record", + "fields": [ + { "name": "string_col", "type":"string" } + ] + }\ + """; + Location schemaFile = tempDir.appendPath("avro_single_column.avsc"); + try (OutputStream out = fileSystem.newOutputFile(schemaFile).create()) { + out.write(schema.getBytes(UTF_8)); + } - assertUpdate(getAvroCreateTableSql(tableName, schemaFile.getAbsolutePath())); + assertUpdate(getAvroCreateTableSql(tableName, schemaFile.toString())); try { if (renameColumn) { @@ -7667,7 +7707,7 @@ protected void testAlterAvroTableWithSchemaUrl(boolean renameColumn, boolean add } finally { assertUpdate("DROP TABLE " + tableName); - verify(schemaFile.delete(), "cannot delete temporary file: %s", schemaFile); + fileSystem.deleteDirectory(tempDir); } } @@ -7687,10 +7727,12 @@ private String getAvroCreateTableSql(String tableName, String schemaFile) schemaFile); } - private static File createAvroSchemaFile() + private Location createAvroSchemaFile() throws Exception { - File schemaFile = File.createTempFile("avro_single_column-", ".avsc"); + TrinoFileSystem fileSystem = getTrinoFileSystem(); + Location tempDir = Location.of("local:///temp_" + UUID.randomUUID()); + fileSystem.createDirectory(tempDir); String schema = "{\n" + " \"namespace\": \"io.trino.test\",\n" + " \"name\": \"single_column\",\n" + @@ -7698,7 +7740,10 @@ private static File createAvroSchemaFile() " \"fields\": [\n" + " { \"name\":\"string_col\", \"type\":\"string\" }\n" + "]}"; - writeString(schemaFile.toPath(), schema); + Location schemaFile = tempDir.appendPath("avro_camelCamelCase_col.avsc"); + try (OutputStream out = fileSystem.newOutputFile(schemaFile).create()) { + out.write(schema.getBytes(UTF_8)); + } return schemaFile; } @@ -7806,16 +7851,10 @@ public void testUseSortedProperties() public void testCreateTableWithCompressionCodec(HiveCompressionCodec compressionCodec) { testWithAllStorageFormats((session, hiveStorageFormat) -> { - if (hiveStorageFormat == HiveStorageFormat.PARQUET && compressionCodec == HiveCompressionCodec.LZ4) { - // TODO (https://github.com/trinodb/trino/issues/9142) Support LZ4 compression with native Parquet writer - assertThatThrownBy(() -> testCreateTableWithCompressionCodec(session, hiveStorageFormat, compressionCodec)) - .hasMessage("Unsupported codec: LZ4"); - return; - } - + // TODO (https://github.com/trinodb/trino/issues/9142) Support LZ4 compression with native Parquet writer if (!isSupportedCodec(hiveStorageFormat, compressionCodec)) { assertThatThrownBy(() -> testCreateTableWithCompressionCodec(session, hiveStorageFormat, compressionCodec)) - .hasMessage("Compression codec " + compressionCodec + " not supported for " + hiveStorageFormat); + .hasMessage("Compression codec " + compressionCodec + " not supported for " + hiveStorageFormat.humanName()); return; } testCreateTableWithCompressionCodec(session, hiveStorageFormat, compressionCodec); @@ -7824,7 +7863,7 @@ public void testCreateTableWithCompressionCodec(HiveCompressionCodec compression private boolean isSupportedCodec(HiveStorageFormat storageFormat, HiveCompressionCodec codec) { - if (storageFormat == HiveStorageFormat.AVRO && codec == HiveCompressionCodec.LZ4) { + if ((storageFormat == HiveStorageFormat.AVRO || storageFormat == HiveStorageFormat.PARQUET) && codec == HiveCompressionCodec.LZ4) { return false; } return true; @@ -8684,13 +8723,22 @@ public void testCollidingMixedCaseProperty() @Test public void testSelectWithShortZoneId() + throws IOException { - String resourceLocation = getPathFromClassPathResource("with_short_zone_id/data"); + URL resourceLocation = Resources.getResource("with_short_zone_id/data/data.orc"); + + TrinoFileSystem fileSystem = getTrinoFileSystem(); + Location tempDir = Location.of("local:///temp_" + UUID.randomUUID()); + fileSystem.createDirectory(tempDir); + Location dataFile = tempDir.appendPath("data.orc"); + try (OutputStream out = fileSystem.newOutputFile(dataFile).create()) { + Resources.copy(resourceLocation, out); + } try (TestTable testTable = new TestTable( getQueryRunner()::execute, "test_select_with_short_zone_id_", - "(id INT, firstName VARCHAR, lastName VARCHAR) WITH (external_location = '%s')".formatted(resourceLocation))) { + "(id INT, firstName VARCHAR, lastName VARCHAR) WITH (external_location = '%s')".formatted(tempDir))) { assertThatThrownBy(() -> query("SELECT * FROM %s".formatted(testTable.getName()))) .hasMessageMatching(".*Failed to read ORC file: .*") .hasStackTraceContaining("Unknown time-zone ID: EST"); @@ -8934,6 +8982,11 @@ private String getTableLocation(String tableName) return (String) computeScalar("SELECT DISTINCT regexp_replace(\"$path\", '/[^/]*$', '') FROM " + tableName); } + private TrinoFileSystem getTrinoFileSystem() + { + return getConnectorService(getQueryRunner(), TrinoFileSystemFactory.class).create(ConnectorIdentity.ofUser("test")); + } + @Override protected boolean supportsPhysicalPushdown() { diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/BaseTestHiveOnDataLake.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/BaseTestHiveOnDataLake.java index 9c21f466b138..a7f5897c9261 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/BaseTestHiveOnDataLake.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/BaseTestHiveOnDataLake.java @@ -15,9 +15,13 @@ import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableSet; import io.airlift.units.DataSize; import io.trino.Session; +import io.trino.plugin.hive.containers.Hive3MinioDataLake; import io.trino.plugin.hive.containers.HiveMinioDataLake; +import io.trino.plugin.hive.metastore.Column; +import io.trino.plugin.hive.metastore.HiveColumnStatistics; import io.trino.plugin.hive.metastore.HiveMetastore; import io.trino.plugin.hive.metastore.Partition; import io.trino.plugin.hive.metastore.PartitionWithStatistics; @@ -34,6 +38,7 @@ import io.trino.testing.minio.MinioClient; import io.trino.testing.sql.TestTable; import org.intellij.lang.annotations.Language; +import org.testng.annotations.AfterClass; import org.testng.annotations.BeforeClass; import org.testng.annotations.DataProvider; import org.testng.annotations.Test; @@ -55,6 +60,7 @@ import static io.airlift.slice.Slices.utf8Slice; import static io.airlift.units.DataSize.Unit.MEGABYTE; import static io.trino.plugin.hive.TestingThriftHiveMetastoreBuilder.testingThriftHiveMetastoreBuilder; +import static io.trino.plugin.hive.metastore.MetastoreUtil.getHiveBasicStatistics; import static io.trino.spi.type.VarcharType.VARCHAR; import static io.trino.testing.MaterializedResult.resultBuilder; import static io.trino.testing.TestingNames.randomNameSuffix; @@ -62,9 +68,9 @@ import static java.nio.charset.StandardCharsets.UTF_8; import static java.time.temporal.ChronoUnit.DAYS; import static java.time.temporal.ChronoUnit.MINUTES; -import static java.util.Objects.requireNonNull; import static java.util.regex.Pattern.quote; import static java.util.stream.Collectors.joining; +import static java.util.stream.Collectors.toSet; import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatThrownBy; import static org.testng.Assert.assertEquals; @@ -74,16 +80,16 @@ public abstract class BaseTestHiveOnDataLake { private static final String HIVE_TEST_SCHEMA = "hive_datalake"; private static final DataSize HIVE_S3_STREAMING_PART_SIZE = DataSize.of(5, MEGABYTE); + private final String hiveHadoopImage; private String bucketName; private HiveMinioDataLake hiveMinioDataLake; private HiveMetastore metastoreClient; - private final String hiveHadoopImage; - - public BaseTestHiveOnDataLake(String hiveHadoopImage) + public BaseTestHiveOnDataLake(String bucketName, String hiveHadoopImage) { - this.hiveHadoopImage = requireNonNull(hiveHadoopImage, "hiveHadoopImage is null"); + this.bucketName = bucketName; + this.hiveHadoopImage = hiveHadoopImage; } @Override @@ -91,13 +97,12 @@ protected QueryRunner createQueryRunner() throws Exception { this.bucketName = "test-hive-insert-overwrite-" + randomNameSuffix(); - this.hiveMinioDataLake = closeAfterClass( - new HiveMinioDataLake(bucketName, hiveHadoopImage)); + this.hiveMinioDataLake = new Hive3MinioDataLake(bucketName, hiveHadoopImage); this.hiveMinioDataLake.start(); this.metastoreClient = new BridgingHiveMetastore( testingThriftHiveMetastoreBuilder() .metastoreClient(this.hiveMinioDataLake.getHiveHadoop().getHiveMetastoreEndpoint()) - .build()); + .build(this::closeAfterClass)); return S3HiveQueryRunner.builder(hiveMinioDataLake) .setHiveProperties( ImmutableMap.builder() @@ -107,10 +112,10 @@ protected QueryRunner createQueryRunner() .put("hive.metastore-cache-ttl", "1d") .put("hive.metastore-refresh-interval", "1d") // This is required to reduce memory pressure to test writing large files - .put("hive.s3.streaming.part-size", HIVE_S3_STREAMING_PART_SIZE.toString()) + .put("s3.streaming.part-size", HIVE_S3_STREAMING_PART_SIZE.toString()) // This is required to enable AWS Athena partition projection .put("hive.partition-projection-enabled", "true") - .put("hive.s3select-pushdown.experimental-textfile-pushdown-enabled", "true") + .put("hive.hive-views.enabled", "true") .buildOrThrow()) .build(); } @@ -124,6 +129,13 @@ public void setUp() bucketName)); } + @AfterClass(alwaysRun = true) + public void tearDown() + throws Exception + { + hiveMinioDataLake.close(); + } + @Test public void testInsertOverwriteInTransaction() { @@ -211,7 +223,7 @@ public void testInsertOverwritePartitionedAndBucketedExternalTable() assertOverwritePartition(externalTableName); } - @Test + @Test(enabled = false) public void testFlushPartitionCache() { String tableName = "nation_" + randomNameSuffix(); @@ -229,7 +241,7 @@ public void testFlushPartitionCache() partitionColumn)); } - @Test + @Test(enabled = false) public void testFlushPartitionCacheWithDeprecatedPartitionParams() { String tableName = "nation_" + randomNameSuffix(); @@ -585,7 +597,7 @@ public void testIntegerPartitionProjectionOnVarcharColumnWithDigitsAlignCreatedO testIntegerPartitionProjectionOnVarcharColumnWithDigitsAlign(tableName); } - @Test + @Test(enabled = false) public void testIntegerPartitionProjectionOnVarcharColumnWithDigitsAlignCreatedOnHive() { String tableName = "nation_" + randomNameSuffix(); @@ -645,7 +657,7 @@ private void testIntegerPartitionProjectionOnVarcharColumnWithDigitsAlign(String "VALUES ('POLAND_1'), ('POLAND_2'), ('CZECH_1'), ('CZECH_2')"); } - @Test + @Test(enabled = false) public void testIntegerPartitionProjectionOnIntegerColumnWithInterval() { String tableName = getRandomTestTableName(); @@ -706,7 +718,7 @@ public void testIntegerPartitionProjectionOnIntegerColumnWithInterval() "VALUES ('POLAND_1'), ('POLAND_2'), ('CZECH_1'), ('CZECH_2')"); } - @Test + @Test(enabled = false) public void testIntegerPartitionProjectionOnIntegerColumnWithDefaults() { String tableName = getRandomTestTableName(); @@ -765,7 +777,7 @@ public void testIntegerPartitionProjectionOnIntegerColumnWithDefaults() "VALUES ('POLAND_1'), ('POLAND_2'), ('CZECH_1'), ('CZECH_2')"); } - @Test + @Test(enabled = false) public void testDatePartitionProjectionOnDateColumnWithDefaults() { String tableName = "nation_" + randomNameSuffix(); @@ -840,7 +852,7 @@ public void testDatePartitionProjectionOnDateColumnWithDefaults() "VALUES ('POLAND_1'), ('POLAND_2'), ('CZECH_1'), ('CZECH_2')"); } - @Test + @Test(enabled = false) public void testDatePartitionProjectionOnTimestampColumnWithInterval() { String tableName = getRandomTestTableName(); @@ -914,7 +926,7 @@ public void testDatePartitionProjectionOnTimestampColumnWithInterval() "VALUES ('POLAND_1'), ('POLAND_2'), ('CZECH_1'), ('CZECH_2')"); } - @Test + @Test(enabled = false) public void testDatePartitionProjectionOnTimestampColumnWithIntervalExpressionCreatedOnTrino() { String tableName = getRandomTestTableName(); @@ -955,7 +967,7 @@ public void testDatePartitionProjectionOnTimestampColumnWithIntervalExpressionCr testDatePartitionProjectionOnTimestampColumnWithIntervalExpression(tableName, dateProjectionFormat); } - @Test + @Test(enabled = false) public void testDatePartitionProjectionOnTimestampColumnWithIntervalExpressionCreatedOnHive() { String tableName = getRandomTestTableName(); @@ -1013,7 +1025,7 @@ private void testDatePartitionProjectionOnTimestampColumnWithIntervalExpression( "VALUES ('POLAND_2'), ('CZECH_1'), ('CZECH_2')"); } - @Test + @Test(enabled = false) public void testDatePartitionProjectionOnVarcharColumnWithHoursInterval() { String tableName = getRandomTestTableName(); @@ -1087,7 +1099,7 @@ public void testDatePartitionProjectionOnVarcharColumnWithHoursInterval() "VALUES ('POLAND_1'), ('POLAND_2'), ('CZECH_1'), ('CZECH_2')"); } - @Test + @Test(enabled = false) public void testDatePartitionProjectionOnVarcharColumnWithDaysInterval() { String tableName = getRandomTestTableName(); @@ -1161,7 +1173,7 @@ public void testDatePartitionProjectionOnVarcharColumnWithDaysInterval() "VALUES ('POLAND_1'), ('POLAND_2'), ('CZECH_1'), ('CZECH_2')"); } - @Test + @Test(enabled = false) public void testDatePartitionProjectionOnVarcharColumnWithIntervalExpression() { String tableName = getRandomTestTableName(); @@ -1268,7 +1280,7 @@ public void testDatePartitionProjectionFormatTextWillNotCauseIntervalRequirement ")"); } - @Test + @Test(enabled = false) public void testInjectedPartitionProjectionOnVarcharColumn() { String tableName = getRandomTestTableName(); @@ -1337,7 +1349,7 @@ public void testPartitionProjectionInvalidTableProperties() ") WITH ( " + " partition_projection_enabled=true " + ")")) - .hasMessage("Partition projection can't be enabled when no partition columns are defined."); + .hasMessage("Partition projection cannot be enabled on a table that is not partitioned"); assertThatThrownBy(() -> getQueryRunner().execute( "CREATE TABLE " + getFullyQualifiedTestTableName("nation_" + randomNameSuffix()) + " ( " + @@ -1350,7 +1362,7 @@ public void testPartitionProjectionInvalidTableProperties() " partitioned_by=ARRAY['short_name1'], " + " partition_projection_enabled=true " + ")")) - .hasMessage("Partition projection can't be defined for non partition column: 'name'"); + .hasMessage("Partition projection cannot be defined for non-partition column: 'name'"); assertThatThrownBy(() -> getQueryRunner().execute( "CREATE TABLE " + getFullyQualifiedTestTableName("nation_" + randomNameSuffix()) + " ( " + @@ -1364,7 +1376,7 @@ public void testPartitionProjectionInvalidTableProperties() " partitioned_by=ARRAY['short_name1', 'short_name2'], " + " partition_projection_enabled=true " + ")")) - .hasMessage("Partition projection definition for column: 'short_name2' missing"); + .hasMessage("Column projection for column 'short_name2' failed. Projection type property missing"); assertThatThrownBy(() -> getQueryRunner().execute( "CREATE TABLE " + getFullyQualifiedTestTableName("nation_" + randomNameSuffix()) + " ( " + @@ -1427,7 +1439,7 @@ public void testPartitionProjectionInvalidTableProperties() " partition_projection_enabled=true " + ")")) .hasMessage("Column projection for column 'short_name1' failed. Property: 'partition_projection_range' needs to be a list of 2 valid dates formatted as 'yyyy-MM-dd HH' " + - "or '^\\s*NOW\\s*(([+-])\\s*([0-9]+)\\s*(DAY|HOUR|MINUTE|SECOND)S?\\s*)?$' that are sequential. Unparseable date: \"2001-01-01\""); + "or '^\\s*NOW\\s*(([+-])\\s*([0-9]+)\\s*(DAY|HOUR|MINUTE|SECOND)S?\\s*)?$' that are sequential: Unparseable date: \"2001-01-01\""); assertThatThrownBy(() -> getQueryRunner().execute( "CREATE TABLE " + getFullyQualifiedTestTableName("nation_" + randomNameSuffix()) + " ( " + @@ -1442,7 +1454,7 @@ public void testPartitionProjectionInvalidTableProperties() " partition_projection_enabled=true " + ")")) .hasMessage("Column projection for column 'short_name1' failed. Property: 'partition_projection_range' needs to be a list of 2 valid dates formatted as 'yyyy-MM-dd' " + - "or '^\\s*NOW\\s*(([+-])\\s*([0-9]+)\\s*(DAY|HOUR|MINUTE|SECOND)S?\\s*)?$' that are sequential. Unparseable date: \"NOW*3DAYS\""); + "or '^\\s*NOW\\s*(([+-])\\s*([0-9]+)\\s*(DAY|HOUR|MINUTE|SECOND)S?\\s*)?$' that are sequential: Unparseable date: \"NOW*3DAYS\""); assertThatThrownBy(() -> getQueryRunner().execute( "CREATE TABLE " + getFullyQualifiedTestTableName("nation_" + randomNameSuffix()) + " ( " + @@ -1502,7 +1514,7 @@ public void testPartitionProjectionInvalidTableProperties() ") WITH ( " + " partitioned_by=ARRAY['short_name1'] " + ")")) - .hasMessage("Columns ['short_name1'] projections are disallowed when partition projection property 'partition_projection_enabled' is missing"); + .hasMessage("Columns partition projection properties cannot be set when 'partition_projection_enabled' is not set"); // Verify that ignored flag is only interpreted for pre-existing tables where configuration is loaded from metastore. // It should not allow creating corrupted config via Trino. It's a kill switch to run away when we have compatibility issues. @@ -1524,7 +1536,7 @@ public void testPartitionProjectionInvalidTableProperties() "Interval defaults to 1 day or 1 month, respectively. Otherwise, interval is required"); } - @Test + @Test(enabled = false) public void testPartitionProjectionIgnore() { String tableName = "nation_" + randomNameSuffix(); @@ -1548,7 +1560,7 @@ public void testPartitionProjectionIgnore() // Expect invalid Partition Projection properties to fail assertThatThrownBy(() -> getQueryRunner().execute("SELECT * FROM " + fullyQualifiedTestTableName)) .hasMessage("Column projection for column 'date_time' failed. Property: 'partition_projection_range' needs to be a list of 2 valid dates formatted as 'yyyy-MM-dd HH' " + - "or '^\\s*NOW\\s*(([+-])\\s*([0-9]+)\\s*(DAY|HOUR|MINUTE|SECOND)S?\\s*)?$' that are sequential. Unparseable date: \"2001-01-01\""); + "or '^\\s*NOW\\s*(([+-])\\s*([0-9]+)\\s*(DAY|HOUR|MINUTE|SECOND)S?\\s*)?$' that are sequential: Unparseable date: \"2001-01-01\""); // Append kill switch table property to ignore Partition Projection properties hiveMinioDataLake.getHiveHadoop().runOnHive( @@ -1793,7 +1805,7 @@ public void testPartitionedTableExternalLocationOnTopOfTheBucket() assertUpdate("DROP TABLE " + tableName); } - @Test(dataProvider = "s3SelectFileFormats") + @Test(dataProvider = "s3SelectFileFormats", enabled = false) public void testS3SelectPushdown(String tableProperties) { Session usingAppendInserts = Session.builder(getSession()) @@ -1870,7 +1882,7 @@ public void testS3SelectPushdown(String tableProperties) } } - @Test(dataProvider = "s3SelectFileFormats") + @Test(dataProvider = "s3SelectFileFormats", enabled = false) public void testS3SelectOnDecimalColumnIsDisabled(String tableProperties) { Session usingAppendInserts = Session.builder(getSession()) @@ -1893,7 +1905,7 @@ public void testS3SelectOnDecimalColumnIsDisabled(String tableProperties) } } - @Test + @Test(enabled = false) public void testJsonS3SelectPushdownWithSpecialCharacters() { Session usingAppendInserts = Session.builder(getSession()) @@ -1917,7 +1929,7 @@ public void testJsonS3SelectPushdownWithSpecialCharacters() } } - @Test + @Test(enabled = false) public void testS3SelectExperimentalPushdown() { // Demonstrate correctness issues which have resulted in pushdown for TEXTFILE @@ -2114,19 +2126,22 @@ private void renamePartitionResourcesOutsideTrino(String tableName, String parti // Delete old partition and update metadata to point to location of new copy Table hiveTable = metastoreClient.getTable(HIVE_TEST_SCHEMA, tableName).get(); - Partition hivePartition = metastoreClient.getPartition(hiveTable, List.of(regionKey)).get(); - Map partitionStatistics = - metastoreClient.getPartitionStatistics(hiveTable, List.of(hivePartition)); + Partition partition = metastoreClient.getPartition(hiveTable, List.of(regionKey)).get(); + Map> partitionStatistics = metastoreClient.getPartitionColumnStatistics( + HIVE_TEST_SCHEMA, + tableName, + ImmutableSet.of(partitionName), + partition.getColumns().stream().map(Column::getName).collect(toSet())); metastoreClient.dropPartition(HIVE_TEST_SCHEMA, tableName, List.of(regionKey), true); metastoreClient.addPartitions(HIVE_TEST_SCHEMA, tableName, List.of( new PartitionWithStatistics( - Partition.builder(hivePartition) + Partition.builder(partition) .withStorage(builder -> builder.setLocation( - hivePartition.getStorage().getLocation() + renamedPartitionSuffix)) + partition.getStorage().getLocation() + renamedPartitionSuffix)) .build(), partitionName, - partitionStatistics.get(partitionName)))); + new PartitionStatistics(getHiveBasicStatistics(partition.getParameters()), partitionStatistics.get(partitionName))))); } protected void assertInsertFailure(String testTable, String expectedMessageRegExp) diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/HiveBenchmarkQueryRunner.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/HiveBenchmarkQueryRunner.java index d91e3629c2d0..3d8375fcc0d9 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/HiveBenchmarkQueryRunner.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/HiveBenchmarkQueryRunner.java @@ -76,7 +76,7 @@ public static LocalQueryRunner createLocalQueryRunner(File tempDir) Map hiveCatalogConfig = ImmutableMap.of("hive.max-split-size", "10GB"); - localQueryRunner.createCatalog("hive", new TestingHiveConnectorFactory(metastore), hiveCatalogConfig); + localQueryRunner.createCatalog("hive", new TestingHiveConnectorFactory(hiveDir.toPath(), Optional.of(metastore)), hiveCatalogConfig); localQueryRunner.execute("CREATE TABLE orders AS SELECT * FROM tpch.sf1.orders"); localQueryRunner.execute("CREATE TABLE lineitem AS SELECT * FROM tpch.sf1.lineitem"); diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/HiveFileSystemTestUtils.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/HiveFileSystemTestUtils.java deleted file mode 100644 index 4ae990494623..000000000000 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/HiveFileSystemTestUtils.java +++ /dev/null @@ -1,211 +0,0 @@ -/* - * 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.hive; - -import com.google.common.collect.ImmutableList; -import io.trino.plugin.hive.AbstractTestHive.HiveTransaction; -import io.trino.plugin.hive.AbstractTestHive.Transaction; -import io.trino.spi.connector.ColumnHandle; -import io.trino.spi.connector.ConnectorMetadata; -import io.trino.spi.connector.ConnectorPageSource; -import io.trino.spi.connector.ConnectorPageSourceProvider; -import io.trino.spi.connector.ConnectorSession; -import io.trino.spi.connector.ConnectorSplit; -import io.trino.spi.connector.ConnectorSplitManager; -import io.trino.spi.connector.ConnectorSplitSource; -import io.trino.spi.connector.ConnectorTableHandle; -import io.trino.spi.connector.DynamicFilter; -import io.trino.spi.connector.SchemaTableName; -import io.trino.spi.type.Type; -import io.trino.testing.MaterializedResult; -import io.trino.testing.MaterializedRow; - -import java.io.Closeable; -import java.io.IOException; -import java.util.List; -import java.util.stream.IntStream; - -import static com.google.common.base.Preconditions.checkArgument; -import static com.google.common.collect.ImmutableList.toImmutableList; -import static io.trino.plugin.hive.AbstractTestHive.getAllSplits; -import static io.trino.plugin.hive.AbstractTestHive.getSplits; -import static io.trino.plugin.hive.HiveTestUtils.getHiveSession; -import static io.trino.plugin.hive.HiveTestUtils.getTypes; -import static io.trino.testing.MaterializedResult.materializeSourceDataStream; - -public class HiveFileSystemTestUtils -{ - private HiveFileSystemTestUtils() {} - - public static MaterializedResult readTable(SchemaTableName tableName, HiveTransactionManager transactionManager, - HiveConfig config, ConnectorPageSourceProvider pageSourceProvider, - ConnectorSplitManager splitManager) - throws IOException - { - ConnectorMetadata metadata = null; - ConnectorSession session = null; - ConnectorSplitSource splitSource = null; - - try (Transaction transaction = newTransaction(transactionManager)) { - metadata = transaction.getMetadata(); - session = newSession(config); - - ConnectorTableHandle table = getTableHandle(metadata, tableName, session); - List columnHandles = ImmutableList.copyOf(metadata.getColumnHandles(session, table).values()); - - metadata.beginQuery(session); - splitSource = getSplits(splitManager, transaction, session, table); - - List allTypes = getTypes(columnHandles); - List dataTypes = getTypes(columnHandles.stream() - .filter(columnHandle -> !((HiveColumnHandle) columnHandle).isHidden()) - .collect(toImmutableList())); - MaterializedResult.Builder result = MaterializedResult.resultBuilder(session, dataTypes); - - List splits = getAllSplits(splitSource); - for (ConnectorSplit split : splits) { - try (ConnectorPageSource pageSource = pageSourceProvider.createPageSource( - transaction.getTransactionHandle(), - session, - split, - table, - columnHandles, - DynamicFilter.EMPTY)) { - MaterializedResult pageSourceResult = materializeSourceDataStream(session, pageSource, allTypes); - for (MaterializedRow row : pageSourceResult.getMaterializedRows()) { - Object[] dataValues = IntStream.range(0, row.getFieldCount()) - .filter(channel -> !((HiveColumnHandle) columnHandles.get(channel)).isHidden()) - .mapToObj(row::getField) - .toArray(); - result.row(dataValues); - } - } - } - return result.build(); - } - finally { - cleanUpQuery(metadata, session); - closeQuietly(splitSource); - } - } - - public static ConnectorTableHandle getTableHandle(ConnectorMetadata metadata, SchemaTableName tableName, ConnectorSession session) - { - ConnectorTableHandle handle = metadata.getTableHandle(session, tableName); - checkArgument(handle != null, "table not found: %s", tableName); - return handle; - } - - public static ConnectorSession newSession(HiveConfig config) - { - return getHiveSession(config); - } - - public static Transaction newTransaction(HiveTransactionManager transactionManager) - { - return new HiveTransaction(transactionManager); - } - - public static MaterializedResult filterTable(SchemaTableName tableName, - List projectedColumns, - HiveTransactionManager transactionManager, - HiveConfig config, - ConnectorPageSourceProvider pageSourceProvider, - ConnectorSplitManager splitManager) - throws IOException - { - ConnectorMetadata metadata = null; - ConnectorSession session = null; - ConnectorSplitSource splitSource = null; - - try (Transaction transaction = newTransaction(transactionManager)) { - metadata = transaction.getMetadata(); - session = newSession(config); - - ConnectorTableHandle table = getTableHandle(metadata, tableName, session); - - metadata.beginQuery(session); - splitSource = getSplits(splitManager, transaction, session, table); - - List allTypes = getTypes(projectedColumns); - List dataTypes = getTypes(projectedColumns.stream() - .filter(columnHandle -> !((HiveColumnHandle) columnHandle).isHidden()) - .collect(toImmutableList())); - MaterializedResult.Builder result = MaterializedResult.resultBuilder(session, dataTypes); - - List splits = getAllSplits(splitSource); - for (ConnectorSplit split : splits) { - try (ConnectorPageSource pageSource = pageSourceProvider.createPageSource(transaction.getTransactionHandle(), - session, split, table, projectedColumns, DynamicFilter.EMPTY)) { - MaterializedResult pageSourceResult = materializeSourceDataStream(session, pageSource, allTypes); - for (MaterializedRow row : pageSourceResult.getMaterializedRows()) { - Object[] dataValues = IntStream.range(0, row.getFieldCount()) - .filter(channel -> !((HiveColumnHandle) projectedColumns.get(channel)).isHidden()) - .mapToObj(row::getField) - .toArray(); - result.row(dataValues); - } - } - } - return result.build(); - } - finally { - cleanUpQuery(metadata, session); - closeQuietly(splitSource); - } - } - - public static int getSplitsCount(SchemaTableName tableName, - HiveTransactionManager transactionManager, - HiveConfig config, - ConnectorSplitManager splitManager) - { - ConnectorMetadata metadata = null; - ConnectorSession session = null; - ConnectorSplitSource splitSource = null; - - try (Transaction transaction = newTransaction(transactionManager)) { - metadata = transaction.getMetadata(); - session = newSession(config); - - ConnectorTableHandle table = getTableHandle(metadata, tableName, session); - - metadata.beginQuery(session); - splitSource = getSplits(splitManager, transaction, session, table); - return getAllSplits(splitSource).size(); - } - finally { - cleanUpQuery(metadata, session); - closeQuietly(splitSource); - } - } - - private static void closeQuietly(Closeable closeable) - { - try { - if (closeable != null) { - closeable.close(); - } - } - catch (IOException ignored) { - } - } - - private static void cleanUpQuery(ConnectorMetadata metadata, ConnectorSession session) - { - if (metadata != null && session != null) { - metadata.cleanupQuery(session); - } - } -} diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/HiveQueryRunner.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/HiveQueryRunner.java index 19fa8bb0cd6e..5da4b1a0250a 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/HiveQueryRunner.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/HiveQueryRunner.java @@ -25,6 +25,7 @@ import io.trino.plugin.hive.fs.DirectoryLister; import io.trino.plugin.hive.metastore.Database; import io.trino.plugin.hive.metastore.HiveMetastore; +import io.trino.plugin.hive.metastore.HiveMetastoreFactory; import io.trino.plugin.tpcds.TpcdsPlugin; import io.trino.plugin.tpch.ColumnNaming; import io.trino.plugin.tpch.DecimalTypeMapping; @@ -38,7 +39,6 @@ import org.intellij.lang.annotations.Language; import org.joda.time.DateTimeZone; -import java.io.File; import java.nio.file.Path; import java.nio.file.Paths; import java.util.HashMap; @@ -50,7 +50,7 @@ import static com.google.inject.util.Modules.EMPTY_MODULE; import static io.airlift.log.Level.WARN; import static io.airlift.units.Duration.nanosSince; -import static io.trino.plugin.hive.metastore.file.TestingFileHiveMetastore.createTestingFileHiveMetastore; +import static io.trino.plugin.hive.TestingHiveUtils.getConnectorService; import static io.trino.plugin.hive.security.HiveSecurityModule.ALLOW_ALL; import static io.trino.plugin.hive.security.HiveSecurityModule.SQL_STANDARD; import static io.trino.plugin.tpch.ColumnNaming.SIMPLIFIED; @@ -102,10 +102,7 @@ public static class Builder> private List> initialTables = ImmutableList.of(); private Optional initialSchemasLocationBase = Optional.empty(); private Function initialTablesSessionMutator = Function.identity(); - private Function metastore = queryRunner -> { - File baseDir = queryRunner.getCoordinator().getBaseDataDir().resolve("hive_data").toFile(); - return createTestingFileHiveMetastore(baseDir); - }; + private Optional> metastore = Optional.empty(); private Optional openTelemetry = Optional.empty(); private Module module = EMPTY_MODULE; private Optional directoryLister = Optional.empty(); @@ -171,7 +168,7 @@ public SELF setInitialTablesSessionMutator(Function initialTab @CanIgnoreReturnValue public SELF setMetastore(Function metastore) { - this.metastore = requireNonNull(metastore, "metastore is null"); + this.metastore = Optional.ofNullable(metastore); return self(); } @@ -252,8 +249,19 @@ public DistributedQueryRunner build() queryRunner.createCatalog("tpcds", "tpcds"); } - HiveMetastore metastore = this.metastore.apply(queryRunner); - queryRunner.installPlugin(new TestingHivePlugin(Optional.of(metastore), openTelemetry, module, directoryLister)); + Optional metastore = this.metastore.map(factory -> factory.apply(queryRunner)); + Path dataDir = queryRunner.getCoordinator().getBaseDataDir().resolve("hive_data"); + + if (metastore.isEmpty() && !hiveProperties.buildOrThrow().containsKey("hive.metastore")) { + hiveProperties.put("hive.metastore", "file"); + hiveProperties.put("hive.metastore.catalog.dir", queryRunner.getCoordinator().getBaseDataDir().resolve("hive_data").toString()); + } + if (hiveProperties.buildOrThrow().keySet().stream().noneMatch(key -> + key.equals("fs.hadoop.enabled") || key.startsWith("fs.native-"))) { + hiveProperties.put("fs.hadoop.enabled", "true"); + } + + queryRunner.installPlugin(new TestingHivePlugin(dataDir, metastore)); Map hiveProperties = new HashMap<>(); if (!skipTimezoneSetup) { @@ -281,7 +289,7 @@ public DistributedQueryRunner build() queryRunner.createCatalog(HIVE_CATALOG, "hive", hiveProperties); if (createTpchSchemas) { - populateData(queryRunner, metastore); + populateData(queryRunner); } return queryRunner; @@ -292,8 +300,11 @@ public DistributedQueryRunner build() } } - private void populateData(DistributedQueryRunner queryRunner, HiveMetastore metastore) + private void populateData(DistributedQueryRunner queryRunner) { + HiveMetastore metastore = getConnectorService(queryRunner, HiveMetastoreFactory.class) + .createMetastore(Optional.empty()); + if (metastore.getDatabase(TPCH_SCHEMA).isEmpty()) { metastore.createDatabase(createDatabaseMetastoreObject(TPCH_SCHEMA, initialSchemasLocationBase)); Session session = initialTablesSessionMutator.apply(queryRunner.getDefaultSession()); diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/HiveTestUtils.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/HiveTestUtils.java index dedbc21e1d67..59046691de6c 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/HiveTestUtils.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/HiveTestUtils.java @@ -17,7 +17,6 @@ import com.google.common.collect.ImmutableSet; import com.google.common.net.HostAndPort; import io.airlift.slice.Slices; -import io.airlift.units.DataSize; import io.trino.filesystem.TrinoFileSystemFactory; import io.trino.filesystem.hdfs.HdfsFileSystemFactory; import io.trino.hdfs.DynamicHdfsConfiguration; @@ -57,8 +56,6 @@ import io.trino.plugin.hive.parquet.ParquetReaderConfig; import io.trino.plugin.hive.parquet.ParquetWriterConfig; import io.trino.plugin.hive.rcfile.RcFilePageSourceFactory; -import io.trino.plugin.hive.s3select.S3SelectRecordCursorProvider; -import io.trino.plugin.hive.s3select.TrinoS3ClientFactory; import io.trino.spi.PageSorter; import io.trino.spi.block.Block; import io.trino.spi.connector.ColumnHandle; @@ -97,7 +94,6 @@ import java.util.UUID; import static com.google.common.collect.ImmutableList.toImmutableList; -import static io.airlift.units.DataSize.Unit.MEGABYTE; import static io.trino.spi.block.ArrayValueBuilder.buildArrayValue; import static io.trino.spi.block.MapValueBuilder.buildMapValue; import static io.trino.spi.block.RowValueBuilder.buildRowValue; @@ -172,7 +168,6 @@ public static HiveSessionProperties getHiveSessionProperties(HiveConfig hiveConf { return new HiveSessionProperties( hiveConfig, - new HiveFormatsConfig(), orcReaderConfig, new OrcWriterConfig(), new ParquetReaderConfig(), @@ -183,7 +178,6 @@ public static HiveSessionProperties getHiveSessionProperties(HiveConfig hiveConf { return new HiveSessionProperties( hiveConfig, - new HiveFormatsConfig(), new OrcReaderConfig(), new OrcWriterConfig(), new ParquetReaderConfig(), @@ -194,39 +188,31 @@ public static HiveSessionProperties getHiveSessionProperties(HiveConfig hiveConf { return new HiveSessionProperties( hiveConfig, - new HiveFormatsConfig(), new OrcReaderConfig(), new OrcWriterConfig(), parquetReaderConfig, new ParquetWriterConfig()); } - public static Set getDefaultHivePageSourceFactories(HdfsEnvironment hdfsEnvironment, HiveConfig hiveConfig) + public static Set getDefaultHivePageSourceFactories(TrinoFileSystemFactory fileSystemFactory, HiveConfig hiveConfig) { - TrinoFileSystemFactory fileSystemFactory = new HdfsFileSystemFactory(hdfsEnvironment, HDFS_FILE_SYSTEM_STATS); FileFormatDataSourceStats stats = new FileFormatDataSourceStats(); return ImmutableSet.builder() - .add(new CsvPageSourceFactory(fileSystemFactory, stats, hiveConfig)) - .add(new JsonPageSourceFactory(fileSystemFactory, stats, hiveConfig)) - .add(new OpenXJsonPageSourceFactory(fileSystemFactory, stats, hiveConfig)) - .add(new RegexPageSourceFactory(fileSystemFactory, stats, hiveConfig)) - .add(new SimpleTextFilePageSourceFactory(fileSystemFactory, stats, hiveConfig)) - .add(new SimpleSequenceFilePageSourceFactory(fileSystemFactory, stats, hiveConfig)) - .add(new AvroPageSourceFactory(fileSystemFactory, stats)) - .add(new RcFilePageSourceFactory(fileSystemFactory, stats, hiveConfig)) + .add(new CsvPageSourceFactory(fileSystemFactory, hiveConfig)) + .add(new JsonPageSourceFactory(fileSystemFactory, hiveConfig)) + .add(new OpenXJsonPageSourceFactory(fileSystemFactory, hiveConfig)) + .add(new RegexPageSourceFactory(fileSystemFactory, hiveConfig)) + .add(new SimpleTextFilePageSourceFactory(fileSystemFactory, hiveConfig)) + .add(new SimpleSequenceFilePageSourceFactory(fileSystemFactory, hiveConfig)) + .add(new AvroPageSourceFactory(fileSystemFactory)) + .add(new RcFilePageSourceFactory(fileSystemFactory, hiveConfig)) .add(new OrcPageSourceFactory(new OrcReaderConfig(), fileSystemFactory, stats, hiveConfig)) .add(new ParquetPageSourceFactory(fileSystemFactory, stats, new ParquetReaderConfig(), hiveConfig)) .build(); } - public static Set getDefaultHiveRecordCursorProviders(HiveConfig hiveConfig, HdfsEnvironment hdfsEnvironment) + public static Set getDefaultHiveFileWriterFactories(HiveConfig hiveConfig, TrinoFileSystemFactory fileSystemFactory) { - return ImmutableSet.of(new S3SelectRecordCursorProvider(hdfsEnvironment, new TrinoS3ClientFactory(hiveConfig), hiveConfig)); - } - - public static Set getDefaultHiveFileWriterFactories(HiveConfig hiveConfig, HdfsEnvironment hdfsEnvironment) - { - TrinoFileSystemFactory fileSystemFactory = new HdfsFileSystemFactory(hdfsEnvironment, HDFS_FILE_SYSTEM_STATS); NodeVersion nodeVersion = new NodeVersion("test_version"); return ImmutableSet.builder() .add(new CsvFileWriterFactory(fileSystemFactory, TESTING_TYPE_MANAGER)) @@ -251,11 +237,6 @@ public static List getTypes(List columnHandles) return types.build(); } - public static HiveRecordCursorProvider createGenericHiveRecordCursorProvider(HdfsEnvironment hdfsEnvironment) - { - return new GenericHiveRecordCursorProvider(hdfsEnvironment, DataSize.of(100, MEGABYTE)); - } - public static MapType mapType(Type keyType, Type valueType) { return (MapType) TESTING_TYPE_MANAGER.getParameterizedType(StandardTypes.MAP, ImmutableList.of( diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestBackgroundHiveSplitLoader.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestBackgroundHiveSplitLoader.java index 9e65561e82be..b066ecc7d52a 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestBackgroundHiveSplitLoader.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestBackgroundHiveSplitLoader.java @@ -19,16 +19,25 @@ import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; import com.google.common.collect.ListMultimap; +import com.google.common.io.Resources; import io.airlift.stats.CounterStat; import io.airlift.units.DataSize; import io.airlift.units.Duration; -import io.trino.filesystem.hdfs.HdfsFileSystemFactory; +import io.trino.filesystem.FileEntry; +import io.trino.filesystem.FileIterator; +import io.trino.filesystem.Location; +import io.trino.filesystem.TrinoFileSystem; +import io.trino.filesystem.TrinoFileSystemFactory; +import io.trino.filesystem.TrinoInputFile; +import io.trino.filesystem.TrinoOutputFile; +import io.trino.filesystem.cache.DefaultCachingHostAddressProvider; +import io.trino.filesystem.memory.MemoryFileSystemFactory; import io.trino.hdfs.DynamicHdfsConfiguration; import io.trino.hdfs.HdfsConfig; import io.trino.hdfs.HdfsConfigurationInitializer; import io.trino.hdfs.HdfsEnvironment; -import io.trino.hdfs.HdfsNamenodeStats; import io.trino.hdfs.authentication.NoHdfsAuthentication; +import io.trino.hive.thrift.metastore.hive_metastoreConstants; import io.trino.plugin.hive.HiveColumnHandle.ColumnType; import io.trino.plugin.hive.fs.CachingDirectoryLister; import io.trino.plugin.hive.fs.DirectoryLister; @@ -37,6 +46,8 @@ import io.trino.plugin.hive.metastore.StorageFormat; import io.trino.plugin.hive.metastore.Table; import io.trino.plugin.hive.util.HiveBucketing.HiveBucketFilter; +import io.trino.plugin.hive.util.InternalHiveSplitFactory; +import io.trino.plugin.hive.util.SerdeConstants; import io.trino.plugin.hive.util.ValidWriteIdList; import io.trino.spi.TrinoException; import io.trino.spi.connector.ColumnHandle; @@ -58,7 +69,6 @@ import org.apache.hadoop.fs.RemoteIterator; import org.apache.hadoop.fs.permission.FsPermission; import org.apache.hadoop.hive.metastore.TableType; -import org.apache.hadoop.hive.ql.io.SymlinkTextInputFormat; import org.apache.hadoop.util.Progressable; import org.testng.annotations.AfterClass; import org.testng.annotations.BeforeClass; @@ -67,16 +77,17 @@ import java.io.File; import java.io.IOException; +import java.io.OutputStream; import java.net.URI; import java.nio.file.Files; -import java.nio.file.Paths; +import java.time.Instant; import java.util.ArrayList; +import java.util.Collection; import java.util.Iterator; import java.util.List; import java.util.Map; import java.util.Optional; import java.util.OptionalInt; -import java.util.Properties; import java.util.Set; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CountDownLatch; @@ -86,12 +97,10 @@ import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; +import static com.google.common.base.Predicates.alwaysTrue; import static com.google.common.base.Throwables.throwIfUnchecked; import static com.google.common.collect.ImmutableList.toImmutableList; -import static com.google.common.io.MoreFiles.deleteRecursively; -import static com.google.common.io.RecursiveDeleteOption.ALLOW_INSECURE; import static com.google.common.io.Resources.getResource; -import static io.airlift.concurrent.MoreFutures.unmodifiableFuture; import static io.airlift.concurrent.Threads.daemonThreadsNamed; import static io.airlift.slice.Slices.utf8Slice; import static io.airlift.units.DataSize.Unit.GIGABYTE; @@ -106,15 +115,17 @@ import static io.trino.plugin.hive.HiveStorageFormat.AVRO; import static io.trino.plugin.hive.HiveStorageFormat.CSV; import static io.trino.plugin.hive.HiveStorageFormat.ORC; -import static io.trino.plugin.hive.HiveTestUtils.HDFS_ENVIRONMENT; -import static io.trino.plugin.hive.HiveTestUtils.HDFS_FILE_SYSTEM_STATS; +import static io.trino.plugin.hive.HiveTableProperties.TRANSACTIONAL; import static io.trino.plugin.hive.HiveTestUtils.SESSION; import static io.trino.plugin.hive.HiveTestUtils.getHiveSession; import static io.trino.plugin.hive.HiveTimestampPrecision.DEFAULT_PRECISION; import static io.trino.plugin.hive.HiveType.HIVE_INT; import static io.trino.plugin.hive.HiveType.HIVE_STRING; import static io.trino.plugin.hive.util.HiveBucketing.BucketingVersion.BUCKETING_V1; +import static io.trino.plugin.hive.util.HiveClassNames.SYMLINK_TEXT_INPUT_FORMAT_CLASS; import static io.trino.plugin.hive.util.HiveUtil.getRegularColumnHandles; +import static io.trino.plugin.hive.util.SerdeConstants.FOOTER_COUNT; +import static io.trino.plugin.hive.util.SerdeConstants.HEADER_COUNT; import static io.trino.spi.predicate.TupleDomain.withColumnDomains; import static io.trino.spi.type.IntegerType.INTEGER; import static io.trino.spi.type.VarcharType.VARCHAR; @@ -124,8 +135,6 @@ import static java.nio.charset.StandardCharsets.UTF_8; import static java.util.concurrent.Executors.newCachedThreadPool; import static java.util.concurrent.TimeUnit.SECONDS; -import static org.apache.hadoop.hive.metastore.api.hive_metastoreConstants.FILE_INPUT_FORMAT; -import static org.apache.hadoop.hive.serde.serdeConstants.SERIALIZATION_LIB; import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatThrownBy; import static org.testng.Assert.assertEquals; @@ -136,31 +145,21 @@ public class TestBackgroundHiveSplitLoader { private static final int BUCKET_COUNT = 2; - private static final String SAMPLE_PATH = "hdfs://VOL1:9000/db_name/table_name/000000_0"; - private static final String SAMPLE_PATH_FILTERED = "hdfs://VOL1:9000/db_name/table_name/000000_1"; + private static final Location LOCATION = Location.of("memory:///db_name/table_name/000000_0"); + private static final Location FILTERED_LOCATION = Location.of("memory:///db_name/table_name/000000_1"); - private static final Path RETURNED_PATH = new Path(SAMPLE_PATH); - private static final Path FILTERED_PATH = new Path(SAMPLE_PATH_FILTERED); + private static final TupleDomain LOCATION_DOMAIN = withColumnDomains(Map.of(pathColumnHandle(), Domain.singleValue(VARCHAR, utf8Slice(LOCATION.toString())))); - private static final TupleDomain RETURNED_PATH_DOMAIN = withColumnDomains( - ImmutableMap.of( - pathColumnHandle(), - Domain.singleValue(VARCHAR, utf8Slice(RETURNED_PATH.toString())))); - - private static final List TEST_FILES = ImmutableList.of( - locatedFileStatus(RETURNED_PATH), - locatedFileStatus(FILTERED_PATH)); + private static final List TEST_LOCATIONS = List.of(LOCATION, FILTERED_LOCATION); private static final List PARTITION_COLUMNS = ImmutableList.of( new Column("partitionColumn", HIVE_INT, Optional.empty())); private static final List BUCKET_COLUMN_HANDLES = ImmutableList.of( createBaseColumn("col1", 0, HIVE_INT, INTEGER, ColumnType.REGULAR, Optional.empty())); - private static final Optional BUCKET_PROPERTY = Optional.of( - new HiveBucketProperty(ImmutableList.of("col1"), BUCKETING_V1, BUCKET_COUNT, ImmutableList.of())); - + private static final String TABLE_PATH = "memory:///db_name/table_name"; private static final Table SIMPLE_TABLE = table(ImmutableList.of(), Optional.empty(), ImmutableMap.of()); - private static final Table PARTITIONED_TABLE = table(PARTITION_COLUMNS, BUCKET_PROPERTY, ImmutableMap.of()); + private static final Table PARTITIONED_TABLE = table(TABLE_PATH, PARTITION_COLUMNS, Optional.of(new HiveBucketProperty(List.of("col1"), BUCKET_COUNT, List.of())), Map.of()); private ExecutorService executor; @@ -181,48 +180,49 @@ public void tearDown() public void testNoPathFilter() throws Exception { - BackgroundHiveSplitLoader backgroundHiveSplitLoader = backgroundHiveSplitLoader( - TEST_FILES, - TupleDomain.none()); + BackgroundHiveSplitLoader backgroundHiveSplitLoader = backgroundHiveSplitLoader(TEST_LOCATIONS, TupleDomain.none()); HiveSplitSource hiveSplitSource = hiveSplitSource(backgroundHiveSplitLoader); backgroundHiveSplitLoader.start(hiveSplitSource); - assertEquals(drain(hiveSplitSource).size(), 2); + assertThat(drain(hiveSplitSource)).hasSize(2); } @Test public void testCsv() throws Exception { - DataSize fileSize = DataSize.of(2, GIGABYTE); - assertSplitCount(CSV, ImmutableMap.of(), fileSize, 33); - assertSplitCount(CSV, ImmutableMap.of("skip.header.line.count", "1"), fileSize, 33); - assertSplitCount(CSV, ImmutableMap.of("skip.header.line.count", "2"), fileSize, 1); - assertSplitCount(CSV, ImmutableMap.of("skip.footer.line.count", "1"), fileSize, 1); - assertSplitCount(CSV, ImmutableMap.of("skip.header.line.count", "1", "skip.footer.line.count", "1"), fileSize, 1); + FileEntry file = new FileEntry(LOCATION, DataSize.of(2, GIGABYTE).toBytes(), Instant.now(), Optional.empty()); + assertCsvSplitCount(file, Map.of(), 33); + assertCsvSplitCount(file, Map.of(HEADER_COUNT, "1"), 33); + assertCsvSplitCount(file, Map.of(HEADER_COUNT, "2"), 1); + assertCsvSplitCount(file, Map.of(FOOTER_COUNT, "1"), 1); + assertCsvSplitCount(file, Map.of(HEADER_COUNT, "1", FOOTER_COUNT, "1"), 1); } - private void assertSplitCount(HiveStorageFormat storageFormat, Map tableProperties, DataSize fileSize, int expectedSplitCount) + private void assertCsvSplitCount(FileEntry file, Map tableProperties, int expectedSplitCount) throws Exception { Table table = table( - ImmutableList.of(), + TABLE_PATH, + List.of(), Optional.empty(), - ImmutableMap.copyOf(tableProperties), - StorageFormat.fromHiveStorageFormat(storageFormat)); + Map.copyOf(tableProperties), + CSV.toStorageFormat()); + TrinoFileSystemFactory fileSystemFactory = new ListSingleFileFileSystemFactory(file); BackgroundHiveSplitLoader backgroundHiveSplitLoader = backgroundHiveSplitLoader( - ImmutableList.of(locatedFileStatus(new Path(SAMPLE_PATH), fileSize.toBytes())), + fileSystemFactory, TupleDomain.all(), Optional.empty(), table, + Optional.empty(), Optional.empty()); HiveSplitSource hiveSplitSource = hiveSplitSource(backgroundHiveSplitLoader); backgroundHiveSplitLoader.start(hiveSplitSource); - assertEquals(drainSplits(hiveSplitSource).size(), expectedSplitCount); + assertThat(drainSplits(hiveSplitSource)).hasSize(expectedSplitCount); } @Test @@ -230,14 +230,14 @@ public void testPathFilter() throws Exception { BackgroundHiveSplitLoader backgroundHiveSplitLoader = backgroundHiveSplitLoader( - TEST_FILES, - RETURNED_PATH_DOMAIN); + TEST_LOCATIONS, + LOCATION_DOMAIN); HiveSplitSource hiveSplitSource = hiveSplitSource(backgroundHiveSplitLoader); backgroundHiveSplitLoader.start(hiveSplitSource); List paths = drain(hiveSplitSource); assertEquals(paths.size(), 1); - assertEquals(paths.get(0), RETURNED_PATH.toString()); + assertEquals(paths.get(0), LOCATION.toString()); } @Test @@ -245,17 +245,17 @@ public void testPathFilterOneBucketMatchPartitionedTable() throws Exception { BackgroundHiveSplitLoader backgroundHiveSplitLoader = backgroundHiveSplitLoader( - TEST_FILES, - RETURNED_PATH_DOMAIN, - Optional.of(new HiveBucketFilter(ImmutableSet.of(0, 1))), + TEST_LOCATIONS, + LOCATION_DOMAIN, + Optional.of(new HiveBucketFilter(Set.of(0, 1))), PARTITIONED_TABLE, - Optional.of(new HiveBucketHandle(BUCKET_COLUMN_HANDLES, BUCKETING_V1, BUCKET_COUNT, BUCKET_COUNT, ImmutableList.of()))); + Optional.of(new HiveTablePartitioning(true, BUCKETING_V1, BUCKET_COUNT, BUCKET_COLUMN_HANDLES, false, List.of(), true))); HiveSplitSource hiveSplitSource = hiveSplitSource(backgroundHiveSplitLoader); backgroundHiveSplitLoader.start(hiveSplitSource); List paths = drain(hiveSplitSource); assertEquals(paths.size(), 1); - assertEquals(paths.get(0), RETURNED_PATH.toString()); + assertEquals(paths.get(0), LOCATION.toString()); } @Test @@ -263,32 +263,42 @@ public void testPathFilterBucketedPartitionedTable() throws Exception { BackgroundHiveSplitLoader backgroundHiveSplitLoader = backgroundHiveSplitLoader( - TEST_FILES, - RETURNED_PATH_DOMAIN, + TEST_LOCATIONS, + LOCATION_DOMAIN, Optional.empty(), PARTITIONED_TABLE, Optional.of( - new HiveBucketHandle( - getRegularColumnHandles(PARTITIONED_TABLE, TESTING_TYPE_MANAGER, DEFAULT_PRECISION), + new HiveTablePartitioning( + true, BUCKETING_V1, BUCKET_COUNT, - BUCKET_COUNT, - ImmutableList.of()))); + getRegularColumnHandles(PARTITIONED_TABLE, TESTING_TYPE_MANAGER, DEFAULT_PRECISION), + false, + List.of(), + true))); HiveSplitSource hiveSplitSource = hiveSplitSource(backgroundHiveSplitLoader); backgroundHiveSplitLoader.start(hiveSplitSource); List paths = drain(hiveSplitSource); assertEquals(paths.size(), 1); - assertEquals(paths.get(0), RETURNED_PATH.toString()); + assertEquals(paths.get(0), LOCATION.toString()); } @Test public void testEmptyFileWithNoBlocks() throws Exception { + TrinoFileSystemFactory fileSystemFactory = new MemoryFileSystemFactory(); + // create an empty file + fileSystemFactory.create(ConnectorIdentity.ofUser("test")).newOutputFile(LOCATION).create().close(); + BackgroundHiveSplitLoader backgroundHiveSplitLoader = backgroundHiveSplitLoader( - ImmutableList.of(locatedFileStatusWithNoBlocks(RETURNED_PATH)), - TupleDomain.none()); + fileSystemFactory, + TupleDomain.none(), + Optional.empty(), + SIMPLE_TABLE, + Optional.empty(), + Optional.empty()); HiveSplitSource hiveSplitSource = hiveSplitSource(backgroundHiveSplitLoader); backgroundHiveSplitLoader.start(hiveSplitSource); @@ -299,6 +309,7 @@ public void testEmptyFileWithNoBlocks() @Test public void testNoHangIfPartitionIsOffline() + throws IOException { BackgroundHiveSplitLoader backgroundHiveSplitLoader = backgroundHiveSplitLoaderOfflinePartitions(); HiveSplitSource hiveSplitSource = hiveSplitSource(backgroundHiveSplitLoader); @@ -316,59 +327,58 @@ public void testNoHangIfPartitionIsOffline() public void testIncompleteDynamicFilterTimeout() throws Exception { - BackgroundHiveSplitLoader backgroundHiveSplitLoader = backgroundHiveSplitLoader( - new DynamicFilter() - { - @Override - public Set getColumnsCovered() + CompletableFuture isBlocked = new CompletableFuture<>(); + try { + BackgroundHiveSplitLoader backgroundHiveSplitLoader = backgroundHiveSplitLoader( + new DynamicFilter() { - return ImmutableSet.of(); - } + @Override + public Set getColumnsCovered() + { + return Set.of(); + } - @Override - public CompletableFuture isBlocked() - { - return unmodifiableFuture(CompletableFuture.runAsync(() -> { - try { - TimeUnit.HOURS.sleep(1); - } - catch (InterruptedException e) { - throw new IllegalStateException(e); - } - })); - } + @Override + public CompletableFuture isBlocked() + { + return isBlocked; + } - @Override - public boolean isComplete() - { - return false; - } + @Override + public boolean isComplete() + { + return false; + } - @Override - public boolean isAwaitable() - { - return true; - } + @Override + public boolean isAwaitable() + { + return true; + } - @Override - public TupleDomain getCurrentPredicate() - { - return TupleDomain.all(); - } - }, - new Duration(1, SECONDS)); - HiveSplitSource hiveSplitSource = hiveSplitSource(backgroundHiveSplitLoader); - backgroundHiveSplitLoader.start(hiveSplitSource); + @Override + public TupleDomain getCurrentPredicate() + { + return TupleDomain.all(); + } + }, + new Duration(1, SECONDS)); + HiveSplitSource hiveSplitSource = hiveSplitSource(backgroundHiveSplitLoader); + backgroundHiveSplitLoader.start(hiveSplitSource); - assertEquals(drain(hiveSplitSource).size(), 2); - assertTrue(hiveSplitSource.isFinished()); + assertThat(drain(hiveSplitSource)).hasSize(2); + assertThat(hiveSplitSource.isFinished()).isTrue(); + } + finally { + isBlocked.complete(null); + } } @Test public void testCachedDirectoryLister() throws Exception { - CachingDirectoryLister cachingDirectoryLister = new CachingDirectoryLister(new Duration(5, TimeUnit.MINUTES), DataSize.of(100, KILOBYTE), ImmutableList.of("test_dbname.test_table")); + CachingDirectoryLister cachingDirectoryLister = new CachingDirectoryLister(new Duration(5, TimeUnit.MINUTES), DataSize.of(100, KILOBYTE), List.of("test_dbname.test_table"), alwaysTrue()); assertEquals(cachingDirectoryLister.getRequestCount(), 0); int totalCount = 100; @@ -376,7 +386,7 @@ public void testCachedDirectoryLister() List>> futures = new ArrayList<>(); futures.add(executor.submit(() -> { - BackgroundHiveSplitLoader backgroundHiveSplitLoader = backgroundHiveSplitLoader(TEST_FILES, cachingDirectoryLister); + BackgroundHiveSplitLoader backgroundHiveSplitLoader = backgroundHiveSplitLoader(TEST_LOCATIONS, cachingDirectoryLister); HiveSplitSource hiveSplitSource = hiveSplitSource(backgroundHiveSplitLoader); backgroundHiveSplitLoader.start(hiveSplitSource); try { @@ -390,7 +400,7 @@ public void testCachedDirectoryLister() for (int i = 0; i < totalCount - 1; i++) { futures.add(executor.submit(() -> { firstVisit.await(); - BackgroundHiveSplitLoader backgroundHiveSplitLoader = backgroundHiveSplitLoader(TEST_FILES, cachingDirectoryLister); + BackgroundHiveSplitLoader backgroundHiveSplitLoader = backgroundHiveSplitLoader(TEST_LOCATIONS, cachingDirectoryLister); HiveSplitSource hiveSplitSource = hiveSplitSource(backgroundHiveSplitLoader); backgroundHiveSplitLoader.start(hiveSplitSource); return drainSplits(hiveSplitSource); @@ -398,7 +408,7 @@ public void testCachedDirectoryLister() } for (Future> future : futures) { - assertEquals(future.get().size(), TEST_FILES.size()); + assertEquals(future.get().size(), TEST_LOCATIONS.size()); } assertEquals(cachingDirectoryLister.getRequestCount(), totalCount); assertEquals(cachingDirectoryLister.getHitCount(), totalCount - 1); @@ -445,12 +455,24 @@ public void testGetAttemptId() assertFalse(hasAttemptId("base_00000_00")); } - @Test(dataProvider = "testPropagateExceptionDataProvider", timeOut = 60_000) - public void testPropagateException(boolean error, int threads) + @Test(timeOut = 60_000) + public void testPropagateException() + throws IOException + { + testPropagateException(false, 1); + testPropagateException(true, 1); + testPropagateException(false, 2); + testPropagateException(true, 2); + testPropagateException(false, 4); + testPropagateException(true, 4); + } + + private void testPropagateException(boolean error, int threads) + throws IOException { AtomicBoolean iteratorUsedAfterException = new AtomicBoolean(); - HdfsEnvironment hdfsEnvironment = new TestingHdfsEnvironment(TEST_FILES); + TrinoFileSystemFactory fileSystemFactory = createTestingFileSystem(TEST_LOCATIONS); BackgroundHiveSplitLoader backgroundHiveSplitLoader = new BackgroundHiveSplitLoader( SIMPLE_TABLE, new Iterator<>() @@ -481,15 +503,12 @@ public HivePartitionMetadata next() TESTING_TYPE_MANAGER, createBucketSplitInfo(Optional.empty(), Optional.empty()), SESSION, - new HdfsFileSystemFactory(hdfsEnvironment, HDFS_FILE_SYSTEM_STATS), - hdfsEnvironment, - new HdfsNamenodeStats(), + fileSystemFactory, new CachingDirectoryLister(new HiveConfig()), executor, threads, false, false, - true, Optional.empty(), Optional.empty(), 100); @@ -525,44 +544,47 @@ public Object[][] testPropagateExceptionDataProvider() public void testMultipleSplitsPerBucket() throws Exception { + TrinoFileSystemFactory fileSystemFactory = new ListSingleFileFileSystemFactory(new FileEntry(LOCATION, DataSize.of(1, GIGABYTE).toBytes(), Instant.now(), Optional.empty())); BackgroundHiveSplitLoader backgroundHiveSplitLoader = backgroundHiveSplitLoader( - ImmutableList.of(locatedFileStatus(new Path(SAMPLE_PATH), DataSize.of(1, GIGABYTE).toBytes())), + fileSystemFactory, TupleDomain.all(), Optional.empty(), SIMPLE_TABLE, - Optional.of(new HiveBucketHandle(BUCKET_COLUMN_HANDLES, BUCKETING_V1, BUCKET_COUNT, BUCKET_COUNT, ImmutableList.of()))); + Optional.of(new HiveTablePartitioning(true, BUCKETING_V1, BUCKET_COUNT, BUCKET_COLUMN_HANDLES, false, List.of(), true)), + Optional.empty()); HiveSplitSource hiveSplitSource = hiveSplitSource(backgroundHiveSplitLoader); backgroundHiveSplitLoader.start(hiveSplitSource); - assertEquals(drainSplits(hiveSplitSource).size(), 17); + assertThat(drainSplits(hiveSplitSource)).hasSize(17); } @Test public void testSplitsGenerationWithAbortedTransactions() throws Exception { - java.nio.file.Path tablePath = Files.createTempDirectory("TestBackgroundHiveSplitLoader"); + TrinoFileSystemFactory fileSystemFactory = new MemoryFileSystemFactory(); + TrinoFileSystem fileSystem = fileSystemFactory.create(ConnectorIdentity.ofUser("test")); + Location tableLocation = Location.of("memory:///my_table"); + Table table = table( - tablePath.toString(), - ImmutableList.of(), + tableLocation.toString(), + List.of(), Optional.empty(), - ImmutableMap.of( - "transactional", "true", + Map.of( + TRANSACTIONAL, "true", "transactional_properties", "insert_only")); - List filePaths = ImmutableList.of( - tablePath + "/delta_0000001_0000001_0000/_orc_acid_version", - tablePath + "/delta_0000001_0000001_0000/bucket_00000", - tablePath + "/delta_0000002_0000002_0000/_orc_acid_version", - tablePath + "/delta_0000002_0000002_0000/bucket_00000", - tablePath + "/delta_0000003_0000003_0000/_orc_acid_version", - tablePath + "/delta_0000003_0000003_0000/bucket_00000"); - - for (String path : filePaths) { - File file = new File(path); - assertTrue(file.getParentFile().exists() || file.getParentFile().mkdirs(), "Failed creating directory " + file.getParentFile()); - createOrcAcidFile(file); + List fileLocations = List.of( + tableLocation.appendPath("delta_0000001_0000001_0000/_orc_acid_version"), + tableLocation.appendPath("delta_0000001_0000001_0000/bucket_00000"), + tableLocation.appendPath("delta_0000002_0000002_0000/_orc_acid_version"), + tableLocation.appendPath("delta_0000002_0000002_0000/bucket_00000"), + tableLocation.appendPath("delta_0000003_0000003_0000/_orc_acid_version"), + tableLocation.appendPath("delta_0000003_0000003_0000/bucket_00000")); + + for (Location fileLocation : fileLocations) { + createOrcAcidFile(fileSystem, fileLocation); } // ValidWriteIdList is of format $.
:::: @@ -570,7 +592,7 @@ public void testSplitsGenerationWithAbortedTransactions() String validWriteIdsList = format("4$%s.%s:3:9223372036854775807::2", table.getDatabaseName(), table.getTableName()); BackgroundHiveSplitLoader backgroundHiveSplitLoader = backgroundHiveSplitLoader( - HDFS_ENVIRONMENT, + fileSystemFactory, TupleDomain.none(), Optional.empty(), table, @@ -580,41 +602,41 @@ public void testSplitsGenerationWithAbortedTransactions() HiveSplitSource hiveSplitSource = hiveSplitSource(backgroundHiveSplitLoader); backgroundHiveSplitLoader.start(hiveSplitSource); List splits = drain(hiveSplitSource); - assertTrue(splits.stream().anyMatch(p -> p.contains(filePaths.get(1))), format("%s not found in splits %s", filePaths.get(1), splits)); - assertTrue(splits.stream().anyMatch(p -> p.contains(filePaths.get(5))), format("%s not found in splits %s", filePaths.get(5), splits)); - - deleteRecursively(tablePath, ALLOW_INSECURE); + assertThat(splits).contains(fileLocations.get(1).toString()); + assertThat(splits).contains(fileLocations.get(5).toString()); } @Test public void testFullAcidTableWithOriginalFiles() throws Exception { - java.nio.file.Path tablePath = Files.createTempDirectory("TestBackgroundHiveSplitLoader"); + TrinoFileSystemFactory fileSystemFactory = new MemoryFileSystemFactory(); + TrinoFileSystem fileSystem = fileSystemFactory.create(ConnectorIdentity.ofUser("test")); + Location tableLocation = Location.of("memory:///my_table"); + Table table = table( - tablePath.toString(), - ImmutableList.of(), + tableLocation.toString(), + List.of(), Optional.empty(), - ImmutableMap.of("transactional", "true")); - - String originalFile = tablePath + "/000000_1"; - List filePaths = ImmutableList.of( - tablePath + "/delta_0000002_0000002_0000/_orc_acid_version", - tablePath + "/delta_0000002_0000002_0000/bucket_00000"); + Map.of(TRANSACTIONAL, "true")); - for (String path : filePaths) { - File file = new File(path); - assertTrue(file.getParentFile().exists() || file.getParentFile().mkdirs(), "Failed creating directory " + file.getParentFile()); - createOrcAcidFile(file); + Location originalFile = tableLocation.appendPath("000000_1"); + try (OutputStream outputStream = fileSystem.newOutputFile(originalFile).create()) { + outputStream.write("test".getBytes(UTF_8)); + } + List fileLocations = List.of( + tableLocation.appendPath("delta_0000002_0000002_0000/_orc_acid_version"), + tableLocation.appendPath("delta_0000002_0000002_0000/bucket_00000")); + for (Location fileLocation : fileLocations) { + createOrcAcidFile(fileSystem, fileLocation); } - Files.write(Paths.get(originalFile), "test".getBytes(UTF_8)); // ValidWriteIdsList is of format $.
:::: // This writeId list has high watermark transaction=3 ValidWriteIdList validWriteIdsList = new ValidWriteIdList(format("4$%s.%s:3:9223372036854775807::", table.getDatabaseName(), table.getTableName())); BackgroundHiveSplitLoader backgroundHiveSplitLoader = backgroundHiveSplitLoader( - HDFS_ENVIRONMENT, + fileSystemFactory, TupleDomain.all(), Optional.empty(), table, @@ -623,30 +645,31 @@ public void testFullAcidTableWithOriginalFiles() HiveSplitSource hiveSplitSource = hiveSplitSource(backgroundHiveSplitLoader); backgroundHiveSplitLoader.start(hiveSplitSource); List splits = drain(hiveSplitSource); - assertTrue(splits.stream().anyMatch(p -> p.contains(originalFile)), format("%s not found in splits %s", filePaths.get(0), splits)); - assertTrue(splits.stream().anyMatch(p -> p.contains(filePaths.get(1))), format("%s not found in splits %s", filePaths.get(1), splits)); + assertThat(splits).contains(originalFile.toString()); + assertThat(splits).contains(fileLocations.get(1).toString()); } @Test public void testVersionValidationNoOrcAcidVersionFile() throws Exception { - java.nio.file.Path tablePath = Files.createTempDirectory("TestBackgroundHiveSplitLoader"); + TrinoFileSystemFactory fileSystemFactory = new MemoryFileSystemFactory(); + TrinoFileSystem fileSystem = fileSystemFactory.create(ConnectorIdentity.ofUser("test")); + Location tableLocation = Location.of("memory:///my_table"); + Table table = table( - tablePath.toString(), - ImmutableList.of(), + tableLocation.toString(), + List.of(), Optional.empty(), - ImmutableMap.of("transactional", "true")); + Map.of(TRANSACTIONAL, "true")); - List filePaths = ImmutableList.of( - tablePath + "/000000_1", + List fileLocations = List.of( + tableLocation.appendPath("000000_1"), // no /delta_0000002_0000002_0000/_orc_acid_version file - tablePath + "/delta_0000002_0000002_0000/bucket_00000"); + tableLocation.appendPath("delta_0000002_0000002_0000/bucket_00000")); - for (String path : filePaths) { - File file = new File(path); - assertTrue(file.getParentFile().exists() || file.getParentFile().mkdirs(), "Failed creating directory " + file.getParentFile()); - createOrcAcidFile(file); + for (Location fileLocation : fileLocations) { + createOrcAcidFile(fileSystem, fileLocation); } // ValidWriteIdsList is of format $.
:::: @@ -654,7 +677,7 @@ public void testVersionValidationNoOrcAcidVersionFile() ValidWriteIdList validWriteIdsList = new ValidWriteIdList(format("4$%s.%s:3:9223372036854775807::", table.getDatabaseName(), table.getTableName())); BackgroundHiveSplitLoader backgroundHiveSplitLoader = backgroundHiveSplitLoader( - HDFS_ENVIRONMENT, + fileSystemFactory, TupleDomain.all(), Optional.empty(), table, @@ -669,30 +692,29 @@ public void testVersionValidationNoOrcAcidVersionFile() .allMatch(Optional::isPresent) .extracting(Optional::get) .noneMatch(AcidInfo::isOrcAcidVersionValidated); - - deleteRecursively(tablePath, ALLOW_INSECURE); } @Test public void testVersionValidationOrcAcidVersionFileHasVersion2() throws Exception { - java.nio.file.Path tablePath = Files.createTempDirectory("TestBackgroundHiveSplitLoader"); + TrinoFileSystemFactory fileSystemFactory = new MemoryFileSystemFactory(); + TrinoFileSystem fileSystem = fileSystemFactory.create(ConnectorIdentity.ofUser("test")); + Location tableLocation = Location.of("memory:///my_table"); + Table table = table( - tablePath.toString(), - ImmutableList.of(), + tableLocation.toString(), + List.of(), Optional.empty(), - ImmutableMap.of("transactional", "true")); + Map.of(TRANSACTIONAL, "true")); - List filePaths = ImmutableList.of( - tablePath + "/000000_1", // _orc_acid_version does not exist so it's assumed to be "ORC ACID version 0" - tablePath + "/delta_0000002_0000002_0000/_orc_acid_version", - tablePath + "/delta_0000002_0000002_0000/bucket_00000"); + List fileLocations = List.of( + tableLocation.appendPath("000000_1"), // _orc_acid_version does not exist, so it's assumed to be "ORC ACID version 0" + tableLocation.appendPath("delta_0000002_0000002_0000/_orc_acid_version"), + tableLocation.appendPath("delta_0000002_0000002_0000/bucket_00000")); - for (String path : filePaths) { - File file = new File(path); - assertTrue(file.getParentFile().exists() || file.getParentFile().mkdirs(), "Failed creating directory " + file.getParentFile()); - createOrcAcidFile(file, 2); + for (Location fileLocation : fileLocations) { + createOrcAcidFile(fileSystem, fileLocation, 2); } // ValidWriteIdsList is of format $.
:::: @@ -700,7 +722,7 @@ public void testVersionValidationOrcAcidVersionFileHasVersion2() ValidWriteIdList validWriteIdsList = new ValidWriteIdList(format("4$%s.%s:3:9223372036854775807::", table.getDatabaseName(), table.getTableName())); BackgroundHiveSplitLoader backgroundHiveSplitLoader = backgroundHiveSplitLoader( - HDFS_ENVIRONMENT, + fileSystemFactory, TupleDomain.all(), Optional.empty(), table, @@ -713,31 +735,30 @@ public void testVersionValidationOrcAcidVersionFileHasVersion2() // We should have it marked in all splits that NO further ORC ACID validation is required assertThat(drainSplits(hiveSplitSource)).extracting(HiveSplit::getAcidInfo) .allMatch(acidInfo -> acidInfo.isEmpty() || acidInfo.get().isOrcAcidVersionValidated()); - - deleteRecursively(tablePath, ALLOW_INSECURE); } @Test public void testVersionValidationOrcAcidVersionFileHasVersion1() throws Exception { - java.nio.file.Path tablePath = Files.createTempDirectory("TestBackgroundHiveSplitLoader"); + TrinoFileSystemFactory fileSystemFactory = new MemoryFileSystemFactory(); + TrinoFileSystem fileSystem = fileSystemFactory.create(ConnectorIdentity.ofUser("test")); + Location tableLocation = Location.of("memory:///my_table"); + Table table = table( - tablePath.toString(), - ImmutableList.of(), + tableLocation.toString(), + List.of(), Optional.empty(), - ImmutableMap.of("transactional", "true")); + Map.of(TRANSACTIONAL, "true")); - List filePaths = ImmutableList.of( - tablePath + "/000000_1", - tablePath + "/delta_0000002_0000002_0000/_orc_acid_version", - tablePath + "/delta_0000002_0000002_0000/bucket_00000"); + List fileLocations = List.of( + tableLocation.appendPath("000000_1"), + tableLocation.appendPath("delta_0000002_0000002_0000/_orc_acid_version"), + tableLocation.appendPath("delta_0000002_0000002_0000/bucket_00000")); - for (String path : filePaths) { - File file = new File(path); - assertTrue(file.getParentFile().exists() || file.getParentFile().mkdirs(), "Failed creating directory " + file.getParentFile()); + for (Location fileLocation : fileLocations) { // _orc_acid_version_exists but has version 1 - createOrcAcidFile(file, 1); + createOrcAcidFile(fileSystem, fileLocation, 1); } // ValidWriteIdsList is of format $.
:::: @@ -745,7 +766,7 @@ public void testVersionValidationOrcAcidVersionFileHasVersion1() ValidWriteIdList validWriteIdsList = new ValidWriteIdList(format("4$%s.%s:3:9223372036854775807::", table.getDatabaseName(), table.getTableName())); BackgroundHiveSplitLoader backgroundHiveSplitLoader = backgroundHiveSplitLoader( - HDFS_ENVIRONMENT, + fileSystemFactory, TupleDomain.all(), Optional.empty(), table, @@ -760,8 +781,6 @@ public void testVersionValidationOrcAcidVersionFileHasVersion1() .allMatch(Optional::isPresent) .extracting(Optional::get) .noneMatch(AcidInfo::isOrcAcidVersionValidated); - - deleteRecursively(tablePath, ALLOW_INSECURE); } @Test @@ -798,35 +817,39 @@ public void testValidateFileBuckets() public void testBuildManifestFileIterator() throws Exception { - CachingDirectoryLister directoryLister = new CachingDirectoryLister(new Duration(0, TimeUnit.MINUTES), DataSize.ofBytes(0), ImmutableList.of()); - Properties schema = new Properties(); - schema.setProperty(FILE_INPUT_FORMAT, SymlinkTextInputFormat.class.getName()); - schema.setProperty(SERIALIZATION_LIB, AVRO.getSerde()); - - Path firstFilePath = new Path("hdfs://VOL1:9000/db_name/table_name/file1"); - Path secondFilePath = new Path("hdfs://VOL1:9000/db_name/table_name/file2"); - List paths = ImmutableList.of(firstFilePath, secondFilePath); - List files = paths.stream() - .map(TestBackgroundHiveSplitLoader::locatedFileStatus) - .collect(toImmutableList()); + CachingDirectoryLister directoryLister = new CachingDirectoryLister(new Duration(0, TimeUnit.MINUTES), DataSize.ofBytes(0), List.of(), alwaysTrue()); + Map schema = ImmutableMap.builder() + .put(hive_metastoreConstants.FILE_INPUT_FORMAT, SYMLINK_TEXT_INPUT_FORMAT_CLASS) + .put(SerdeConstants.SERIALIZATION_LIB, AVRO.getSerde()) + .buildOrThrow(); - BackgroundHiveSplitLoader backgroundHiveSplitLoader = backgroundHiveSplitLoader( - files, - directoryLister); - Optional> splitIterator = backgroundHiveSplitLoader.buildManifestFileIterator( - AVRO, + Location firstFilePath = Location.of("memory:///db_name/table_name/file1"); + Location secondFilePath = Location.of("memory:///db_name/table_name/file2"); + List locations = List.of(firstFilePath, secondFilePath); + + InternalHiveSplitFactory splitFactory = new InternalHiveSplitFactory( "partition", + AVRO, schema, - ImmutableList.of(), + List.of(), TupleDomain.all(), () -> true, + ImmutableMap.of(), + Optional.empty(), + Optional.empty(), + DataSize.of(512, MEGABYTE), false, - TableToPartitionMapping.empty(), - new Path("hdfs://VOL1:9000/db_name/table_name"), - paths, + Optional.empty()); + + BackgroundHiveSplitLoader backgroundHiveSplitLoader = backgroundHiveSplitLoader( + locations, + directoryLister); + Iterator splitIterator = backgroundHiveSplitLoader.buildManifestFileIterator( + splitFactory, + Location.of(TABLE_PATH), + locations, true); - assertTrue(splitIterator.isPresent()); - List splits = ImmutableList.copyOf(splitIterator.get()); + List splits = ImmutableList.copyOf(splitIterator); assertEquals(splits.size(), 2); assertEquals(splits.get(0).getPath(), firstFilePath.toString()); assertEquals(splits.get(1).getPath(), secondFilePath.toString()); @@ -836,41 +859,49 @@ public void testBuildManifestFileIterator() public void testBuildManifestFileIteratorNestedDirectory() throws Exception { - CachingDirectoryLister directoryLister = new CachingDirectoryLister(new Duration(5, TimeUnit.MINUTES), DataSize.of(100, KILOBYTE), ImmutableList.of()); - Properties schema = new Properties(); - schema.setProperty(FILE_INPUT_FORMAT, SymlinkTextInputFormat.class.getName()); - schema.setProperty(SERIALIZATION_LIB, AVRO.getSerde()); - - Path filePath = new Path("hdfs://VOL1:9000/db_name/table_name/file1"); - Path directoryPath = new Path("hdfs://VOL1:9000/db_name/table_name/dir/file2"); - List paths = ImmutableList.of(filePath, directoryPath); - List files = ImmutableList.of( - locatedFileStatus(filePath), - locatedFileStatus(directoryPath)); + CachingDirectoryLister directoryLister = new CachingDirectoryLister(new Duration(5, TimeUnit.MINUTES), DataSize.of(100, KILOBYTE), List.of(), alwaysTrue()); + Map schema = ImmutableMap.builder() + .put(hive_metastoreConstants.FILE_INPUT_FORMAT, SYMLINK_TEXT_INPUT_FORMAT_CLASS) + .put(SerdeConstants.SERIALIZATION_LIB, AVRO.getSerde()) + .buildOrThrow(); - BackgroundHiveSplitLoader backgroundHiveSplitLoader = backgroundHiveSplitLoader( - files, - directoryLister); - Optional> splitIterator = backgroundHiveSplitLoader.buildManifestFileIterator( - AVRO, + Location filePath = Location.of("memory:///db_name/table_name/file1"); + Location directoryPath = Location.of("memory:///db_name/table_name/dir/file2"); + List locations = List.of(filePath, directoryPath); + + InternalHiveSplitFactory splitFactory = new InternalHiveSplitFactory( "partition", + AVRO, schema, - ImmutableList.of(), + List.of(), TupleDomain.all(), () -> true, + ImmutableMap.of(), + Optional.empty(), + Optional.empty(), + DataSize.of(512, MEGABYTE), false, - TableToPartitionMapping.empty(), - new Path("hdfs://VOL1:9000/db_name/table_name"), - paths, + Optional.empty()); + + BackgroundHiveSplitLoader backgroundHiveSplitLoader = backgroundHiveSplitLoader( + locations, + directoryLister); + Iterator splitIterator = backgroundHiveSplitLoader.buildManifestFileIterator( + splitFactory, + Location.of(TABLE_PATH), + locations, false); - assertTrue(splitIterator.isEmpty()); + List splits = ImmutableList.copyOf(splitIterator); + assertThat(splits).hasSize(2); + assertThat(splits.get(0).getPath()).isEqualTo(filePath.toString()); + assertThat(splits.get(1).getPath()).isEqualTo(directoryPath.toString()); } @Test public void testMaxPartitions() throws Exception { - CachingDirectoryLister directoryLister = new CachingDirectoryLister(new Duration(0, TimeUnit.MINUTES), DataSize.ofBytes(0), ImmutableList.of()); + CachingDirectoryLister directoryLister = new CachingDirectoryLister(new Duration(0, TimeUnit.MINUTES), DataSize.ofBytes(0), List.of(), alwaysTrue()); // zero partitions { BackgroundHiveSplitLoader backgroundHiveSplitLoader = backgroundHiveSplitLoader( @@ -887,12 +918,12 @@ public void testMaxPartitions() { BackgroundHiveSplitLoader backgroundHiveSplitLoader = backgroundHiveSplitLoader( ImmutableList.of(createPartitionMetadata()), - TEST_FILES, + TEST_LOCATIONS, directoryLister, 1); HiveSplitSource hiveSplitSource = hiveSplitSource(backgroundHiveSplitLoader); backgroundHiveSplitLoader.start(hiveSplitSource); - assertThat(drainSplits(hiveSplitSource)).hasSize(TEST_FILES.size()); + assertThat(drainSplits(hiveSplitSource)).hasSize(TEST_LOCATIONS.size()); } // single partition, crossing the limit @@ -900,7 +931,7 @@ public void testMaxPartitions() int partitionLimit = 0; BackgroundHiveSplitLoader backgroundHiveSplitLoader = backgroundHiveSplitLoader( ImmutableList.of(createPartitionMetadata()), - TEST_FILES, + TEST_LOCATIONS, directoryLister, partitionLimit); HiveSplitSource hiveSplitSource = hiveSplitSource(backgroundHiveSplitLoader); @@ -919,12 +950,12 @@ public void testMaxPartitions() List partitions = ImmutableList.of(createPartitionMetadata(), createPartitionMetadata(), createPartitionMetadata()); BackgroundHiveSplitLoader backgroundHiveSplitLoader = backgroundHiveSplitLoader( partitions, - TEST_FILES, + TEST_LOCATIONS, directoryLister, partitionLimit); HiveSplitSource hiveSplitSource = hiveSplitSource(backgroundHiveSplitLoader); backgroundHiveSplitLoader.start(hiveSplitSource); - assertThat(drainSplits(hiveSplitSource)).hasSize(TEST_FILES.size() * partitions.size()); + assertThat(drainSplits(hiveSplitSource)).hasSize(TEST_LOCATIONS.size() * partitions.size()); } // multiple partitions, crossing the limit @@ -932,7 +963,7 @@ public void testMaxPartitions() int partitionLimit = 3; BackgroundHiveSplitLoader backgroundHiveSplitLoader = backgroundHiveSplitLoader( ImmutableList.of(createPartitionMetadata(), createPartitionMetadata(), createPartitionMetadata(), createPartitionMetadata()), - TEST_FILES, + TEST_LOCATIONS, directoryLister, partitionLimit); HiveSplitSource hiveSplitSource = hiveSplitSource(backgroundHiveSplitLoader); @@ -951,7 +982,7 @@ private static HivePartitionMetadata createPartitionMetadata() return new HivePartitionMetadata( new HivePartition(SIMPLE_TABLE.getSchemaTableName()), Optional.empty(), - TableToPartitionMapping.empty()); + ImmutableMap.of()); } private static void createOrcAcidFile(File file) @@ -970,6 +1001,24 @@ private static void createOrcAcidFile(File file, int orcAcidVersion) Files.copy(getResource("fullacidNationTableWithOriginalFiles/000000_0").openStream(), file.toPath()); } + private static void createOrcAcidFile(TrinoFileSystem fileSystem, Location location) + throws IOException + { + createOrcAcidFile(fileSystem, location, 2); + } + + private static void createOrcAcidFile(TrinoFileSystem fileSystem, Location location, int orcAcidVersion) + throws IOException + { + try (OutputStream outputStream = fileSystem.newOutputFile(location).create()) { + if (location.fileName().equals("_orc_acid_version")) { + outputStream.write(String.valueOf(orcAcidVersion).getBytes(UTF_8)); + return; + } + Resources.copy(getResource("fullacidNationTableWithOriginalFiles/000000_0"), outputStream); + } + } + private static List drain(HiveSplitSource source) throws Exception { @@ -998,12 +1047,33 @@ private static List drainSplits(HiveSplitSource source) return splits.build(); } + private BackgroundHiveSplitLoader backgroundHiveSplitLoader( + TrinoFileSystemFactory fileSystemFactory, + TupleDomain compactEffectivePredicate, + Optional hiveBucketFilter, + Table table, + Optional tablePartitioning, + Optional validWriteIds) + { + return backgroundHiveSplitLoader( + fileSystemFactory, + compactEffectivePredicate, + DynamicFilter.EMPTY, + new Duration(0, SECONDS), + hiveBucketFilter, + table, + tablePartitioning, + validWriteIds); + } + private BackgroundHiveSplitLoader backgroundHiveSplitLoader( DynamicFilter dynamicFilter, Duration dynamicFilteringProbeBlockingTimeoutMillis) + throws IOException { + TrinoFileSystemFactory fileSystemFactory = createTestingFileSystem(TEST_LOCATIONS); return backgroundHiveSplitLoader( - new TestingHdfsEnvironment(TEST_FILES), + fileSystemFactory, TupleDomain.all(), dynamicFilter, dynamicFilteringProbeBlockingTimeoutMillis, @@ -1013,12 +1083,26 @@ private BackgroundHiveSplitLoader backgroundHiveSplitLoader( Optional.empty()); } + private static TrinoFileSystemFactory createTestingFileSystem(Collection locations) + throws IOException + { + TrinoFileSystemFactory fileSystemFactory = new MemoryFileSystemFactory(); + TrinoFileSystem fileSystem = fileSystemFactory.create(ConnectorIdentity.ofUser("test")); + for (Location location : locations) { + try (OutputStream outputStream = fileSystem.newOutputFile(location).create()) { + outputStream.write(new byte[10]); + } + } + return fileSystemFactory; + } + private BackgroundHiveSplitLoader backgroundHiveSplitLoader( - List files, + List locations, TupleDomain tupleDomain) + throws IOException { return backgroundHiveSplitLoader( - files, + locations, tupleDomain, Optional.empty(), SIMPLE_TABLE, @@ -1026,65 +1110,49 @@ private BackgroundHiveSplitLoader backgroundHiveSplitLoader( } private BackgroundHiveSplitLoader backgroundHiveSplitLoader( - List files, + List locations, TupleDomain compactEffectivePredicate, Optional hiveBucketFilter, Table table, - Optional bucketHandle) + Optional tablePartitioning) + throws IOException { return backgroundHiveSplitLoader( - files, + locations, compactEffectivePredicate, hiveBucketFilter, table, - bucketHandle, + tablePartitioning, Optional.empty()); } private BackgroundHiveSplitLoader backgroundHiveSplitLoader( - List files, - TupleDomain compactEffectivePredicate, - Optional hiveBucketFilter, - Table table, - Optional bucketHandle, - Optional validWriteIds) - { - return backgroundHiveSplitLoader( - new TestingHdfsEnvironment(files), - compactEffectivePredicate, - hiveBucketFilter, - table, - bucketHandle, - validWriteIds); - } - - private BackgroundHiveSplitLoader backgroundHiveSplitLoader( - HdfsEnvironment hdfsEnvironment, + List locations, TupleDomain compactEffectivePredicate, Optional hiveBucketFilter, Table table, - Optional bucketHandle, + Optional tablePartitioning, Optional validWriteIds) + throws IOException { + TrinoFileSystemFactory fileSystemFactory = createTestingFileSystem(locations); return backgroundHiveSplitLoader( - hdfsEnvironment, + fileSystemFactory, compactEffectivePredicate, - DynamicFilter.EMPTY, - new Duration(0, SECONDS), hiveBucketFilter, table, - bucketHandle, + tablePartitioning, validWriteIds); } private BackgroundHiveSplitLoader backgroundHiveSplitLoader( - HdfsEnvironment hdfsEnvironment, + TrinoFileSystemFactory fileSystemFactory, TupleDomain compactEffectivePredicate, DynamicFilter dynamicFilter, Duration dynamicFilteringProbeBlockingTimeout, Optional hiveBucketFilter, Table table, - Optional bucketHandle, + Optional tablePartitioning, Optional validWriteIds) { List hivePartitionMetadatas = @@ -1092,7 +1160,7 @@ private BackgroundHiveSplitLoader backgroundHiveSplitLoader( new HivePartitionMetadata( new HivePartition(new SchemaTableName("testSchema", "table_name")), Optional.empty(), - TableToPartitionMapping.empty())); + ImmutableMap.of())); return new BackgroundHiveSplitLoader( table, @@ -1101,44 +1169,43 @@ private BackgroundHiveSplitLoader backgroundHiveSplitLoader( dynamicFilter, dynamicFilteringProbeBlockingTimeout, TESTING_TYPE_MANAGER, - createBucketSplitInfo(bucketHandle, hiveBucketFilter), + createBucketSplitInfo(tablePartitioning, hiveBucketFilter), SESSION, - new HdfsFileSystemFactory(hdfsEnvironment, HDFS_FILE_SYSTEM_STATS), - hdfsEnvironment, - new HdfsNamenodeStats(), + fileSystemFactory, new CachingDirectoryLister(new HiveConfig()), executor, 2, false, false, - true, validWriteIds, Optional.empty(), 100); } private BackgroundHiveSplitLoader backgroundHiveSplitLoader( - List files, + List locations, DirectoryLister directoryLister) + throws IOException { - List partitions = ImmutableList.of( + List partitions = List.of( new HivePartitionMetadata( new HivePartition(new SchemaTableName("testSchema", "table_name")), Optional.empty(), - TableToPartitionMapping.empty())); - return backgroundHiveSplitLoader(partitions, files, directoryLister, 100); + ImmutableMap.of())); + return backgroundHiveSplitLoader(partitions, locations, directoryLister, 100); } private BackgroundHiveSplitLoader backgroundHiveSplitLoader( List partitions, - List files, + List locations, DirectoryLister directoryLister, int maxPartitions) + throws IOException { ConnectorSession connectorSession = getHiveSession(new HiveConfig() .setMaxSplitSize(DataSize.of(1, GIGABYTE))); - HdfsEnvironment hdfsEnvironment = new TestingHdfsEnvironment(files); + TrinoFileSystemFactory fileSystemFactory = createTestingFileSystem(locations); return new BackgroundHiveSplitLoader( SIMPLE_TABLE, partitions.iterator(), @@ -1148,26 +1215,24 @@ private BackgroundHiveSplitLoader backgroundHiveSplitLoader( TESTING_TYPE_MANAGER, Optional.empty(), connectorSession, - new HdfsFileSystemFactory(hdfsEnvironment, HDFS_FILE_SYSTEM_STATS), - hdfsEnvironment, - new HdfsNamenodeStats(), + fileSystemFactory, directoryLister, executor, 2, false, false, - true, Optional.empty(), Optional.empty(), maxPartitions); } private BackgroundHiveSplitLoader backgroundHiveSplitLoaderOfflinePartitions() + throws IOException { ConnectorSession connectorSession = getHiveSession(new HiveConfig() .setMaxSplitSize(DataSize.of(1, GIGABYTE))); - HdfsEnvironment hdfsEnvironment = new TestingHdfsEnvironment(TEST_FILES); + TrinoFileSystemFactory fileSystemFactory = createTestingFileSystem(TEST_LOCATIONS); return new BackgroundHiveSplitLoader( SIMPLE_TABLE, createPartitionMetadataWithOfflinePartitions(), @@ -1177,15 +1242,12 @@ private BackgroundHiveSplitLoader backgroundHiveSplitLoaderOfflinePartitions() TESTING_TYPE_MANAGER, createBucketSplitInfo(Optional.empty(), Optional.empty()), connectorSession, - new HdfsFileSystemFactory(hdfsEnvironment, HDFS_FILE_SYSTEM_STATS), - hdfsEnvironment, - new HdfsNamenodeStats(), + fileSystemFactory, new CachingDirectoryLister(new HiveConfig()), executor, 2, false, false, - true, Optional.empty(), Optional.empty(), 100); @@ -1209,7 +1271,7 @@ protected HivePartitionMetadata computeNext() return new HivePartitionMetadata( new HivePartition(new SchemaTableName("testSchema", "table_name")), Optional.empty(), - TableToPartitionMapping.empty()); + ImmutableMap.of()); case 1: throw new RuntimeException("OFFLINE"); default: @@ -1232,9 +1294,23 @@ private HiveSplitSource hiveSplitSource(HiveSplitLoader hiveSplitLoader) hiveSplitLoader, executor, new CounterStat(), + new DefaultCachingHostAddressProvider(), false); } + private static Table table( + String location, + List partitionColumns, + Optional bucketProperty, + Map tableParameters) + { + return table(location, + partitionColumns, + bucketProperty, + tableParameters, + ORC.toStorageFormat()); + } + private static Table table( List partitionColumns, Optional bucketProperty, @@ -1265,7 +1341,7 @@ private static Table table( Map tableParameters, StorageFormat storageFormat) { - return table("hdfs://VOL1:9000/db_name/table_name", + return table(TABLE_PATH, partitionColumns, bucketProperty, tableParameters, @@ -1474,4 +1550,102 @@ public URI getUri() return URI.create("hdfs://VOL1:9000/"); } } + + private record ListSingleFileFileSystemFactory(FileEntry fileEntry) + implements TrinoFileSystemFactory + { + @Override + public TrinoFileSystem create(ConnectorIdentity identity) + { + return new TrinoFileSystem() + { + @Override + public Optional directoryExists(Location location) + { + return Optional.empty(); + } + + @Override + public FileIterator listFiles(Location location) + { + Iterator iterator = List.of(fileEntry).iterator(); + return new FileIterator() + { + @Override + public boolean hasNext() + { + return iterator.hasNext(); + } + + @Override + public FileEntry next() + { + return iterator.next(); + } + }; + } + + @Override + public TrinoInputFile newInputFile(Location location) + { + throw new UnsupportedOperationException(); + } + + @Override + public TrinoInputFile newInputFile(Location location, long length) + { + throw new UnsupportedOperationException(); + } + + @Override + public TrinoOutputFile newOutputFile(Location location) + { + throw new UnsupportedOperationException(); + } + + @Override + public void deleteFile(Location location) + { + throw new UnsupportedOperationException(); + } + + @Override + public void deleteDirectory(Location location) + { + throw new UnsupportedOperationException(); + } + + @Override + public void renameFile(Location source, Location target) + { + throw new UnsupportedOperationException(); + } + + @Override + public void createDirectory(Location location) + { + throw new UnsupportedOperationException(); + } + + @Override + public void renameDirectory(Location source, Location target) + { + throw new UnsupportedOperationException(); + } + + @Override + public Set listDirectories(Location location) + { + throw new UnsupportedOperationException(); + } + + @Override + public Optional createTemporaryDirectory(Location targetPath, String temporaryPrefix, String relativePrefix) + throws IOException + { + throw new UnsupportedOperationException(); + } + }; + } + } } diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHive2OnDataLake.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHive2OnDataLake.java index b55338274da4..23b9fc0bbaaf 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHive2OnDataLake.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHive2OnDataLake.java @@ -15,11 +15,15 @@ import io.trino.plugin.hive.containers.HiveHadoop; +import static io.trino.testing.TestingNames.randomNameSuffix; + public class TestHive2OnDataLake extends BaseTestHiveOnDataLake { + private static final String BUCKET_NAME = "test-hive-insert-overwrite-" + randomNameSuffix(); + public TestHive2OnDataLake() { - super(HiveHadoop.DEFAULT_IMAGE); + super(BUCKET_NAME, HiveHadoop.HIVE3_IMAGE); } } diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHive3OnDataLake.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHive3OnDataLake.java index e952b6146bf9..c47c2d3b6180 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHive3OnDataLake.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHive3OnDataLake.java @@ -16,12 +16,16 @@ import io.trino.plugin.hive.containers.HiveHadoop; import org.testng.annotations.Test; +import static io.trino.testing.TestingNames.randomNameSuffix; + public class TestHive3OnDataLake extends BaseTestHiveOnDataLake { + private static final String BUCKET_NAME = "test-hive-insert-overwrite-" + randomNameSuffix(); + public TestHive3OnDataLake() { - super(HiveHadoop.HIVE3_IMAGE); + super(BUCKET_NAME, HiveHadoop.HIVE3_IMAGE); } @Test diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHiveAnalyzeCorruptStatistics.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHiveAnalyzeCorruptStatistics.java index ff1d8992eff3..d5fee159477d 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHiveAnalyzeCorruptStatistics.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHiveAnalyzeCorruptStatistics.java @@ -14,6 +14,7 @@ package io.trino.plugin.hive; import io.airlift.units.Duration; +import io.trino.plugin.hive.containers.Hive3MinioDataLake; import io.trino.plugin.hive.containers.HiveMinioDataLake; import io.trino.plugin.hive.s3.S3HiveQueryRunner; import io.trino.testing.AbstractTestQueryFramework; @@ -35,7 +36,7 @@ public class TestHiveAnalyzeCorruptStatistics protected QueryRunner createQueryRunner() throws Exception { - hiveMinioDataLake = closeAfterClass(new HiveMinioDataLake("test-analyze")); + hiveMinioDataLake = closeAfterClass(new Hive3MinioDataLake("test-analyze")); hiveMinioDataLake.start(); return S3HiveQueryRunner.builder(hiveMinioDataLake) diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHiveConfig.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHiveConfig.java index de8264a4cbc8..d901c8e6e2a5 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHiveConfig.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHiveConfig.java @@ -119,7 +119,9 @@ public void testDefaults() .setDeltaLakeCatalogName(null) .setHudiCatalogName(null) .setAutoPurge(false) - .setPartitionProjectionEnabled(false)); + .setPartitionProjectionEnabled(false) + .setS3StorageClassFilter(S3StorageClassFilter.READ_ALL) + .setMetadataParallelism(8)); } @Test @@ -207,6 +209,8 @@ public void testExplicitPropertyMappings() .put("hive.hudi-catalog-name", "hudi") .put("hive.auto-purge", "true") .put(CONFIGURATION_HIVE_PARTITION_PROJECTION_ENABLED, "true") + .put("hive.s3.storage-class-filter", "READ_NON_GLACIER_AND_RESTORED") + .put("hive.metadata.parallelism", "10") .buildOrThrow(); HiveConfig expected = new HiveConfig() @@ -290,7 +294,9 @@ public void testExplicitPropertyMappings() .setDeltaLakeCatalogName("delta") .setHudiCatalogName("hudi") .setAutoPurge(true) - .setPartitionProjectionEnabled(true); + .setPartitionProjectionEnabled(true) + .setS3StorageClassFilter(S3StorageClassFilter.READ_NON_GLACIER_AND_RESTORED) + .setMetadataParallelism(10); assertFullMapping(properties, expected); } diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHiveCreateExternalTable.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHiveCreateExternalTable.java index 6ca1d28257b5..034b1d6a910f 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHiveCreateExternalTable.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHiveCreateExternalTable.java @@ -15,6 +15,7 @@ import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; +import io.trino.filesystem.Location; import io.trino.testing.AbstractTestQueryFramework; import io.trino.testing.MaterializedResult; import io.trino.testing.QueryRunner; @@ -23,9 +24,8 @@ import java.io.IOException; import java.nio.file.Path; +import java.util.UUID; -import static com.google.common.io.MoreFiles.deleteRecursively; -import static com.google.common.io.RecursiveDeleteOption.ALLOW_INSECURE; import static io.trino.testing.QueryAssertions.assertEqualsIgnoreOrder; import static io.trino.tpch.TpchTable.CUSTOMER; import static io.trino.tpch.TpchTable.ORDERS; @@ -50,14 +50,13 @@ protected QueryRunner createQueryRunner() public void testCreateExternalTableWithData() throws IOException { - Path tempDir = createTempDirectory(null); - Path tableLocation = tempDir.resolve("data"); + Location tempDir = Location.of("local:///temp_" + UUID.randomUUID()); @Language("SQL") String createTableSql = format("" + "CREATE TABLE test_create_external " + "WITH (external_location = '%s') AS " + "SELECT * FROM tpch.tiny.nation", - tableLocation.toUri().toASCIIString()); + tempDir); assertUpdate(createTableSql, 25); @@ -67,10 +66,9 @@ public void testCreateExternalTableWithData() MaterializedResult result = computeActual("SELECT DISTINCT regexp_replace(\"$path\", '/[^/]*$', '/') FROM test_create_external"); String tablePath = (String) result.getOnlyValue(); - assertThat(tablePath).startsWith(tableLocation.toFile().toURI().toString()); + assertThat(tablePath).startsWith(tempDir.toString()); assertUpdate("DROP TABLE test_create_external"); - deleteRecursively(tempDir, ALLOW_INSECURE); } @Test diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHiveCreateSchemaInternalRetry.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHiveCreateSchemaInternalRetry.java deleted file mode 100644 index 746be09823fd..000000000000 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHiveCreateSchemaInternalRetry.java +++ /dev/null @@ -1,102 +0,0 @@ -/* - * 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.hive; - -import com.google.common.collect.ImmutableMap; -import io.trino.plugin.hive.metastore.Database; -import io.trino.plugin.hive.metastore.HiveMetastore; -import io.trino.plugin.hive.metastore.HiveMetastoreConfig; -import io.trino.plugin.hive.metastore.file.FileHiveMetastore; -import io.trino.plugin.hive.metastore.file.FileHiveMetastoreConfig; -import io.trino.testing.AbstractTestQueryFramework; -import io.trino.testing.QueryRunner; -import org.testng.annotations.AfterClass; -import org.testng.annotations.Test; - -import java.io.IOException; - -import static io.trino.plugin.hive.HiveMetadata.PRESTO_QUERY_ID_NAME; -import static io.trino.plugin.hive.HiveTestUtils.HDFS_ENVIRONMENT; -import static io.trino.testing.TestingNames.randomNameSuffix; - -public class TestHiveCreateSchemaInternalRetry - extends AbstractTestQueryFramework -{ - private static final String TEST_SCHEMA_TIMEOUT = "test_hive_schema_" + randomNameSuffix(); - private static final String TEST_SCHEMA_DIFFERENT_SESSION = "test_hive_schema_" + randomNameSuffix(); - - private HiveMetastore metastore; - - @Override - protected QueryRunner createQueryRunner() - throws Exception - { - return HiveQueryRunner.builder() - .setCreateTpchSchemas(false) - .setMetastore(distributedQueryRunner -> metastore = createMetastore(distributedQueryRunner.getCoordinator().getBaseDataDir().resolve("hive_data").toString())) - .build(); - } - - private FileHiveMetastore createMetastore(String dataDirectory) - { - return new FileHiveMetastore( - new NodeVersion("testversion"), - HDFS_ENVIRONMENT, - new HiveMetastoreConfig().isHideDeltaLakeTables(), - new FileHiveMetastoreConfig() - .setCatalogDirectory(dataDirectory) - .setMetastoreUser("test")) - { - @Override - public synchronized void createDatabase(Database database) - { - if (database.getDatabaseName().equals(TEST_SCHEMA_DIFFERENT_SESSION)) { - // By modifying query id test simulates that schema was created from different session. - database = Database.builder(database) - .setParameters(ImmutableMap.of(PRESTO_QUERY_ID_NAME, "new_query_id")) - .build(); - } - // Simulate retry mechanism with timeout failure. - // 1. createDatabase correctly create schema but timeout is triggered - // 2. Retry to createDatabase throws SchemaAlreadyExistsException - super.createDatabase(database); - throw new SchemaAlreadyExistsException(database.getDatabaseName()); - } - }; - } - - @AfterClass(alwaysRun = true) - public void tearDown() - throws IOException - { - if (metastore != null) { - metastore.dropDatabase(TEST_SCHEMA_TIMEOUT, false); - metastore.dropDatabase(TEST_SCHEMA_DIFFERENT_SESSION, false); - } - } - - @Test - public void testSchemaCreationWithTimeout() - { - assertQuerySucceeds("CREATE SCHEMA " + TEST_SCHEMA_TIMEOUT); - assertQuery("SHOW SCHEMAS LIKE '" + TEST_SCHEMA_TIMEOUT + "'", "VALUES ('" + TEST_SCHEMA_TIMEOUT + "')"); - } - - @Test - public void testSchemaCreationFailsWhenCreatedWithDifferentSession() - { - assertQueryFails("CREATE SCHEMA " + TEST_SCHEMA_DIFFERENT_SESSION, "Schema already exists: '" + TEST_SCHEMA_DIFFERENT_SESSION + "'"); - assertQuery("SHOW SCHEMAS LIKE '" + TEST_SCHEMA_DIFFERENT_SESSION + "'", "VALUES ('" + TEST_SCHEMA_DIFFERENT_SESSION + "')"); - } -} diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHiveFileFormats.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHiveFileFormats.java index 80bb13ecc15f..6d0f84852f00 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHiveFileFormats.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHiveFileFormats.java @@ -14,13 +14,14 @@ package io.trino.plugin.hive; import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; import com.google.common.collect.Lists; -import io.airlift.compress.lzo.LzoCodec; -import io.airlift.compress.lzo.LzopCodec; import io.trino.filesystem.Location; +import io.trino.filesystem.TrinoFileSystem; import io.trino.filesystem.TrinoFileSystemFactory; import io.trino.filesystem.hdfs.HdfsFileSystemFactory; +import io.trino.filesystem.memory.MemoryFileSystemFactory; import io.trino.hive.formats.compression.CompressionKind; import io.trino.orc.OrcReaderOptions; import io.trino.orc.OrcWriterOptions; @@ -47,45 +48,32 @@ import io.trino.plugin.hive.rcfile.RcFilePageSourceFactory; import io.trino.spi.connector.ConnectorPageSource; import io.trino.spi.connector.ConnectorSession; -import io.trino.spi.connector.RecordCursor; -import io.trino.spi.connector.RecordPageSource; import io.trino.spi.predicate.TupleDomain; +import io.trino.spi.security.ConnectorIdentity; import io.trino.spi.type.Type; +import io.trino.spi.type.VarcharType; import io.trino.testing.TestingConnectorSession; -import org.apache.hadoop.conf.Configuration; import org.apache.hadoop.hive.common.type.HiveVarchar; -import org.apache.hadoop.hive.ql.io.SymlinkTextInputFormat; -import org.apache.hadoop.hive.serde2.objectinspector.ListObjectInspector; -import org.apache.hadoop.hive.serde2.objectinspector.MapObjectInspector; -import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspector; -import org.apache.hadoop.hive.serde2.objectinspector.PrimitiveObjectInspector; -import org.apache.hadoop.hive.serde2.objectinspector.PrimitiveObjectInspector.PrimitiveCategory; -import org.apache.hadoop.hive.serde2.objectinspector.StructField; -import org.apache.hadoop.hive.serde2.objectinspector.StructObjectInspector; -import org.apache.hadoop.hive.serde2.typeinfo.VarcharTypeInfo; -import org.apache.hadoop.mapred.FileSplit; import org.testng.annotations.BeforeClass; import org.testng.annotations.DataProvider; import org.testng.annotations.Test; -import java.io.File; import java.io.IOException; import java.time.Instant; -import java.util.ArrayList; +import java.util.Collections; import java.util.HashSet; import java.util.List; import java.util.Locale; +import java.util.Map; import java.util.Optional; import java.util.OptionalInt; -import java.util.Properties; import java.util.Set; import java.util.TimeZone; -import java.util.stream.Collectors; +import java.util.function.Function; -import static com.google.common.base.Preconditions.checkState; +import static com.google.common.base.Verify.verify; import static com.google.common.collect.ImmutableList.toImmutableList; import static io.airlift.slice.Slices.utf8Slice; -import static io.trino.hadoop.ConfigurationInstantiator.newEmptyConfiguration; import static io.trino.plugin.hive.HivePageSourceProvider.ColumnMapping.buildColumnMappings; import static io.trino.plugin.hive.HiveStorageFormat.AVRO; import static io.trino.plugin.hive.HiveStorageFormat.CSV; @@ -98,13 +86,18 @@ import static io.trino.plugin.hive.HiveStorageFormat.SEQUENCEFILE; import static io.trino.plugin.hive.HiveStorageFormat.TEXTFILE; import static io.trino.plugin.hive.HiveTestUtils.HDFS_ENVIRONMENT; -import static io.trino.plugin.hive.HiveTestUtils.HDFS_FILE_SYSTEM_FACTORY; import static io.trino.plugin.hive.HiveTestUtils.HDFS_FILE_SYSTEM_STATS; import static io.trino.plugin.hive.HiveTestUtils.SESSION; -import static io.trino.plugin.hive.HiveTestUtils.createGenericHiveRecordCursorProvider; import static io.trino.plugin.hive.HiveTestUtils.getHiveSession; -import static io.trino.plugin.hive.HiveTestUtils.getTypes; import static io.trino.plugin.hive.acid.AcidTransaction.NO_ACID_TRANSACTION; +import static io.trino.plugin.hive.util.HiveTypeTranslator.toHiveType; +import static io.trino.plugin.hive.util.SerdeConstants.LIST_COLUMNS; +import static io.trino.plugin.hive.util.SerdeConstants.LIST_COLUMN_TYPES; +import static io.trino.spi.type.RowType.field; +import static io.trino.spi.type.RowType.rowType; +import static io.trino.spi.type.TinyintType.TINYINT; +import static io.trino.spi.type.VarcharType.VARCHAR; +import static io.trino.spi.type.VarcharType.createVarcharType; import static io.trino.testing.StructuralTestUtil.rowBlockOf; import static io.trino.testing.assertions.TrinoExceptionAssert.assertTrinoExceptionThrownBy; import static io.trino.type.InternalTypeManager.TESTING_TYPE_MANAGER; @@ -112,13 +105,9 @@ import static java.util.Objects.requireNonNull; import static java.util.stream.Collectors.toList; import static org.apache.hadoop.hive.metastore.api.hive_metastoreConstants.FILE_INPUT_FORMAT; -import static org.apache.hadoop.hive.serde.serdeConstants.SERIALIZATION_LIB; -import static org.apache.hadoop.hive.serde2.objectinspector.ObjectInspectorFactory.getStandardStructObjectInspector; -import static org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorFactory.getPrimitiveJavaObjectInspector; -import static org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorFactory.javaStringObjectInspector; +import static org.assertj.core.api.Assertions.assertThat; import static org.joda.time.DateTimeZone.UTC; import static org.testng.Assert.assertEquals; -import static org.testng.Assert.assertNotNull; import static org.testng.Assert.assertTrue; // Failing on multiple threads because of org.apache.hadoop.hive.ql.io.parquet.write.ParquetRecordWriterWrapper @@ -169,9 +158,8 @@ public void testTextFile(int rowCount, long fileSizePadding) .withColumns(testColumns) .withRowsCount(rowCount) .withFileSizePadding(fileSizePadding) - .withFileWriterFactory(new SimpleTextFileWriterFactory(HDFS_FILE_SYSTEM_FACTORY, TESTING_TYPE_MANAGER)) - .isReadableByRecordCursor(createGenericHiveRecordCursorProvider(HDFS_ENVIRONMENT)) - .isReadableByPageSource(new SimpleTextFilePageSourceFactory(HDFS_FILE_SYSTEM_FACTORY, STATS, new HiveConfig())); + .withFileWriterFactory(fileSystemFactory -> new SimpleTextFileWriterFactory(fileSystemFactory, TESTING_TYPE_MANAGER)) + .isReadableByPageSource(fileSystemFactory -> new SimpleTextFilePageSourceFactory(fileSystemFactory, new HiveConfig())); } @Test(dataProvider = "validRowAndFileSizePadding") @@ -187,9 +175,8 @@ public void testSequenceFile(int rowCount, long fileSizePadding) .withColumns(testColumns) .withRowsCount(rowCount) .withFileSizePadding(fileSizePadding) - .withFileWriterFactory(new SimpleSequenceFileWriterFactory(HDFS_FILE_SYSTEM_FACTORY, TESTING_TYPE_MANAGER, new NodeVersion("test"))) - .isReadableByRecordCursor(createGenericHiveRecordCursorProvider(HDFS_ENVIRONMENT)) - .isReadableByPageSource(new SimpleSequenceFilePageSourceFactory(HDFS_FILE_SYSTEM_FACTORY, STATS, new HiveConfig())); + .withFileWriterFactory(fileSystemFactory -> new SimpleSequenceFileWriterFactory(fileSystemFactory, TESTING_TYPE_MANAGER, new NodeVersion("test"))) + .isReadableByPageSource(fileSystemFactory -> new SimpleSequenceFilePageSourceFactory(fileSystemFactory, new HiveConfig())); } @Test(dataProvider = "validRowAndFileSizePadding") @@ -198,7 +185,7 @@ public void testCsvFile(int rowCount, long fileSizePadding) { List testColumns = TEST_COLUMNS.stream() // CSV table only support Hive string columns. Notice that CSV does not allow to store null, it uses an empty string instead. - .filter(column -> column.isPartitionKey() || ("string".equals(column.getType()) && !column.getName().contains("_null_"))) + .filter(column -> column.partitionKey() || (column.type() instanceof VarcharType varcharType && varcharType.isUnbounded() && !column.name().contains("_null_"))) .collect(toImmutableList()); assertTrue(testColumns.size() > 5); @@ -207,9 +194,8 @@ public void testCsvFile(int rowCount, long fileSizePadding) .withColumns(testColumns) .withRowsCount(rowCount) .withFileSizePadding(fileSizePadding) - .withFileWriterFactory(new CsvFileWriterFactory(HDFS_FILE_SYSTEM_FACTORY, TESTING_TYPE_MANAGER)) - .isReadableByRecordCursor(createGenericHiveRecordCursorProvider(HDFS_ENVIRONMENT)) - .isReadableByPageSource(new CsvPageSourceFactory(HDFS_FILE_SYSTEM_FACTORY, STATS, new HiveConfig())); + .withFileWriterFactory(fileSystemFactory -> new CsvFileWriterFactory(fileSystemFactory, TESTING_TYPE_MANAGER)) + .isReadableByPageSource(fileSystemFactory -> new CsvPageSourceFactory(fileSystemFactory, new HiveConfig())); } @Test @@ -218,12 +204,11 @@ public void testCsvFileWithNullAndValue() { assertThatFileFormat(CSV) .withColumns(ImmutableList.of( - new TestColumn("t_null_string", javaStringObjectInspector, null, utf8Slice("")), // null was converted to empty string! - new TestColumn("t_string", javaStringObjectInspector, "test", utf8Slice("test")))) + new TestColumn("t_null_string", VARCHAR, null, utf8Slice("")), // null was converted to empty string! + new TestColumn("t_string", VARCHAR, "test", utf8Slice("test")))) .withRowsCount(2) - .withFileWriterFactory(new CsvFileWriterFactory(HDFS_FILE_SYSTEM_FACTORY, TESTING_TYPE_MANAGER)) - .isReadableByRecordCursor(createGenericHiveRecordCursorProvider(HDFS_ENVIRONMENT)) - .isReadableByPageSource(new CsvPageSourceFactory(HDFS_FILE_SYSTEM_FACTORY, STATS, new HiveConfig())); + .withFileWriterFactory(fileSystemFactory -> new CsvFileWriterFactory(fileSystemFactory, TESTING_TYPE_MANAGER)) + .isReadableByPageSource(fileSystemFactory -> new CsvPageSourceFactory(fileSystemFactory, new HiveConfig())); } @Test(dataProvider = "validRowAndFileSizePadding") @@ -232,29 +217,28 @@ public void testJson(int rowCount, long fileSizePadding) { List testColumns = TEST_COLUMNS.stream() // binary is not supported - .filter(column -> !column.getName().equals("t_binary")) + .filter(column -> !column.name().equals("t_binary")) // non-string map keys are not supported - .filter(column -> !column.getName().equals("t_map_tinyint")) - .filter(column -> !column.getName().equals("t_map_smallint")) - .filter(column -> !column.getName().equals("t_map_int")) - .filter(column -> !column.getName().equals("t_map_bigint")) - .filter(column -> !column.getName().equals("t_map_float")) - .filter(column -> !column.getName().equals("t_map_double")) + .filter(column -> !column.name().equals("t_map_tinyint")) + .filter(column -> !column.name().equals("t_map_smallint")) + .filter(column -> !column.name().equals("t_map_int")) + .filter(column -> !column.name().equals("t_map_bigint")) + .filter(column -> !column.name().equals("t_map_float")) + .filter(column -> !column.name().equals("t_map_double")) // null map keys are not supported .filter(TestHiveFileFormats::withoutNullMapKeyTests) // decimal(38) is broken or not supported - .filter(column -> !column.getName().equals("t_decimal_precision_38")) - .filter(column -> !column.getName().equals("t_map_decimal_precision_38")) - .filter(column -> !column.getName().equals("t_array_decimal_precision_38")) + .filter(column -> !column.name().equals("t_decimal_precision_38")) + .filter(column -> !column.name().equals("t_map_decimal_precision_38")) + .filter(column -> !column.name().equals("t_array_decimal_precision_38")) .collect(toList()); assertThatFileFormat(JSON) .withColumns(testColumns) .withRowsCount(rowCount) .withFileSizePadding(fileSizePadding) - .withFileWriterFactory(new JsonFileWriterFactory(HDFS_FILE_SYSTEM_FACTORY, TESTING_TYPE_MANAGER)) - .isReadableByRecordCursor(createGenericHiveRecordCursorProvider(HDFS_ENVIRONMENT)) - .isReadableByPageSource(new JsonPageSourceFactory(HDFS_FILE_SYSTEM_FACTORY, STATS, new HiveConfig())); + .withFileWriterFactory(fileSystemFactory -> new JsonFileWriterFactory(fileSystemFactory, TESTING_TYPE_MANAGER)) + .isReadableByPageSource(fileSystemFactory -> new JsonPageSourceFactory(fileSystemFactory, new HiveConfig())); } @Test(dataProvider = "validRowAndFileSizePadding") @@ -272,8 +256,8 @@ public void testOpenXJson(int rowCount, long fileSizePadding) .withFileSizePadding(fileSizePadding) // openx serde is not available for testing .withSkipGenericWriterTest() - .withFileWriterFactory(new OpenXJsonFileWriterFactory(HDFS_FILE_SYSTEM_FACTORY, TESTING_TYPE_MANAGER)) - .isReadableByPageSource(new OpenXJsonPageSourceFactory(HDFS_FILE_SYSTEM_FACTORY, STATS, new HiveConfig())); + .withFileWriterFactory(fileSystemFactory -> new OpenXJsonFileWriterFactory(fileSystemFactory, TESTING_TYPE_MANAGER)) + .isReadableByPageSource(fileSystemFactory -> new OpenXJsonPageSourceFactory(fileSystemFactory, new HiveConfig())); } @Test(dataProvider = "validRowAndFileSizePadding") @@ -284,7 +268,7 @@ public void testRcTextPageSource(int rowCount, long fileSizePadding) .withColumns(TEST_COLUMNS) .withRowsCount(rowCount) .withFileSizePadding(fileSizePadding) - .isReadableByPageSource(new RcFilePageSourceFactory(FILE_SYSTEM_FACTORY, STATS, new HiveConfig())); + .isReadableByPageSource(fileSystemFactory -> new RcFilePageSourceFactory(fileSystemFactory, new HiveConfig())); } @Test(dataProvider = "rowCount") @@ -299,9 +283,8 @@ public void testRcTextOptimizedWriter(int rowCount) assertThatFileFormat(RCTEXT) .withColumns(testColumns) .withRowsCount(rowCount) - .withFileWriterFactory(new RcFileFileWriterFactory(FILE_SYSTEM_FACTORY, TESTING_TYPE_MANAGER, new NodeVersion("test"), HIVE_STORAGE_TIME_ZONE)) - .isReadableByRecordCursor(createGenericHiveRecordCursorProvider(HDFS_ENVIRONMENT)) - .isReadableByPageSource(new RcFilePageSourceFactory(FILE_SYSTEM_FACTORY, STATS, new HiveConfig())); + .withFileWriterFactory(fileSystemFactory -> new RcFileFileWriterFactory(fileSystemFactory, TESTING_TYPE_MANAGER, new NodeVersion("test"), HIVE_STORAGE_TIME_ZONE)) + .isReadableByPageSource(fileSystemFactory -> new RcFilePageSourceFactory(fileSystemFactory, new HiveConfig())); } @Test(dataProvider = "rowCount") @@ -311,14 +294,14 @@ public void testRcBinaryPageSource(int rowCount) // RCBinary does not support complex type as key of a map and interprets empty VARCHAR as nulls // Hive binary writers are broken for timestamps List testColumns = TEST_COLUMNS.stream() - .filter(testColumn -> !testColumn.getName().equals("t_empty_varchar")) + .filter(testColumn -> !testColumn.name().equals("t_empty_varchar")) .filter(TestHiveFileFormats::withoutTimestamps) .collect(toList()); assertThatFileFormat(RCBINARY) .withColumns(testColumns) .withRowsCount(rowCount) - .isReadableByPageSource(new RcFilePageSourceFactory(FILE_SYSTEM_FACTORY, STATS, new HiveConfig())); + .isReadableByPageSource(fileSystemFactory -> new RcFilePageSourceFactory(fileSystemFactory, new HiveConfig())); } @Test(dataProvider = "rowCount") @@ -327,7 +310,7 @@ public void testRcBinaryOptimizedWriter(int rowCount) { List testColumns = TEST_COLUMNS.stream() // RCBinary interprets empty VARCHAR as nulls - .filter(testColumn -> !testColumn.getName().equals("t_empty_varchar")) + .filter(testColumn -> !testColumn.name().equals("t_empty_varchar")) // t_map_null_key_* must be disabled because Trino cannot produce maps with null keys so the writer will throw .filter(TestHiveFileFormats::withoutNullMapKeyTests) .collect(toList()); @@ -342,10 +325,9 @@ public void testRcBinaryOptimizedWriter(int rowCount) .withRowsCount(rowCount) // generic Hive writer corrupts timestamps .withSkipGenericWriterTest() - .withFileWriterFactory(new RcFileFileWriterFactory(FILE_SYSTEM_FACTORY, TESTING_TYPE_MANAGER, new NodeVersion("test"), HIVE_STORAGE_TIME_ZONE)) - .isReadableByPageSource(new RcFilePageSourceFactory(FILE_SYSTEM_FACTORY, STATS, new HiveConfig())) - .withColumns(testColumnsNoTimestamps) - .isReadableByRecordCursor(createGenericHiveRecordCursorProvider(HDFS_ENVIRONMENT)); + .withFileWriterFactory(fileSystemFactory -> new RcFileFileWriterFactory(fileSystemFactory, TESTING_TYPE_MANAGER, new NodeVersion("test"), HIVE_STORAGE_TIME_ZONE)) + .isReadableByPageSource(fileSystemFactory -> new RcFilePageSourceFactory(fileSystemFactory, new HiveConfig())) + .withColumns(testColumnsNoTimestamps); } @Test(dataProvider = "validRowAndFileSizePadding") @@ -356,7 +338,7 @@ public void testOrc(int rowCount, long fileSizePadding) .withColumns(TEST_COLUMNS) .withRowsCount(rowCount) .withFileSizePadding(fileSizePadding) - .isReadableByPageSource(new OrcPageSourceFactory(new OrcReaderOptions(), HDFS_FILE_SYSTEM_FACTORY, STATS, UTC)); + .isReadableByPageSource(fileSystemFactory -> new OrcPageSourceFactory(new OrcReaderOptions(), fileSystemFactory, STATS, UTC)); } @Test(dataProvider = "validRowAndFileSizePadding") @@ -365,7 +347,6 @@ public void testOrcOptimizedWriter(int rowCount, long fileSizePadding) { HiveSessionProperties hiveSessionProperties = new HiveSessionProperties( new HiveConfig(), - new HiveFormatsConfig(), new OrcReaderConfig(), new OrcWriterConfig() .setValidationPercentage(100.0), @@ -385,9 +366,8 @@ public void testOrcOptimizedWriter(int rowCount, long fileSizePadding) .withRowsCount(rowCount) .withSession(session) .withFileSizePadding(fileSizePadding) - .withFileWriterFactory(new OrcFileWriterFactory(TESTING_TYPE_MANAGER, new NodeVersion("test"), STATS, new OrcWriterOptions(), HDFS_FILE_SYSTEM_FACTORY)) - .isReadableByRecordCursor(createGenericHiveRecordCursorProvider(HDFS_ENVIRONMENT)) - .isReadableByPageSource(new OrcPageSourceFactory(new OrcReaderOptions(), HDFS_FILE_SYSTEM_FACTORY, STATS, UTC)); + .withFileWriterFactory(fileSystemFactory -> new OrcFileWriterFactory(TESTING_TYPE_MANAGER, new NodeVersion("test"), STATS, new OrcWriterOptions(), fileSystemFactory)) + .isReadableByPageSource(fileSystemFactory -> new OrcPageSourceFactory(new OrcReaderOptions(), fileSystemFactory, STATS, UTC)); } @Test(dataProvider = "rowCount") @@ -406,7 +386,7 @@ public void testOrcUseColumnNames(int rowCount) .withRowsCount(rowCount) .withReadColumns(Lists.reverse(testColumns)) .withSession(session) - .isReadableByPageSource(new OrcPageSourceFactory(new OrcReaderOptions(), HDFS_FILE_SYSTEM_FACTORY, STATS, UTC)); + .isReadableByPageSource(fileSystemFactory -> new OrcPageSourceFactory(new OrcReaderOptions(), fileSystemFactory, STATS, UTC)); } @Test(dataProvider = "rowCount") @@ -414,7 +394,7 @@ public void testOrcUseColumnNameLowerCaseConversion(int rowCount) throws Exception { List testColumnsUpperCase = TEST_COLUMNS.stream() - .map(testColumn -> new TestColumn(testColumn.getName().toUpperCase(Locale.ENGLISH), testColumn.getObjectInspector(), testColumn.getWriteValue(), testColumn.getExpectedValue(), testColumn.isPartitionKey())) + .map(testColumn -> testColumn.withName(testColumn.name().toUpperCase(Locale.ENGLISH))) .collect(toList()); ConnectorSession session = getHiveSession(new HiveConfig(), new OrcReaderConfig().setUseColumnNames(true)); @@ -423,7 +403,7 @@ public void testOrcUseColumnNameLowerCaseConversion(int rowCount) .withRowsCount(rowCount) .withReadColumns(TEST_COLUMNS) .withSession(session) - .isReadableByPageSource(new OrcPageSourceFactory(new OrcReaderOptions(), HDFS_FILE_SYSTEM_FACTORY, STATS, UTC)); + .isReadableByPageSource(fileSystemFactory -> new OrcPageSourceFactory(new OrcReaderOptions(), fileSystemFactory, STATS, UTC)); } @Test(dataProvider = "validRowAndFileSizePadding") @@ -434,39 +414,28 @@ public void testAvro(int rowCount, long fileSizePadding) .withColumns(getTestColumnsSupportedByAvro()) .withRowsCount(rowCount) .withFileSizePadding(fileSizePadding) - .withFileWriterFactory(new AvroFileWriterFactory(FILE_SYSTEM_FACTORY, TESTING_TYPE_MANAGER, new NodeVersion("test_version"))) - .isReadableByPageSource(new AvroPageSourceFactory(FILE_SYSTEM_FACTORY, STATS)) - .isReadableByRecordCursor(createGenericHiveRecordCursorProvider(HDFS_ENVIRONMENT)); + .withFileWriterFactory(fileSystemFactory -> new AvroFileWriterFactory(fileSystemFactory, TESTING_TYPE_MANAGER, new NodeVersion("test_version"))) + .isReadableByPageSource(AvroPageSourceFactory::new); } @Test(dataProvider = "rowCount") public void testAvroFileInSymlinkTable(int rowCount) throws Exception { - File file = File.createTempFile("trino_test", AVRO.name()); - //noinspection ResultOfMethodCallIgnored - file.delete(); - try { - FileSplit split = createTestFileHive(file.getAbsolutePath(), AVRO, HiveCompressionCodec.NONE, getTestColumnsSupportedByAvro(), rowCount); - Properties splitProperties = new Properties(); - splitProperties.setProperty(FILE_INPUT_FORMAT, SymlinkTextInputFormat.class.getName()); - splitProperties.setProperty(SERIALIZATION_LIB, AVRO.getSerde()); - testCursorProvider(createGenericHiveRecordCursorProvider(HDFS_ENVIRONMENT), split, splitProperties, getTestColumnsSupportedByAvro(), SESSION, file.length(), rowCount); - testPageSourceFactory(new AvroPageSourceFactory(FILE_SYSTEM_FACTORY, STATS), split, AVRO, getTestColumnsSupportedByAvro(), SESSION, file.length(), rowCount); - } - finally { - //noinspection ResultOfMethodCallIgnored - file.delete(); - } + TrinoFileSystemFactory fileSystemFactory = new MemoryFileSystemFactory(); + Location location = Location.of("memory:///avro-test"); + createTestFileHive(fileSystemFactory, location, AVRO, HiveCompressionCodec.NONE, getTestColumnsSupportedByAvro(), rowCount); + long fileSize = fileSystemFactory.create(ConnectorIdentity.ofUser("test")).newInputFile(location).length(); + testPageSourceFactory(new AvroPageSourceFactory(fileSystemFactory), location, AVRO, getTestColumnsSupportedByAvro(), SESSION, fileSize, fileSize, rowCount); } private static List getTestColumnsSupportedByAvro() { // Avro only supports String for Map keys, and doesn't support smallint or tinyint. return TEST_COLUMNS.stream() - .filter(column -> !column.getName().startsWith("t_map_") || column.getName().equals("t_map_string")) - .filter(column -> !column.getName().endsWith("_smallint")) - .filter(column -> !column.getName().endsWith("_tinyint")) + .filter(column -> !column.name().startsWith("t_map_") || column.name().equals("t_map_string")) + .filter(column -> !column.name().endsWith("_smallint")) + .filter(column -> !column.name().endsWith("_tinyint")) .collect(toList()); } @@ -480,7 +449,7 @@ public void testParquetPageSource(int rowCount, long fileSizePadding) .withSession(PARQUET_SESSION) .withRowsCount(rowCount) .withFileSizePadding(fileSizePadding) - .isReadableByPageSource(PARQUET_PAGE_SOURCE_FACTORY); + .isReadableByPageSource(fileSystemFactory -> new ParquetPageSourceFactory(fileSystemFactory, STATS, new ParquetReaderConfig(), new HiveConfig())); } @Test(dataProvider = "validRowAndFileSizePadding") @@ -494,7 +463,7 @@ public void testParquetPageSourceGzip(int rowCount, long fileSizePadding) .withCompressionCodec(HiveCompressionCodec.GZIP) .withFileSizePadding(fileSizePadding) .withRowsCount(rowCount) - .isReadableByPageSource(PARQUET_PAGE_SOURCE_FACTORY); + .isReadableByPageSource(fileSystemFactory -> new ParquetPageSourceFactory(fileSystemFactory, STATS, new ParquetReaderConfig(), new HiveConfig())); } @Test(dataProvider = "rowCount") @@ -508,8 +477,8 @@ public void testParquetWriter(int rowCount) .withSession(session) .withColumns(testColumns) .withRowsCount(rowCount) - .withFileWriterFactory(new ParquetFileWriterFactory(HDFS_FILE_SYSTEM_FACTORY, new NodeVersion("test-version"), TESTING_TYPE_MANAGER, new HiveConfig(), STATS)) - .isReadableByPageSource(PARQUET_PAGE_SOURCE_FACTORY); + .withFileWriterFactory(fileSystemFactory -> new ParquetFileWriterFactory(fileSystemFactory, new NodeVersion("test-version"), TESTING_TYPE_MANAGER, new HiveConfig(), STATS)) + .isReadableByPageSource(fileSystemFactory -> new ParquetPageSourceFactory(fileSystemFactory, STATS, new ParquetReaderConfig(), new HiveConfig())); } @Test(dataProvider = "rowCount") @@ -518,29 +487,24 @@ public void testParquetPageSourceSchemaEvolution(int rowCount) { List writeColumns = getTestColumnsSupportedByParquet(); - // test index-based access + // test the index-based access List readColumns = writeColumns.stream() - .map(column -> new TestColumn( - column.getName() + "_new", - column.getObjectInspector(), - column.getWriteValue(), - column.getExpectedValue(), - column.isPartitionKey())) + .map(column -> column.withName(column.name() + "_new")) .collect(toList()); assertThatFileFormat(PARQUET) .withWriteColumns(writeColumns) .withReadColumns(readColumns) .withSession(PARQUET_SESSION) .withRowsCount(rowCount) - .isReadableByPageSource(PARQUET_PAGE_SOURCE_FACTORY); + .isReadableByPageSource(fileSystemFactory -> new ParquetPageSourceFactory(fileSystemFactory, STATS, new ParquetReaderConfig(), new HiveConfig())); - // test name-based access + // test the name-based access readColumns = Lists.reverse(writeColumns); assertThatFileFormat(PARQUET) .withWriteColumns(writeColumns) .withReadColumns(readColumns) .withSession(PARQUET_SESSION_USE_NAME) - .isReadableByPageSource(PARQUET_PAGE_SOURCE_FACTORY); + .isReadableByPageSource(fileSystemFactory -> new ParquetPageSourceFactory(fileSystemFactory, STATS, new ParquetReaderConfig(), new HiveConfig())); } private static List getTestColumnsSupportedByParquet() @@ -552,12 +516,9 @@ private static List getTestColumnsSupportedByParquet() return TEST_COLUMNS.stream() .filter(TestHiveFileFormats::withoutTimestamps) .filter(TestHiveFileFormats::withoutNullMapKeyTests) - .filter(column -> !column.getName().equals("t_null_array_int")) - .filter(column -> !column.getName().equals("t_array_empty")) - .filter(column -> column.isPartitionKey() || ( - !hasType(column.getObjectInspector(), PrimitiveCategory.DATE)) && - !hasType(column.getObjectInspector(), PrimitiveCategory.SHORT) && - !hasType(column.getObjectInspector(), PrimitiveCategory.BYTE)) + .filter(column -> !column.name().equals("t_null_array_int")) + .filter(column -> !column.name().equals("t_array_empty")) + .filter(column -> column.partitionKey() || !hasType(column.type(), TINYINT)) .collect(toList()); } @@ -565,52 +526,47 @@ private static List getTestColumnsSupportedByParquet() public void testTruncateVarcharColumn() throws Exception { - TestColumn writeColumn = new TestColumn("varchar_column", getPrimitiveJavaObjectInspector(new VarcharTypeInfo(4)), new HiveVarchar("test", 4), utf8Slice("test")); - TestColumn readColumn = new TestColumn("varchar_column", getPrimitiveJavaObjectInspector(new VarcharTypeInfo(3)), new HiveVarchar("tes", 3), utf8Slice("tes")); + TestColumn writeColumn = new TestColumn("varchar_column", createVarcharType(4), new HiveVarchar("test", 4), utf8Slice("test")); + TestColumn readColumn = new TestColumn("varchar_column", createVarcharType(3), new HiveVarchar("tes", 3), utf8Slice("tes")); assertThatFileFormat(RCTEXT) .withWriteColumns(ImmutableList.of(writeColumn)) .withReadColumns(ImmutableList.of(readColumn)) - .isReadableByPageSource(new RcFilePageSourceFactory(FILE_SYSTEM_FACTORY, STATS, new HiveConfig())) - .isReadableByRecordCursor(createGenericHiveRecordCursorProvider(HDFS_ENVIRONMENT)); + .isReadableByPageSource(fileSystemFactory -> new RcFilePageSourceFactory(fileSystemFactory, new HiveConfig())); assertThatFileFormat(RCBINARY) .withWriteColumns(ImmutableList.of(writeColumn)) .withReadColumns(ImmutableList.of(readColumn)) - .isReadableByPageSource(new RcFilePageSourceFactory(FILE_SYSTEM_FACTORY, STATS, new HiveConfig())) - .isReadableByRecordCursor(createGenericHiveRecordCursorProvider(HDFS_ENVIRONMENT)); + .isReadableByPageSource(fileSystemFactory -> new RcFilePageSourceFactory(fileSystemFactory, new HiveConfig())); assertThatFileFormat(ORC) .withWriteColumns(ImmutableList.of(writeColumn)) .withReadColumns(ImmutableList.of(readColumn)) - .isReadableByPageSource(new OrcPageSourceFactory(new OrcReaderOptions(), HDFS_FILE_SYSTEM_FACTORY, STATS, UTC)); + .isReadableByPageSource(fileSystemFactory -> new OrcPageSourceFactory(new OrcReaderOptions(), fileSystemFactory, STATS, UTC)); assertThatFileFormat(PARQUET) .withWriteColumns(ImmutableList.of(writeColumn)) .withReadColumns(ImmutableList.of(readColumn)) .withSession(PARQUET_SESSION) - .isReadableByPageSource(PARQUET_PAGE_SOURCE_FACTORY); + .isReadableByPageSource(fileSystemFactory -> new ParquetPageSourceFactory(fileSystemFactory, STATS, new ParquetReaderConfig(), new HiveConfig())); assertThatFileFormat(AVRO) .withWriteColumns(ImmutableList.of(writeColumn)) .withReadColumns(ImmutableList.of(readColumn)) - .withFileWriterFactory(new AvroFileWriterFactory(FILE_SYSTEM_FACTORY, TESTING_TYPE_MANAGER, new NodeVersion("test_version"))) - .isReadableByRecordCursor(createGenericHiveRecordCursorProvider(HDFS_ENVIRONMENT)) - .isReadableByPageSource(new AvroPageSourceFactory(FILE_SYSTEM_FACTORY, STATS)); + .withFileWriterFactory(fileSystemFactory -> new AvroFileWriterFactory(fileSystemFactory, TESTING_TYPE_MANAGER, new NodeVersion("test_version"))) + .isReadableByPageSource(AvroPageSourceFactory::new); assertThatFileFormat(SEQUENCEFILE) .withWriteColumns(ImmutableList.of(writeColumn)) .withReadColumns(ImmutableList.of(readColumn)) - .withFileWriterFactory(new SimpleSequenceFileWriterFactory(HDFS_FILE_SYSTEM_FACTORY, TESTING_TYPE_MANAGER, new NodeVersion("test"))) - .isReadableByRecordCursor(createGenericHiveRecordCursorProvider(HDFS_ENVIRONMENT)) - .isReadableByPageSource(new SimpleSequenceFilePageSourceFactory(HDFS_FILE_SYSTEM_FACTORY, STATS, new HiveConfig())); + .withFileWriterFactory(fileSystemFactory -> new SimpleSequenceFileWriterFactory(fileSystemFactory, TESTING_TYPE_MANAGER, new NodeVersion("test"))) + .isReadableByPageSource(fileSystemFactory -> new SimpleSequenceFilePageSourceFactory(fileSystemFactory, new HiveConfig())); assertThatFileFormat(TEXTFILE) .withWriteColumns(ImmutableList.of(writeColumn)) .withReadColumns(ImmutableList.of(readColumn)) - .withFileWriterFactory(new SimpleTextFileWriterFactory(HDFS_FILE_SYSTEM_FACTORY, TESTING_TYPE_MANAGER)) - .isReadableByRecordCursor(createGenericHiveRecordCursorProvider(HDFS_ENVIRONMENT)) - .isReadableByPageSource(new SimpleTextFilePageSourceFactory(HDFS_FILE_SYSTEM_FACTORY, STATS, new HiveConfig())); + .withFileWriterFactory(fileSystemFactory -> new SimpleTextFileWriterFactory(fileSystemFactory, TESTING_TYPE_MANAGER)) + .isReadableByPageSource(fileSystemFactory -> new SimpleTextFilePageSourceFactory(fileSystemFactory, new HiveConfig())); } @Test(dataProvider = "rowCount") @@ -633,9 +589,8 @@ public void testAvroProjectedColumns(int rowCount) .withWriteColumns(writeColumns) .withReadColumns(readColumns) .withRowsCount(rowCount) - .withFileWriterFactory(new AvroFileWriterFactory(FILE_SYSTEM_FACTORY, TESTING_TYPE_MANAGER, new NodeVersion("test_version"))) - .isReadableByRecordCursorPageSource(createGenericHiveRecordCursorProvider(HDFS_ENVIRONMENT)) - .isReadableByPageSource(new AvroPageSourceFactory(FILE_SYSTEM_FACTORY, STATS)); + .withFileWriterFactory(fileSystemFactory -> new AvroFileWriterFactory(fileSystemFactory, TESTING_TYPE_MANAGER, new NodeVersion("test_version"))) + .isReadableByPageSource(AvroPageSourceFactory::new); } @Test(dataProvider = "rowCount") @@ -659,14 +614,14 @@ public void testParquetProjectedColumns(int rowCount) .withReadColumns(readColumns) .withRowsCount(rowCount) .withSession(PARQUET_SESSION) - .isReadableByPageSource(PARQUET_PAGE_SOURCE_FACTORY); + .isReadableByPageSource(fileSystemFactory -> new ParquetPageSourceFactory(fileSystemFactory, STATS, new ParquetReaderConfig(), new HiveConfig())); assertThatFileFormat(PARQUET) .withWriteColumns(writeColumns) .withReadColumns(readColumns) .withRowsCount(rowCount) .withSession(PARQUET_SESSION_USE_NAME) - .isReadableByPageSource(PARQUET_PAGE_SOURCE_FACTORY); + .isReadableByPageSource(fileSystemFactory -> new ParquetPageSourceFactory(fileSystemFactory, STATS, new ParquetReaderConfig(), new HiveConfig())); } @Test(dataProvider = "rowCount") @@ -691,13 +646,13 @@ public void testORCProjectedColumns(int rowCount) .withReadColumns(readColumns) .withRowsCount(rowCount) .withSession(session) - .isReadableByPageSource(new OrcPageSourceFactory(new OrcReaderOptions(), HDFS_FILE_SYSTEM_FACTORY, STATS, UTC)); + .isReadableByPageSource(fileSystemFactory -> new OrcPageSourceFactory(new OrcReaderOptions(), fileSystemFactory, STATS, UTC)); assertThatFileFormat(ORC) .withWriteColumns(writeColumns) .withReadColumns(readColumns) .withRowsCount(rowCount) - .isReadableByPageSource(new OrcPageSourceFactory(new OrcReaderOptions(), HDFS_FILE_SYSTEM_FACTORY, STATS, UTC)); + .isReadableByPageSource(fileSystemFactory -> new OrcPageSourceFactory(new OrcReaderOptions(), fileSystemFactory, STATS, UTC)); } @Test(dataProvider = "rowCount") @@ -705,7 +660,7 @@ public void testSequenceFileProjectedColumns(int rowCount) throws Exception { List supportedColumns = TEST_COLUMNS.stream() - .filter(column -> !column.getName().equals("t_map_null_key_complex_key_value")) + .filter(column -> !column.name().equals("t_map_null_key_complex_key_value")) .collect(toList()); List regularColumns = getRegularColumns(supportedColumns); @@ -723,9 +678,8 @@ public void testSequenceFileProjectedColumns(int rowCount) .withWriteColumns(writeColumns) .withReadColumns(readColumns) .withRowsCount(rowCount) - .withFileWriterFactory(new SimpleSequenceFileWriterFactory(HDFS_FILE_SYSTEM_FACTORY, TESTING_TYPE_MANAGER, new NodeVersion("test"))) - .isReadableByRecordCursor(createGenericHiveRecordCursorProvider(HDFS_ENVIRONMENT)) - .isReadableByPageSource(new SimpleSequenceFilePageSourceFactory(HDFS_FILE_SYSTEM_FACTORY, STATS, new HiveConfig())); + .withFileWriterFactory(fileSystemFactory -> new SimpleSequenceFileWriterFactory(fileSystemFactory, TESTING_TYPE_MANAGER, new NodeVersion("test"))) + .isReadableByPageSource(fileSystemFactory -> new SimpleSequenceFilePageSourceFactory(fileSystemFactory, new HiveConfig())); } @Test(dataProvider = "rowCount") @@ -752,40 +706,8 @@ public void testTextFileProjectedColumns(int rowCount) .withWriteColumns(writeColumns) .withReadColumns(readColumns) .withRowsCount(rowCount) - .withFileWriterFactory(new SimpleTextFileWriterFactory(HDFS_FILE_SYSTEM_FACTORY, TESTING_TYPE_MANAGER)) - .isReadableByRecordCursor(createGenericHiveRecordCursorProvider(HDFS_ENVIRONMENT)) - .isReadableByPageSource(new SimpleTextFilePageSourceFactory(HDFS_FILE_SYSTEM_FACTORY, STATS, new HiveConfig())); - } - - @Test(dataProvider = "rowCount") - public void testRCTextProjectedColumns(int rowCount) - throws Exception - { - List supportedColumns = TEST_COLUMNS.stream() - .filter(testColumn -> { - // TODO: This is a bug in the RC text reader - // RC file does not support complex type as key of a map - return !testColumn.getName().equals("t_struct_null") - && !testColumn.getName().equals("t_map_null_key_complex_key_value"); - }) - .collect(toImmutableList()); - - List regularColumns = getRegularColumns(supportedColumns); - List partitionColumns = getPartitionColumns(supportedColumns); - - // Created projected columns for all regular supported columns - ImmutableList.Builder writeColumnsBuilder = ImmutableList.builder(); - ImmutableList.Builder readeColumnsBuilder = ImmutableList.builder(); - generateProjectedColumns(regularColumns, writeColumnsBuilder, readeColumnsBuilder); - - List writeColumns = writeColumnsBuilder.addAll(partitionColumns).build(); - List readColumns = readeColumnsBuilder.addAll(partitionColumns).build(); - - assertThatFileFormat(RCTEXT) - .withWriteColumns(writeColumns) - .withReadColumns(readColumns) - .withRowsCount(rowCount) - .isReadableByRecordCursorPageSource(createGenericHiveRecordCursorProvider(HDFS_ENVIRONMENT)); + .withFileWriterFactory(fileSystemFactory -> new SimpleTextFileWriterFactory(fileSystemFactory, TESTING_TYPE_MANAGER)) + .isReadableByPageSource(fileSystemFactory -> new SimpleTextFilePageSourceFactory(fileSystemFactory, new HiveConfig())); } @Test(dataProvider = "rowCount") @@ -808,7 +730,7 @@ public void testRCTextProjectedColumnsPageSource(int rowCount) .withWriteColumns(writeColumns) .withReadColumns(readColumns) .withRowsCount(rowCount) - .isReadableByPageSource(new RcFilePageSourceFactory(FILE_SYSTEM_FACTORY, STATS, new HiveConfig())); + .isReadableByPageSource(fileSystemFactory -> new RcFilePageSourceFactory(fileSystemFactory, new HiveConfig())); } @Test(dataProvider = "rowCount") @@ -818,7 +740,7 @@ public void testRCBinaryProjectedColumns(int rowCount) // RCBinary does not support complex type as key of a map and interprets empty VARCHAR as nulls List supportedColumns = TEST_COLUMNS.stream() .filter(testColumn -> { - String name = testColumn.getName(); + String name = testColumn.name(); return !name.equals("t_map_null_key_complex_key_value") && !name.equals("t_empty_varchar"); }) .collect(toList()); @@ -840,8 +762,8 @@ public void testRCBinaryProjectedColumns(int rowCount) .withRowsCount(rowCount) // generic Hive writer corrupts timestamps .withSkipGenericWriterTest() - .withFileWriterFactory(new RcFileFileWriterFactory(FILE_SYSTEM_FACTORY, TESTING_TYPE_MANAGER, new NodeVersion("test"), HIVE_STORAGE_TIME_ZONE)) - .isReadableByPageSource(new RcFilePageSourceFactory(FILE_SYSTEM_FACTORY, STATS, new HiveConfig())); + .withFileWriterFactory(fileSystemFactory -> new RcFileFileWriterFactory(fileSystemFactory, TESTING_TYPE_MANAGER, new NodeVersion("test"), HIVE_STORAGE_TIME_ZONE)) + .isReadableByPageSource(fileSystemFactory -> new RcFilePageSourceFactory(fileSystemFactory, new HiveConfig())); } @Test(dataProvider = "rowCount") @@ -850,7 +772,7 @@ public void testRCBinaryProjectedColumnsPageSource(int rowCount) { // RCBinary does not support complex type as key of a map and interprets empty VARCHAR as nulls List supportedColumns = TEST_COLUMNS.stream() - .filter(testColumn -> !testColumn.getName().equals("t_empty_varchar")) + .filter(testColumn -> !testColumn.name().equals("t_empty_varchar")) .collect(toList()); List regularColumns = getRegularColumns(supportedColumns); @@ -870,16 +792,16 @@ public void testRCBinaryProjectedColumnsPageSource(int rowCount) .withRowsCount(rowCount) // generic Hive writer corrupts timestamps .withSkipGenericWriterTest() - .withFileWriterFactory(new RcFileFileWriterFactory(FILE_SYSTEM_FACTORY, TESTING_TYPE_MANAGER, new NodeVersion("test"), HIVE_STORAGE_TIME_ZONE)) - .isReadableByPageSource(new RcFilePageSourceFactory(FILE_SYSTEM_FACTORY, STATS, new HiveConfig())); + .withFileWriterFactory(fileSystemFactory -> new RcFileFileWriterFactory(fileSystemFactory, TESTING_TYPE_MANAGER, new NodeVersion("test"), HIVE_STORAGE_TIME_ZONE)) + .isReadableByPageSource(fileSystemFactory -> new RcFilePageSourceFactory(fileSystemFactory, new HiveConfig())); } @Test public void testFailForLongVarcharPartitionColumn() throws Exception { - TestColumn partitionColumn = new TestColumn("partition_column", getPrimitiveJavaObjectInspector(new VarcharTypeInfo(3)), "test", utf8Slice("tes"), true); - TestColumn varcharColumn = new TestColumn("varchar_column", getPrimitiveJavaObjectInspector(new VarcharTypeInfo(3)), new HiveVarchar("tes", 3), utf8Slice("tes")); + TestColumn partitionColumn = new TestColumn("partition_column", createVarcharType(3), "test", utf8Slice("tes"), true); + TestColumn varcharColumn = new TestColumn("varchar_column", createVarcharType(3), new HiveVarchar("tes", 3), utf8Slice("tes")); List columns = ImmutableList.of(partitionColumn, varcharColumn); @@ -888,271 +810,104 @@ public void testFailForLongVarcharPartitionColumn() assertThatFileFormat(RCTEXT) .withColumns(columns) - .isFailingForPageSource(new RcFilePageSourceFactory(FILE_SYSTEM_FACTORY, STATS, new HiveConfig()), expectedErrorCode, expectedMessage) - .isFailingForRecordCursor(createGenericHiveRecordCursorProvider(HDFS_ENVIRONMENT), expectedErrorCode, expectedMessage); + .isFailingForPageSource(fileSystemFactory -> new RcFilePageSourceFactory(fileSystemFactory, new HiveConfig()), expectedErrorCode, expectedMessage); assertThatFileFormat(RCBINARY) .withColumns(columns) - .isFailingForPageSource(new RcFilePageSourceFactory(FILE_SYSTEM_FACTORY, STATS, new HiveConfig()), expectedErrorCode, expectedMessage) - .isFailingForRecordCursor(createGenericHiveRecordCursorProvider(HDFS_ENVIRONMENT), expectedErrorCode, expectedMessage); + .isFailingForPageSource(fileSystemFactory -> new RcFilePageSourceFactory(fileSystemFactory, new HiveConfig()), expectedErrorCode, expectedMessage); assertThatFileFormat(ORC) .withColumns(columns) - .isFailingForPageSource(new OrcPageSourceFactory(new OrcReaderOptions(), HDFS_FILE_SYSTEM_FACTORY, STATS, UTC), expectedErrorCode, expectedMessage); + .isFailingForPageSource(fileSystemFactory -> new OrcPageSourceFactory(new OrcReaderOptions(), fileSystemFactory, STATS, UTC), expectedErrorCode, expectedMessage); assertThatFileFormat(PARQUET) .withColumns(columns) .withSession(PARQUET_SESSION) - .isFailingForPageSource(PARQUET_PAGE_SOURCE_FACTORY, expectedErrorCode, expectedMessage); - - assertThatFileFormat(SEQUENCEFILE) - .withColumns(columns) - .isFailingForRecordCursor(createGenericHiveRecordCursorProvider(HDFS_ENVIRONMENT), expectedErrorCode, expectedMessage); - - assertThatFileFormat(TEXTFILE) - .withColumns(columns) - .isFailingForRecordCursor(createGenericHiveRecordCursorProvider(HDFS_ENVIRONMENT), expectedErrorCode, expectedMessage); - } - - private void testRecordPageSource( - HiveRecordCursorProvider cursorProvider, - FileSplit split, - HiveStorageFormat storageFormat, - List testReadColumns, - ConnectorSession session, - long fileSize, - int rowCount) - throws Exception - { - Properties splitProperties = new Properties(); - splitProperties.setProperty(FILE_INPUT_FORMAT, storageFormat.getInputFormat()); - splitProperties.setProperty(SERIALIZATION_LIB, storageFormat.getSerde()); - ConnectorPageSource pageSource = createPageSourceFromCursorProvider(cursorProvider, split, splitProperties, fileSize, testReadColumns, session); - checkPageSource(pageSource, testReadColumns, getTypes(getColumnHandles(testReadColumns)), rowCount); - } - - private void testCursorProvider( - HiveRecordCursorProvider cursorProvider, - FileSplit split, - HiveStorageFormat storageFormat, - List testReadColumns, - ConnectorSession session, - long fileSize, - int rowCount) - { - Properties splitProperties = new Properties(); - splitProperties.setProperty(FILE_INPUT_FORMAT, storageFormat.getInputFormat()); - splitProperties.setProperty(SERIALIZATION_LIB, storageFormat.getSerde()); - testCursorProvider(cursorProvider, split, splitProperties, testReadColumns, session, fileSize, rowCount); - } - - private void testCursorProvider( - HiveRecordCursorProvider cursorProvider, - FileSplit split, - Properties splitProperties, - List testReadColumns, - ConnectorSession session, - long fileSize, - int rowCount) - { - ConnectorPageSource pageSource = createPageSourceFromCursorProvider(cursorProvider, split, splitProperties, fileSize, testReadColumns, session); - RecordCursor cursor = ((RecordPageSource) pageSource).getCursor(); - checkCursor(cursor, testReadColumns, rowCount); - } - - private ConnectorPageSource createPageSourceFromCursorProvider( - HiveRecordCursorProvider cursorProvider, - FileSplit split, - Properties splitProperties, - long fileSize, - List testReadColumns, - ConnectorSession session) - { - // Use full columns in split properties - ImmutableList.Builder splitPropertiesColumnNames = ImmutableList.builder(); - ImmutableList.Builder splitPropertiesColumnTypes = ImmutableList.builder(); - Set baseColumnNames = new HashSet<>(); - - for (TestColumn testReadColumn : testReadColumns) { - String name = testReadColumn.getBaseName(); - if (!baseColumnNames.contains(name) && !testReadColumn.isPartitionKey()) { - baseColumnNames.add(name); - splitPropertiesColumnNames.add(name); - splitPropertiesColumnTypes.add(testReadColumn.getBaseObjectInspector().getTypeName()); - } - } - - splitProperties.setProperty( - "columns", - splitPropertiesColumnNames.build().stream() - .collect(Collectors.joining(","))); - - splitProperties.setProperty( - "columns.types", - splitPropertiesColumnTypes.build().stream() - .collect(Collectors.joining(","))); - - List partitionKeys = testReadColumns.stream() - .filter(TestColumn::isPartitionKey) - .map(input -> new HivePartitionKey(input.getName(), (String) input.getWriteValue())) - .collect(toList()); - - String partitionName = String.join("/", partitionKeys.stream() - .map(partitionKey -> format("%s=%s", partitionKey.getName(), partitionKey.getValue())) - .collect(toImmutableList())); - - Configuration configuration = newEmptyConfiguration(); - configuration.set("io.compression.codecs", LzoCodec.class.getName() + "," + LzopCodec.class.getName()); - - List columnHandles = getColumnHandles(testReadColumns); - List columnMappings = buildColumnMappings( - partitionName, - partitionKeys, - columnHandles, - ImmutableList.of(), - TableToPartitionMapping.empty(), - split.getPath().toString(), - OptionalInt.empty(), - fileSize, - Instant.now().toEpochMilli()); - - Optional pageSource = HivePageSourceProvider.createHivePageSource( - ImmutableSet.of(), - ImmutableSet.of(cursorProvider), - configuration, - session, - Location.of(split.getPath().toString()), - OptionalInt.empty(), - split.getStart(), - split.getLength(), - fileSize, - splitProperties, - TupleDomain.all(), - columnHandles, - TESTING_TYPE_MANAGER, - Optional.empty(), - Optional.empty(), - false, - Optional.empty(), - false, - NO_ACID_TRANSACTION, - columnMappings); - - return pageSource.get(); + .isFailingForPageSource(fileSystemFactory -> new ParquetPageSourceFactory(fileSystemFactory, STATS, new ParquetReaderConfig(), new HiveConfig()), expectedErrorCode, expectedMessage); } private void testPageSourceFactory( HivePageSourceFactory sourceFactory, - FileSplit split, + Location location, HiveStorageFormat storageFormat, List testReadColumns, ConnectorSession session, long fileSize, + long paddedFileSize, int rowCount) throws IOException { - Properties splitProperties = new Properties(); - splitProperties.setProperty(FILE_INPUT_FORMAT, storageFormat.getInputFormat()); - splitProperties.setProperty(SERIALIZATION_LIB, storageFormat.getSerde()); - // Use full columns in split properties ImmutableList.Builder splitPropertiesColumnNames = ImmutableList.builder(); ImmutableList.Builder splitPropertiesColumnTypes = ImmutableList.builder(); Set baseColumnNames = new HashSet<>(); for (TestColumn testReadColumn : testReadColumns) { - String name = testReadColumn.getBaseName(); - if (!baseColumnNames.contains(name) && !testReadColumn.isPartitionKey()) { + String name = testReadColumn.baseName(); + if (!baseColumnNames.contains(name) && !testReadColumn.partitionKey()) { baseColumnNames.add(name); splitPropertiesColumnNames.add(name); - splitPropertiesColumnTypes.add(testReadColumn.getBaseObjectInspector().getTypeName()); + splitPropertiesColumnTypes.add(toHiveType(testReadColumn.baseType()).toString()); } } - splitProperties.setProperty("columns", splitPropertiesColumnNames.build().stream().collect(Collectors.joining(","))); - splitProperties.setProperty("columns.types", splitPropertiesColumnTypes.build().stream().collect(Collectors.joining(","))); + Map splitProperties = ImmutableMap.builder() + .put(FILE_INPUT_FORMAT, storageFormat.getInputFormat()) + .put(LIST_COLUMNS, String.join(",", splitPropertiesColumnNames.build())) + .put(LIST_COLUMN_TYPES, String.join(",", splitPropertiesColumnTypes.build())) + .buildOrThrow(); List partitionKeys = testReadColumns.stream() - .filter(TestColumn::isPartitionKey) - .map(input -> new HivePartitionKey(input.getName(), (String) input.getWriteValue())) + .filter(TestColumn::partitionKey) + .map(input -> new HivePartitionKey(input.name(), (String) input.writeValue())) .collect(toList()); String partitionName = String.join("/", partitionKeys.stream() .map(partitionKey -> format("%s=%s", partitionKey.getName(), partitionKey.getValue())) .collect(toImmutableList())); - List columnHandles = getColumnHandles(testReadColumns); - List columnMappings = buildColumnMappings( partitionName, partitionKeys, - columnHandles, + getColumnHandles(testReadColumns), ImmutableList.of(), - TableToPartitionMapping.empty(), - split.getPath().toString(), + ImmutableMap.of(), + location.toString(), OptionalInt.empty(), - fileSize, + paddedFileSize, Instant.now().toEpochMilli()); - Optional pageSource = HivePageSourceProvider.createHivePageSource( + ConnectorPageSource pageSource = HivePageSourceProvider.createHivePageSource( ImmutableSet.of(sourceFactory), - ImmutableSet.of(), - newEmptyConfiguration(), session, - Location.of(split.getPath().toString()), + location, OptionalInt.empty(), - split.getStart(), - split.getLength(), + 0, fileSize, - splitProperties, + paddedFileSize, + 12345, + new Schema(storageFormat.getSerde(), false, splitProperties), TupleDomain.all(), - columnHandles, TESTING_TYPE_MANAGER, Optional.empty(), Optional.empty(), - false, Optional.empty(), false, NO_ACID_TRANSACTION, - columnMappings); + columnMappings) + .orElseThrow(); - assertTrue(pageSource.isPresent()); - - checkPageSource(pageSource.get(), testReadColumns, getTypes(columnHandles), rowCount); + checkPageSource(pageSource, testReadColumns, rowCount); } - public static boolean hasType(ObjectInspector objectInspector, PrimitiveCategory... types) + private static boolean hasType(Type actualType, Type testType) { - if (objectInspector instanceof PrimitiveObjectInspector primitiveInspector) { - PrimitiveCategory primitiveCategory = primitiveInspector.getPrimitiveCategory(); - for (PrimitiveCategory type : types) { - if (primitiveCategory == type) { - return true; - } - } - return false; - } - if (objectInspector instanceof ListObjectInspector listInspector) { - return hasType(listInspector.getListElementObjectInspector(), types); - } - if (objectInspector instanceof MapObjectInspector mapInspector) { - return hasType(mapInspector.getMapKeyObjectInspector(), types) || - hasType(mapInspector.getMapValueObjectInspector(), types); - } - if (objectInspector instanceof StructObjectInspector structObjectInspector) { - for (StructField field : structObjectInspector.getAllStructFieldRefs()) { - if (hasType(field.getFieldObjectInspector(), types)) { - return true; - } - } - return false; - } - throw new IllegalArgumentException("Unknown object inspector type " + objectInspector); + return actualType.equals(testType) || actualType.getTypeParameters().stream().anyMatch(type -> hasType(type, testType)); } private static boolean withoutNullMapKeyTests(TestColumn testColumn) { - String name = testColumn.getName(); + String name = testColumn.name(); return !name.equals("t_map_null_key") && !name.equals("t_map_null_key_complex_key_value") && !name.equals("t_map_null_key_complex_value"); @@ -1160,7 +915,7 @@ private static boolean withoutNullMapKeyTests(TestColumn testColumn) private static boolean withoutTimestamps(TestColumn testColumn) { - String name = testColumn.getName(); + String name = testColumn.name(); return !name.equals("t_timestamp") && !name.equals("t_map_timestamp") && !name.equals("t_array_timestamp"); @@ -1178,48 +933,36 @@ private static HiveConfig createParquetHiveConfig(boolean useParquetColumnNames) .setUseParquetColumnNames(useParquetColumnNames); } - private void generateProjectedColumns(List childColumns, ImmutableList.Builder testFullColumnsBuilder, ImmutableList.Builder testDereferencedColumnsBuilder) + private void generateProjectedColumns(List testColumns, ImmutableList.Builder testFullColumnsBuilder, ImmutableList.Builder testDereferencedColumnsBuilder) { - for (int i = 0; i < childColumns.size(); i++) { - TestColumn childColumn = childColumns.get(i); - checkState(childColumn.getDereferenceIndices().size() == 0); - ObjectInspector newObjectInspector = getStandardStructObjectInspector( - ImmutableList.of("field0"), - ImmutableList.of(childColumn.getObjectInspector())); - - HiveType hiveType = (HiveType.valueOf(childColumn.getObjectInspector().getTypeName())); - Type trinoType = hiveType.getType(TESTING_TYPE_MANAGER); - - List list = new ArrayList<>(); - list.add(childColumn.getWriteValue()); - - TestColumn newProjectedColumn = new TestColumn( - "new_col" + i, newObjectInspector, - ImmutableList.of("field0"), - ImmutableList.of(0), - childColumn.getObjectInspector(), - childColumn.getWriteValue(), - childColumn.getExpectedValue(), - false); - - TestColumn newFullColumn = new TestColumn("new_col" + i, newObjectInspector, list, rowBlockOf(ImmutableList.of(trinoType), childColumn.getExpectedValue())); - - testFullColumnsBuilder.add(newFullColumn); - testDereferencedColumnsBuilder.add(newProjectedColumn); + // wrapper every test column in a ROW type with one field, and then dereference that field + for (int i = 0; i < testColumns.size(); i++) { + TestColumn testColumn = testColumns.get(i); + verify(!testColumn.dereference()); + + TestColumn baseColumn = new TestColumn( + "new_col" + i, + rowType(field("field0", testColumn.type())), + Collections.singletonList(testColumn.writeValue()), + rowBlockOf(ImmutableList.of(testColumn.type()), testColumn.expectedValue())); + + TestColumn projectedColumn = baseColumn.withDereferenceFirstField(testColumn.writeValue(), testColumn.expectedValue()); + testFullColumnsBuilder.add(baseColumn); + testDereferencedColumnsBuilder.add(projectedColumn); } } private List getRegularColumns(List columns) { return columns.stream() - .filter(column -> !column.isPartitionKey()) + .filter(column -> !column.partitionKey()) .collect(toImmutableList()); } private List getPartitionColumns(List columns) { return columns.stream() - .filter(TestColumn::isPartitionKey) + .filter(TestColumn::partitionKey) .collect(toImmutableList()); } @@ -1236,6 +979,8 @@ private class FileFormatAssertion private HiveFileWriterFactory fileWriterFactory; private long fileSizePadding; + private final TrinoFileSystemFactory fileSystemFactory = new MemoryFileSystemFactory(); + private FileFormatAssertion(String formatName) { this.formatName = requireNonNull(formatName, "formatName is null"); @@ -1259,9 +1004,9 @@ public FileFormatAssertion withSkipGenericWriterTest() return this; } - public FileFormatAssertion withFileWriterFactory(HiveFileWriterFactory fileWriterFactory) + public FileFormatAssertion withFileWriterFactory(Function fileWriterFactoryBuilder) { - this.fileWriterFactory = requireNonNull(fileWriterFactory, "fileWriterFactory is null"); + this.fileWriterFactory = fileWriterFactoryBuilder.apply(fileSystemFactory); return this; } @@ -1302,102 +1047,72 @@ public FileFormatAssertion withFileSizePadding(long fileSizePadding) return this; } - public FileFormatAssertion isReadableByPageSource(HivePageSourceFactory pageSourceFactory) - throws Exception - { - assertRead(Optional.of(pageSourceFactory), Optional.empty(), false); - return this; - } - - public FileFormatAssertion isReadableByRecordCursorPageSource(HiveRecordCursorProvider cursorProvider) - throws Exception - { - assertRead(Optional.empty(), Optional.of(cursorProvider), true); - return this; - } - - public FileFormatAssertion isReadableByRecordCursor(HiveRecordCursorProvider cursorProvider) - throws Exception - { - assertRead(Optional.empty(), Optional.of(cursorProvider), false); - return this; - } - - public FileFormatAssertion isFailingForPageSource(HivePageSourceFactory pageSourceFactory, HiveErrorCode expectedErrorCode, String expectedMessage) + public FileFormatAssertion isReadableByPageSource(Function pageSourceFactoryBuilder) throws Exception { - assertFailure(Optional.of(pageSourceFactory), Optional.empty(), expectedErrorCode, expectedMessage, false); + assertRead(pageSourceFactoryBuilder.apply(fileSystemFactory)); return this; } - public FileFormatAssertion isFailingForRecordCursor(HiveRecordCursorProvider cursorProvider, HiveErrorCode expectedErrorCode, String expectedMessage) - throws Exception + public void isFailingForPageSource(Function pageSourceFactoryBuilder, HiveErrorCode expectedErrorCode, String expectedMessage) { - assertFailure(Optional.empty(), Optional.of(cursorProvider), expectedErrorCode, expectedMessage, false); - return this; + HivePageSourceFactory pageSourceFactory = pageSourceFactoryBuilder.apply(fileSystemFactory); + assertTrinoExceptionThrownBy(() -> assertRead(pageSourceFactory)) + .hasErrorCode(expectedErrorCode) + .hasMessage(expectedMessage); } - private void assertRead(Optional pageSourceFactory, Optional cursorProvider, boolean withRecordPageSource) + private void assertRead(HivePageSourceFactory pageSourceFactory) throws Exception { - assertNotNull(storageFormat, "storageFormat must be specified"); - assertNotNull(writeColumns, "writeColumns must be specified"); - assertNotNull(readColumns, "readColumns must be specified"); - assertNotNull(session, "session must be specified"); - assertTrue(rowsCount >= 0, "rowsCount must be non-negative"); + assertThat(storageFormat) + .describedAs("storageFormat must be specified") + .isNotNull(); + assertThat(writeColumns) + .describedAs("writeColumns must be specified") + .isNotNull(); + assertThat(readColumns) + .describedAs("readColumns must be specified") + .isNotNull(); + assertThat(session) + .describedAs("session must be specified") + .isNotNull(); + assertThat(rowsCount >= 0) + .describedAs("rowsCount must be non-negative") + .isTrue(); String compressionSuffix = compressionCodec.getHiveCompressionKind() .map(CompressionKind::getFileExtension) .orElse(""); - File file = File.createTempFile("trino_test", formatName + compressionSuffix); - file.delete(); + Location location = Location.of("memory:///%s-test%s".formatted(formatName, compressionSuffix)); + TrinoFileSystem fileSystem = fileSystemFactory.create(ConnectorIdentity.ofUser("test")); for (boolean testFileWriter : ImmutableList.of(false, true)) { try { - FileSplit split; if (testFileWriter) { if (fileWriterFactory == null) { continue; } - split = createTestFileTrino(file.getAbsolutePath(), storageFormat, compressionCodec, writeColumns, session, rowsCount, fileWriterFactory); + createTestFileTrino(location, storageFormat, compressionCodec, writeColumns, session, rowsCount, fileWriterFactory); } else { if (skipGenericWrite) { continue; } - split = createTestFileHive(file.getAbsolutePath(), storageFormat, compressionCodec, writeColumns, rowsCount); + createTestFileHive(fileSystemFactory, location, storageFormat, compressionCodec, writeColumns, rowsCount); } - long fileSize = split.getLength() + fileSizePadding; - if (pageSourceFactory.isPresent()) { - testPageSourceFactory(pageSourceFactory.get(), split, storageFormat, readColumns, session, fileSize, rowsCount); - } - if (cursorProvider.isPresent()) { - if (withRecordPageSource) { - testRecordPageSource(cursorProvider.get(), split, storageFormat, readColumns, session, fileSize, rowsCount); - } - else { - testCursorProvider(cursorProvider.get(), split, storageFormat, readColumns, session, fileSize, rowsCount); - } - } + long fileSize = fileSystem.newInputFile(location).length(); + testPageSourceFactory(pageSourceFactory, location, storageFormat, readColumns, session, fileSize, fileSize + fileSizePadding, rowsCount); } finally { - //noinspection ResultOfMethodCallIgnored - file.delete(); + try { + fileSystem.deleteFile(location); + } + catch (IOException ignore) { + } } } } - - private void assertFailure( - Optional pageSourceFactory, - Optional cursorProvider, - HiveErrorCode expectedErrorCode, - String expectedMessage, - boolean withRecordPageSource) - { - assertTrinoExceptionThrownBy(() -> assertRead(pageSourceFactory, cursorProvider, withRecordPageSource)) - .hasErrorCode(expectedErrorCode) - .hasMessage(expectedMessage); - } } } diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHiveFileMetastore.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHiveFileMetastore.java deleted file mode 100644 index 67c763642c22..000000000000 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHiveFileMetastore.java +++ /dev/null @@ -1,88 +0,0 @@ -/* - * 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.hive; - -import io.trino.plugin.hive.metastore.HiveMetastore; -import io.trino.plugin.hive.metastore.file.FileHiveMetastore; -import io.trino.plugin.hive.metastore.file.FileHiveMetastoreConfig; -import org.testng.SkipException; -import org.testng.annotations.Test; - -import java.io.File; - -import static io.trino.plugin.hive.HiveTestUtils.HDFS_ENVIRONMENT; - -// staging directory is shared mutable state -@Test(singleThreaded = true) -public class TestHiveFileMetastore - extends AbstractTestHiveLocal -{ - @Override - protected HiveMetastore createMetastore(File tempDir) - { - File baseDir = new File(tempDir, "metastore"); - return new FileHiveMetastore( - new NodeVersion("test_version"), - HDFS_ENVIRONMENT, - true, - new FileHiveMetastoreConfig() - .setCatalogDirectory(baseDir.toURI().toString()) - .setMetastoreUser("test")); - } - - @Test - public void forceTestNgToRespectSingleThreaded() - { - // TODO: Remove after updating TestNG to 7.4.0+ (https://github.com/trinodb/trino/issues/8571) - // TestNG doesn't enforce @Test(singleThreaded = true) when tests are defined in base class. According to - // https://github.com/cbeust/testng/issues/2361#issuecomment-688393166 a workaround it to add a dummy test to the leaf test class. - } - - @Override - public void testMismatchSchemaTable() - { - // FileHiveMetastore only supports replaceTable() for views - } - - @Override - public void testPartitionSchemaMismatch() - { - // test expects an exception to be thrown - throw new SkipException("FileHiveMetastore only supports replaceTable() for views"); - } - - @Override - public void testBucketedTableEvolution() - { - // FileHiveMetastore only supports replaceTable() for views - } - - @Override - public void testBucketedTableEvolutionWithDifferentReadBucketCount() - { - // FileHiveMetastore has various incompatibilities - } - - @Override - public void testTransactionDeleteInsert() - { - // FileHiveMetastore has various incompatibilities - } - - @Override - public void testInsertOverwriteUnpartitioned() - { - // FileHiveMetastore has various incompatibilities - } -} diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHiveInMemoryMetastore.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHiveInMemoryMetastore.java deleted file mode 100644 index 7de52af417b9..000000000000 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHiveInMemoryMetastore.java +++ /dev/null @@ -1,90 +0,0 @@ -/* - * 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.hive; - -import io.trino.plugin.hive.metastore.HiveMetastore; -import io.trino.plugin.hive.metastore.thrift.BridgingHiveMetastore; -import io.trino.plugin.hive.metastore.thrift.InMemoryThriftMetastore; -import io.trino.plugin.hive.metastore.thrift.ThriftMetastoreConfig; -import org.testng.SkipException; -import org.testng.annotations.Test; - -import java.io.File; - -import static org.assertj.core.api.Assertions.assertThatThrownBy; - -// staging directory is shared mutable state -@Test(singleThreaded = true) -public class TestHiveInMemoryMetastore - extends AbstractTestHiveLocal -{ - @Override - protected HiveMetastore createMetastore(File tempDir) - { - File baseDir = new File(tempDir, "metastore"); - ThriftMetastoreConfig metastoreConfig = new ThriftMetastoreConfig(); - InMemoryThriftMetastore hiveMetastore = new InMemoryThriftMetastore(baseDir, metastoreConfig); - return new BridgingHiveMetastore(hiveMetastore); - } - - @Test - public void forceTestNgToRespectSingleThreaded() - { - // TODO: Remove after updating TestNG to 7.4.0+ (https://github.com/trinodb/trino/issues/8571) - // TestNG doesn't enforce @Test(singleThreaded = true) when tests are defined in base class. According to - // https://github.com/cbeust/testng/issues/2361#issuecomment-688393166 a workaround it to add a dummy test to the leaf test class. - } - - @Override - public void testMetadataDelete() - { - // InMemoryHiveMetastore ignores "removeData" flag in dropPartition - } - - @Override - public void testTransactionDeleteInsert() - { - // InMemoryHiveMetastore does not check whether partition exist in createPartition and dropPartition - } - - @Override - public void testHideDeltaLakeTables() - { - throw new SkipException("not supported"); - } - - @Override - public void testDisallowQueryingOfIcebergTables() - { - throw new SkipException("not supported"); - } - - @Override - public void testDataColumnProperties() - { - // Column properties are currently not supported in ThriftHiveMetastore - assertThatThrownBy(super::testDataColumnProperties) - .isInstanceOf(IllegalArgumentException.class) - .hasMessage("Persisting column properties is not supported: Column{name=id, type=bigint}"); - } - - @Override - public void testPartitionColumnProperties() - { - // Column properties are currently not supported in ThriftHiveMetastore - assertThatThrownBy(super::testPartitionColumnProperties) - .isInstanceOf(IllegalArgumentException.class) - .hasMessage("Persisting column properties is not supported: Column{name=part_key, type=varchar(256)}"); - } -} diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHiveLocationService.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHiveLocationService.java index 1fdc4355b738..099cb54f586d 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHiveLocationService.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHiveLocationService.java @@ -13,14 +13,12 @@ */ package io.trino.plugin.hive; -import com.google.common.collect.ImmutableList; import io.trino.filesystem.Location; -import io.trino.hdfs.HdfsEnvironment; import io.trino.plugin.hive.LocationService.WriteInfo; -import io.trino.plugin.hive.TestBackgroundHiveSplitLoader.TestingHdfsEnvironment; import io.trino.spi.TrinoException; import org.testng.annotations.Test; +import static io.trino.plugin.hive.HiveTestUtils.HDFS_FILE_SYSTEM_FACTORY; import static io.trino.plugin.hive.LocationHandle.WriteMode.DIRECT_TO_TARGET_EXISTING_DIRECTORY; import static io.trino.plugin.hive.LocationHandle.WriteMode.DIRECT_TO_TARGET_NEW_DIRECTORY; import static io.trino.plugin.hive.LocationHandle.WriteMode.STAGE_AND_MOVE_TO_TARGET_DIRECTORY; @@ -80,8 +78,7 @@ public static class Assertion public Assertion(LocationHandle locationHandle, boolean overwrite) { - HdfsEnvironment hdfsEnvironment = new TestingHdfsEnvironment(ImmutableList.of()); - LocationService service = new HiveLocationService(hdfsEnvironment, new HiveConfig()); + LocationService service = new HiveLocationService(HDFS_FILE_SYSTEM_FACTORY, new HiveConfig()); this.actual = service.getTableWriteInfo(locationHandle, overwrite); } diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHivePageSink.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHivePageSink.java index 8e633e11a727..51458ce1af49 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHivePageSink.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHivePageSink.java @@ -18,7 +18,12 @@ import com.google.common.collect.ImmutableMap; import io.airlift.json.JsonCodec; import io.airlift.slice.Slices; +import io.trino.filesystem.FileEntry; +import io.trino.filesystem.FileIterator; import io.trino.filesystem.Location; +import io.trino.filesystem.TrinoFileSystemFactory; +import io.trino.filesystem.memory.MemoryFileSystemFactory; +import io.trino.hive.thrift.metastore.hive_metastoreConstants; import io.trino.operator.GroupByHashPageIndexerFactory; import io.trino.plugin.hive.metastore.HiveMetastore; import io.trino.plugin.hive.metastore.HiveMetastoreFactory; @@ -33,11 +38,11 @@ import io.trino.spi.connector.ConnectorTableHandle; import io.trino.spi.connector.DynamicFilter; import io.trino.spi.connector.SchemaTableName; +import io.trino.spi.security.ConnectorIdentity; import io.trino.spi.type.Type; import io.trino.spi.type.TypeOperators; import io.trino.sql.gen.JoinCompiler; import io.trino.testing.MaterializedResult; -import io.trino.testing.TestingNodeManager; import io.trino.tpch.LineItem; import io.trino.tpch.LineItemColumn; import io.trino.tpch.LineItemGenerator; @@ -47,32 +52,28 @@ import org.testng.annotations.Test; import java.io.File; -import java.nio.file.Files; +import java.io.IOException; import java.util.ArrayList; import java.util.List; +import java.util.Map; import java.util.Optional; import java.util.OptionalInt; -import java.util.Properties; +import java.util.function.Function; import java.util.stream.Stream; import static com.google.common.collect.ImmutableList.toImmutableList; -import static com.google.common.collect.Iterables.getOnlyElement; -import static com.google.common.io.MoreFiles.deleteRecursively; -import static com.google.common.io.RecursiveDeleteOption.ALLOW_INSECURE; import static io.airlift.concurrent.MoreFutures.getFutureValue; -import static io.airlift.testing.Assertions.assertGreaterThan; import static io.trino.plugin.hive.HiveColumnHandle.ColumnType.REGULAR; import static io.trino.plugin.hive.HiveColumnHandle.createBaseColumn; import static io.trino.plugin.hive.HiveCompressionOption.LZ4; import static io.trino.plugin.hive.HiveCompressionOption.NONE; -import static io.trino.plugin.hive.HiveTestUtils.HDFS_ENVIRONMENT; +import static io.trino.plugin.hive.HiveStorageFormat.AVRO; +import static io.trino.plugin.hive.HiveStorageFormat.PARQUET; import static io.trino.plugin.hive.HiveTestUtils.HDFS_FILE_SYSTEM_FACTORY; import static io.trino.plugin.hive.HiveTestUtils.PAGE_SORTER; import static io.trino.plugin.hive.HiveTestUtils.getDefaultHiveFileWriterFactories; import static io.trino.plugin.hive.HiveTestUtils.getDefaultHivePageSourceFactories; -import static io.trino.plugin.hive.HiveTestUtils.getDefaultHiveRecordCursorProviders; import static io.trino.plugin.hive.HiveTestUtils.getHiveSession; -import static io.trino.plugin.hive.HiveTestUtils.getHiveSessionProperties; import static io.trino.plugin.hive.HiveType.HIVE_DATE; import static io.trino.plugin.hive.HiveType.HIVE_DOUBLE; import static io.trino.plugin.hive.HiveType.HIVE_INT; @@ -81,22 +82,21 @@ import static io.trino.plugin.hive.LocationHandle.WriteMode.DIRECT_TO_TARGET_NEW_DIRECTORY; import static io.trino.plugin.hive.acid.AcidTransaction.NO_ACID_TRANSACTION; import static io.trino.plugin.hive.metastore.file.TestingFileHiveMetastore.createTestingFileHiveMetastore; +import static io.trino.plugin.hive.util.SerdeConstants.LIST_COLUMNS; +import static io.trino.plugin.hive.util.SerdeConstants.LIST_COLUMN_TYPES; import static io.trino.spi.type.BigintType.BIGINT; import static io.trino.spi.type.DateType.DATE; import static io.trino.spi.type.DoubleType.DOUBLE; import static io.trino.spi.type.IntegerType.INTEGER; +import static io.trino.spi.type.VarcharType.VARCHAR; import static io.trino.spi.type.VarcharType.createUnboundedVarcharType; import static io.trino.testing.TestingPageSinkId.TESTING_PAGE_SINK_ID; import static io.trino.type.InternalTypeManager.TESTING_TYPE_MANAGER; import static java.lang.Math.round; import static java.lang.String.format; import static java.util.stream.Collectors.toList; -import static org.apache.hadoop.hive.metastore.api.hive_metastoreConstants.FILE_INPUT_FORMAT; -import static org.apache.hadoop.hive.serde.serdeConstants.SERIALIZATION_LIB; import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatThrownBy; -import static org.testng.Assert.assertEquals; -import static org.testng.Assert.assertTrue; public class TestHivePageSink { @@ -110,77 +110,113 @@ public void testAllFormats() { HiveConfig config = new HiveConfig(); SortingFileWriterConfig sortingFileWriterConfig = new SortingFileWriterConfig(); - File tempDir = Files.createTempDirectory(null).toFile(); - try { - HiveMetastore metastore = createTestingFileHiveMetastore(new File(tempDir, "metastore")); - for (HiveStorageFormat format : HiveStorageFormat.values()) { - if (format == HiveStorageFormat.CSV) { - // CSV supports only unbounded VARCHAR type, which is not provided by lineitem + + TrinoFileSystemFactory fileSystemFactory = new MemoryFileSystemFactory(); + HiveMetastore metastore = createTestingFileHiveMetastore(fileSystemFactory, Location.of("memory:///metastore")); + for (HiveStorageFormat format : HiveStorageFormat.values()) { + if (format == HiveStorageFormat.CSV) { + // CSV supports only the unbounded VARCHAR type, which is not provided by lineitem + continue; + } + if (format == HiveStorageFormat.REGEX) { + // REGEX format is readonly + continue; + } + config.setHiveStorageFormat(format); + config.setHiveCompressionCodec(NONE); + long uncompressedLength = writeTestFile(fileSystemFactory, config, sortingFileWriterConfig, metastore, makeFileName(config)); + assertThat(uncompressedLength).isGreaterThan(0L); + + for (HiveCompressionOption codec : HiveCompressionOption.values()) { + if (codec == NONE) { continue; } - if (format == HiveStorageFormat.REGEX) { - // REGEX format is readonly + config.setHiveCompressionCodec(codec); + + // TODO (https://github.com/trinodb/trino/issues/9142) LZ4 is not supported with native Parquet writer + if (!isSupportedCodec(format, codec)) { + assertThatThrownBy(() -> writeTestFile(fileSystemFactory, config, sortingFileWriterConfig, metastore, makeFileName(config))) + .hasMessage("Compression codec " + codec + " not supported for " + format.humanName()); continue; } - config.setHiveStorageFormat(format); - config.setHiveCompressionCodec(NONE); - long uncompressedLength = writeTestFile(config, sortingFileWriterConfig, metastore, makeFileName(tempDir, config)); - assertGreaterThan(uncompressedLength, 0L); - - for (HiveCompressionOption codec : HiveCompressionOption.values()) { - if (codec == NONE) { - continue; - } - if ((format == HiveStorageFormat.PARQUET) && (codec == LZ4)) { - // TODO (https://github.com/trinodb/trino/issues/9142) LZ4 is not supported with native Parquet writer - continue; - } - config.setHiveCompressionCodec(codec); - if (!isSupportedCodec(format, codec)) { - assertThatThrownBy(() -> writeTestFile(config, sortingFileWriterConfig, metastore, makeFileName(tempDir, config))) - .hasMessage("Compression codec " + codec + " not supported for " + format); - continue; - } - - long length = writeTestFile(config, sortingFileWriterConfig, metastore, makeFileName(tempDir, config)); - assertTrue(uncompressedLength > length, format("%s with %s compressed to %s which is not less than %s", format, codec, length, uncompressedLength)); - } + long length = writeTestFile(fileSystemFactory, config, sortingFileWriterConfig, metastore, makeFileName(config)); + assertThat(uncompressedLength > length) + .describedAs(format("%s with %s compressed to %s which is not less than %s", format, codec, length, uncompressedLength)) + .isTrue(); } } - finally { - deleteRecursively(tempDir.toPath(), ALLOW_INSECURE); - } } private boolean isSupportedCodec(HiveStorageFormat storageFormat, HiveCompressionOption compressionOption) { - if (storageFormat == HiveStorageFormat.AVRO && compressionOption == LZ4) { + if ((storageFormat == AVRO || storageFormat == PARQUET) && compressionOption == LZ4) { return false; } return true; } + private static Location makeFileName(HiveConfig config) + { + return Location.of("memory:///" + config.getHiveStorageFormat().name() + "." + config.getHiveCompressionCodec().name()); + } + private static String makeFileName(File tempDir, HiveConfig config) { return tempDir.getAbsolutePath() + "/" + config.getHiveStorageFormat().name() + "." + config.getHiveCompressionCodec().name(); } - private static long writeTestFile(HiveConfig config, SortingFileWriterConfig sortingFileWriterConfig, HiveMetastore metastore, String outputPath) + private static long writeTestFile(TrinoFileSystemFactory fileSystemFactory, HiveConfig config, SortingFileWriterConfig sortingFileWriterConfig, HiveMetastore metastore, Location location) + throws IOException { HiveTransactionHandle transaction = new HiveTransactionHandle(false); HiveWriterStats stats = new HiveWriterStats(); - ConnectorPageSink pageSink = createPageSink(transaction, config, sortingFileWriterConfig, metastore, Location.of("file:///" + outputPath), stats); + ConnectorPageSink pageSink = createPageSink(fileSystemFactory, transaction, config, sortingFileWriterConfig, metastore, location, stats, getColumnHandles()); List columns = getTestColumns(); List columnTypes = columns.stream() .map(LineItemColumn::getType) - .map(TestHivePageSink::getHiveType) - .map(hiveType -> hiveType.getType(TESTING_TYPE_MANAGER)) + .map(TestHivePageSink::getType) + .map(hiveType -> TESTING_TYPE_MANAGER.getType(hiveType.getTypeSignature())) .collect(toList()); + Page page = createPage(lineItem -> true); + pageSink.appendPage(page); + getFutureValue(pageSink.finish()); + + FileIterator fileIterator = fileSystemFactory.create(ConnectorIdentity.ofUser("test")).listFiles(location); + FileEntry fileEntry = fileIterator.next(); + assertThat(fileIterator.hasNext()).isFalse(); + + List pages = new ArrayList<>(); + try (ConnectorPageSource pageSource = createPageSource(fileSystemFactory, transaction, config, fileEntry.location())) { + while (!pageSource.isFinished()) { + Page nextPage = pageSource.getNextPage(); + if (nextPage != null) { + pages.add(nextPage.getLoadedPage()); + } + } + } + + MaterializedResult expectedResults = toMaterializedResult(getHiveSession(config), columnTypes, ImmutableList.of(page)); + MaterializedResult results = toMaterializedResult(getHiveSession(config), columnTypes, pages); + assertThat(results).containsExactlyElementsOf(expectedResults); + assertThat(round(stats.getInputPageSizeInBytes().getAllTime().getMax())).isEqualTo(page.getRetainedSizeInBytes()); + return fileEntry.length(); + } + private static Page createPage(Function filter) + { + List columns = getTestColumns(); + List columnTypes = columns.stream() + .map(LineItemColumn::getType) + .map(TestHivePageSink::getType) + .map(hiveType -> TESTING_TYPE_MANAGER.getType(hiveType.getTypeSignature())) + .collect(toList()); PageBuilder pageBuilder = new PageBuilder(columnTypes); int rows = 0; for (LineItem lineItem : new LineItemGenerator(0.01, 1, 1)) { + if (!filter.apply(lineItem)) { + continue; + } rows++; if (rows >= NUM_ROWS) { break; @@ -210,29 +246,7 @@ private static long writeTestFile(HiveConfig config, SortingFileWriterConfig sor } } } - Page page = pageBuilder.build(); - pageSink.appendPage(page); - getFutureValue(pageSink.finish()); - - File outputDir = new File(outputPath); - List files = ImmutableList.copyOf(outputDir.listFiles((dir, name) -> !name.endsWith(".crc"))); - File outputFile = getOnlyElement(files); - long length = outputFile.length(); - - ConnectorPageSource pageSource = createPageSource(transaction, config, outputFile); - - List pages = new ArrayList<>(); - while (!pageSource.isFinished()) { - Page nextPage = pageSource.getNextPage(); - if (nextPage != null) { - pages.add(nextPage.getLoadedPage()); - } - } - MaterializedResult expectedResults = toMaterializedResult(getHiveSession(config), columnTypes, ImmutableList.of(page)); - MaterializedResult results = toMaterializedResult(getHiveSession(config), columnTypes, pages); - assertThat(results).containsExactlyElementsOf(expectedResults); - assertEquals(round(stats.getInputPageSizeInBytes().getAllTime().getMax()), page.getRetainedSizeInBytes()); - return length; + return pageBuilder.build(); } public static MaterializedResult toMaterializedResult(ConnectorSession session, List types, List pages) @@ -245,50 +259,56 @@ public static MaterializedResult toMaterializedResult(ConnectorSession session, return resultBuilder.build(); } - private static ConnectorPageSource createPageSource(HiveTransactionHandle transaction, HiveConfig config, File outputFile) + private static ConnectorPageSource createPageSource(TrinoFileSystemFactory fileSystemFactory, HiveTransactionHandle transaction, HiveConfig config, Location location) + throws IOException { - Properties splitProperties = new Properties(); - splitProperties.setProperty(FILE_INPUT_FORMAT, config.getHiveStorageFormat().getInputFormat()); - splitProperties.setProperty(SERIALIZATION_LIB, config.getHiveStorageFormat().getSerde()); - splitProperties.setProperty("columns", Joiner.on(',').join(getColumnHandles().stream().map(HiveColumnHandle::getName).collect(toImmutableList()))); - splitProperties.setProperty("columns.types", Joiner.on(',').join(getColumnHandles().stream().map(HiveColumnHandle::getHiveType).map(hiveType -> hiveType.getHiveTypeName().toString()).collect(toImmutableList()))); + long length = fileSystemFactory.create(ConnectorIdentity.ofUser("test")).newInputFile(location).length(); + Map splitProperties = ImmutableMap.builder() + .put(hive_metastoreConstants.FILE_INPUT_FORMAT, config.getHiveStorageFormat().getInputFormat()) + .put(LIST_COLUMNS, Joiner.on(',').join(getColumnHandles().stream().map(HiveColumnHandle::getName).collect(toImmutableList()))) + .put(LIST_COLUMN_TYPES, Joiner.on(',').join(getColumnHandles().stream().map(HiveColumnHandle::getHiveType).map(hiveType -> hiveType.getHiveTypeName().toString()).collect(toImmutableList()))) + .buildOrThrow(); HiveSplit split = new HiveSplit( "", - "file:///" + outputFile.getAbsolutePath(), + location.toString(), 0, - outputFile.length(), - outputFile.length(), - outputFile.lastModified(), - splitProperties, + length, + length, + 0, + new Schema(config.getHiveStorageFormat().getSerde(), false, splitProperties), ImmutableList.of(), ImmutableList.of(), OptionalInt.empty(), OptionalInt.empty(), false, - TableToPartitionMapping.empty(), + ImmutableMap.of(), Optional.empty(), Optional.empty(), - false, Optional.empty(), SplitWeight.standard()); ConnectorTableHandle table = new HiveTableHandle(SCHEMA_NAME, TABLE_NAME, ImmutableMap.of(), ImmutableList.of(), ImmutableList.of(), Optional.empty()); HivePageSourceProvider provider = new HivePageSourceProvider( TESTING_TYPE_MANAGER, - HDFS_ENVIRONMENT, config, - getDefaultHivePageSourceFactories(HDFS_ENVIRONMENT, config), - getDefaultHiveRecordCursorProviders(config, HDFS_ENVIRONMENT), - new GenericHiveRecordCursorProvider(HDFS_ENVIRONMENT, config)); + getDefaultHivePageSourceFactories(fileSystemFactory, config)); return provider.createPageSource(transaction, getHiveSession(config), split, table, ImmutableList.copyOf(getColumnHandles()), DynamicFilter.EMPTY); } - private static ConnectorPageSink createPageSink(HiveTransactionHandle transaction, HiveConfig config, SortingFileWriterConfig sortingFileWriterConfig, HiveMetastore metastore, Location outputPath, HiveWriterStats stats) + private static ConnectorPageSink createPageSink( + TrinoFileSystemFactory fileSystemFactory, + HiveTransactionHandle transaction, + HiveConfig config, + SortingFileWriterConfig sortingFileWriterConfig, + HiveMetastore metastore, + Location location, + HiveWriterStats stats, + List columnHandles) { - LocationHandle locationHandle = new LocationHandle(outputPath, outputPath, DIRECT_TO_TARGET_NEW_DIRECTORY); + LocationHandle locationHandle = new LocationHandle(location, location, DIRECT_TO_TARGET_NEW_DIRECTORY); HiveOutputTableHandle handle = new HiveOutputTableHandle( SCHEMA_NAME, TABLE_NAME, - getColumnHandles(), + columnHandles, new HivePageSinkMetadata(new SchemaTableName(SCHEMA_NAME, TABLE_NAME), metastore.getTable(SCHEMA_NAME, TABLE_NAME), ImmutableMap.of()), locationHandle, config.getHiveStorageFormat(), @@ -304,20 +324,16 @@ private static ConnectorPageSink createPageSink(HiveTransactionHandle transactio TypeOperators typeOperators = new TypeOperators(); BlockTypeOperators blockTypeOperators = new BlockTypeOperators(typeOperators); HivePageSinkProvider provider = new HivePageSinkProvider( - getDefaultHiveFileWriterFactories(config, HDFS_ENVIRONMENT), + getDefaultHiveFileWriterFactories(config, fileSystemFactory), HDFS_FILE_SYSTEM_FACTORY, - HDFS_ENVIRONMENT, PAGE_SORTER, HiveMetastoreFactory.ofInstance(metastore), - new GroupByHashPageIndexerFactory(new JoinCompiler(typeOperators), blockTypeOperators), + new GroupByHashPageIndexerFactory(new JoinCompiler(new TypeOperators()), blockTypeOperators), TESTING_TYPE_MANAGER, config, sortingFileWriterConfig, - new HiveLocationService(HDFS_ENVIRONMENT, config), + new HiveLocationService(HDFS_FILE_SYSTEM_FACTORY, config), partitionUpdateCodec, - new TestingNodeManager("fake-environment"), - new HiveEventClient(), - getHiveSessionProperties(config), stats); return provider.createPageSink(transaction, getHiveSession(config), handle, TESTING_PAGE_SINK_ID); } @@ -358,4 +374,15 @@ private static HiveType getHiveType(TpchColumnType type) } throw new UnsupportedOperationException(); } + + private static Type getType(TpchColumnType type) + { + return switch (type.getBase()) { + case IDENTIFIER -> BIGINT; + case INTEGER -> INTEGER; + case DATE -> DATE; + case DOUBLE -> DOUBLE; + case VARCHAR -> VARCHAR; + }; + } } diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHiveQueryFailureRecoveryTest.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHiveQueryFailureRecoveryTest.java index 221391cae726..530f3572ea61 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHiveQueryFailureRecoveryTest.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHiveQueryFailureRecoveryTest.java @@ -17,6 +17,7 @@ import io.trino.operator.RetryPolicy; import io.trino.plugin.exchange.filesystem.FileSystemExchangePlugin; import io.trino.plugin.exchange.filesystem.containers.MinioStorage; +import io.trino.plugin.hive.containers.Hive3MinioDataLake; import io.trino.plugin.hive.containers.HiveMinioDataLake; import io.trino.plugin.hive.s3.S3HiveQueryRunner; import io.trino.testing.QueryRunner; @@ -48,7 +49,7 @@ protected QueryRunner createQueryRunner( throws Exception { String bucketName = "test-hive-insert-overwrite-" + randomNameSuffix(); // randomizing bucket name to ensure cached TrinoS3FileSystem objects are not reused - this.hiveMinioDataLake = new HiveMinioDataLake(bucketName); + this.hiveMinioDataLake = new Hive3MinioDataLake(bucketName); hiveMinioDataLake.start(); this.minioStorage = new MinioStorage("test-exchange-spooling-" + randomNameSuffix()); diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHiveS3AndGlueMetastoreTest.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHiveS3AndGlueMetastoreTest.java index 5f22419b9204..2163474b4b2a 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHiveS3AndGlueMetastoreTest.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHiveS3AndGlueMetastoreTest.java @@ -15,13 +15,12 @@ import com.google.common.collect.ImmutableMap; import io.trino.Session; +import io.trino.plugin.hive.metastore.glue.GlueHiveMetastore; import io.trino.spi.security.Identity; import io.trino.spi.security.SelectedRole; -import io.trino.testing.DistributedQueryRunner; import io.trino.testing.QueryRunner; import org.testng.annotations.Test; -import java.nio.file.Path; import java.util.HashSet; import java.util.Optional; import java.util.Set; @@ -29,7 +28,7 @@ import static io.trino.plugin.hive.BaseS3AndGlueMetastoreTest.LocationPattern.DOUBLE_SLASH; import static io.trino.plugin.hive.BaseS3AndGlueMetastoreTest.LocationPattern.TRIPLE_SLASH; import static io.trino.plugin.hive.BaseS3AndGlueMetastoreTest.LocationPattern.TWO_TRAILING_SLASHES; -import static io.trino.plugin.hive.metastore.glue.GlueHiveMetastore.createTestingGlueHiveMetastore; +import static io.trino.plugin.hive.TestingHiveUtils.getConnectorService; import static io.trino.spi.security.SelectedRole.Type.ROLE; import static io.trino.testing.TestingNames.randomNameSuffix; import static io.trino.testing.TestingSession.testSessionBuilder; @@ -49,16 +48,25 @@ public TestHiveS3AndGlueMetastoreTest() protected QueryRunner createQueryRunner() throws Exception { - metastore = createTestingGlueHiveMetastore(Path.of(schemaPath())); - Session session = createSession(Optional.of(new SelectedRole(ROLE, Optional.of("admin")))); - DistributedQueryRunner queryRunner = HiveQueryRunner.builder(session) + QueryRunner queryRunner = HiveQueryRunner.builder(session) + .addExtraProperty("sql.path", "hive.functions") + .addExtraProperty("sql.default-function-catalog", "hive") + .addExtraProperty("sql.default-function-schema", "functions") .setCreateTpchSchemas(false) + .addHiveProperty("hive.metastore", "glue") + .addHiveProperty("hive.metastore.glue.default-warehouse-dir", schemaPath()) .addHiveProperty("hive.security", "allow-all") .addHiveProperty("hive.non-managed-table-writes-enabled", "true") - .setMetastore(runner -> metastore) + .addHiveProperty("hive.partition-projection-enabled", "true") + .addHiveProperty("fs.hadoop.enabled", "false") + .addHiveProperty("fs.native-s3.enabled", "true") .build(); queryRunner.execute("CREATE SCHEMA " + schemaName + " WITH (location = '" + schemaPath() + "')"); + queryRunner.execute("CREATE SCHEMA IF NOT EXISTS functions"); + + metastore = getConnectorService(queryRunner, GlueHiveMetastore.class); + return queryRunner; } diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHiveSplit.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHiveSplit.java index d90fd803da74..c63bbcd36b6f 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHiveSplit.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHiveSplit.java @@ -28,9 +28,9 @@ import org.testng.annotations.Test; import java.time.Instant; +import java.util.Map; import java.util.Optional; import java.util.OptionalInt; -import java.util.Properties; import static io.trino.plugin.hive.HiveColumnHandle.createBaseColumn; import static io.trino.plugin.hive.HiveType.HIVE_LONG; @@ -47,9 +47,10 @@ public void testJsonRoundTrip() objectMapperProvider.setJsonDeserializers(ImmutableMap.of(Type.class, new TypeDeserializer(new TestingTypeManager()))); JsonCodec codec = new JsonCodecFactory(objectMapperProvider).jsonCodec(HiveSplit.class); - Properties schema = new Properties(); - schema.setProperty("foo", "bar"); - schema.setProperty("bar", "baz"); + Map schema = ImmutableMap.builder() + .put("foo", "bar") + .put("bar", "baz") + .buildOrThrow(); ImmutableList partitionKeys = ImmutableList.of(new HivePartitionKey("a", "apple"), new HivePartitionKey("b", "42")); ImmutableList addresses = ImmutableList.of(HostAddress.fromParts("127.0.0.1", 44), HostAddress.fromParts("127.0.0.1", 45)); @@ -66,20 +67,19 @@ public void testJsonRoundTrip() 87, 88, Instant.now().toEpochMilli(), - schema, + new Schema("abc", true, schema), partitionKeys, addresses, OptionalInt.empty(), OptionalInt.empty(), true, - TableToPartitionMapping.mapColumnsByIndex(ImmutableMap.of(1, new HiveTypeName("string"))), + ImmutableMap.of(1, new HiveTypeName("string")), Optional.of(new HiveSplit.BucketConversion( BUCKETING_V1, 32, 16, ImmutableList.of(createBaseColumn("col", 5, HIVE_LONG, BIGINT, ColumnType.REGULAR, Optional.of("comment"))))), Optional.empty(), - false, Optional.of(acidInfo), SplitWeight.fromProportion(2.0)); // some non-standard value @@ -93,12 +93,8 @@ public void testJsonRoundTrip() assertEquals(actual.getEstimatedFileSize(), expected.getEstimatedFileSize()); assertEquals(actual.getSchema(), expected.getSchema()); assertEquals(actual.getPartitionKeys(), expected.getPartitionKeys()); - assertEquals(actual.getAddresses(), expected.getAddresses()); - assertEquals(actual.getTableToPartitionMapping().getPartitionColumnCoercions(), expected.getTableToPartitionMapping().getPartitionColumnCoercions()); - assertEquals(actual.getTableToPartitionMapping().getTableToPartitionColumns(), expected.getTableToPartitionMapping().getTableToPartitionColumns()); assertEquals(actual.getBucketConversion(), expected.getBucketConversion()); assertEquals(actual.isForceLocalScheduling(), expected.isForceLocalScheduling()); - assertEquals(actual.isS3SelectPushdownEnabled(), expected.isS3SelectPushdownEnabled()); assertEquals(actual.getAcidInfo().get(), expected.getAcidInfo().get()); assertEquals(actual.getSplitWeight(), expected.getSplitWeight()); } diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHiveSplitSource.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHiveSplitSource.java index 22e008077bfe..2e82ed05eb62 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHiveSplitSource.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHiveSplitSource.java @@ -14,9 +14,11 @@ package io.trino.plugin.hive; import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; import com.google.common.util.concurrent.SettableFuture; import io.airlift.stats.CounterStat; import io.airlift.units.DataSize; +import io.trino.filesystem.cache.DefaultCachingHostAddressProvider; import io.trino.spi.connector.ConnectorSplit; import io.trino.spi.connector.ConnectorSplitSource; import org.testng.annotations.Test; @@ -38,6 +40,7 @@ import static io.trino.plugin.hive.HiveTestUtils.SESSION; import static io.trino.testing.assertions.TrinoExceptionAssert.assertTrinoExceptionThrownBy; import static java.lang.Math.toIntExact; +import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatThrownBy; import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertTrue; @@ -58,6 +61,7 @@ public void testOutstandingSplitCount() new TestingHiveSplitLoader(), Executors.newFixedThreadPool(5), new CounterStat(), + new DefaultCachingHostAddressProvider(), false); // add 10 splits @@ -93,6 +97,7 @@ public void testDynamicPartitionPruning() new TestingHiveSplitLoader(), Executors.newFixedThreadPool(5), new CounterStat(), + new DefaultCachingHostAddressProvider(), false); // add two splits, one of the splits is dynamically pruned @@ -120,6 +125,7 @@ public void testEvenlySizedSplitRemainder() new TestingHiveSplitLoader(), Executors.newSingleThreadExecutor(), new CounterStat(), + new DefaultCachingHostAddressProvider(), false); // One byte larger than the initial split max size @@ -148,6 +154,7 @@ public void testFail() new TestingHiveSplitLoader(), Executors.newFixedThreadPool(5), new CounterStat(), + new DefaultCachingHostAddressProvider(), false); // add some splits @@ -199,6 +206,7 @@ public void testReaderWaitsForSplits() new TestingHiveSplitLoader(), Executors.newFixedThreadPool(5), new CounterStat(), + new DefaultCachingHostAddressProvider(), false); SettableFuture splits = SettableFuture.create(); @@ -231,7 +239,7 @@ public void testReaderWaitsForSplits() // wait for thread to get the split ConnectorSplit split = splits.get(800, TimeUnit.MILLISECONDS); - assertEquals(((HiveSplit) split).getSchema().getProperty("id"), "33"); + assertThat(((HiveSplit) split).getSchema().serdeProperties()).containsEntry("id", "33"); } finally { // make sure the thread exits @@ -254,6 +262,7 @@ public void testOutstandingSplitSize() new TestingHiveSplitLoader(), Executors.newFixedThreadPool(5), new CounterStat(), + new DefaultCachingHostAddressProvider(), false); int testSplitSizeInBytes = new TestSplit(0).getEstimatedSizeInBytes(); @@ -325,17 +334,16 @@ private TestSplit(int id, OptionalInt bucketNumber, DataSize fileSize, BooleanSu fileSize.toBytes(), fileSize.toBytes(), Instant.now().toEpochMilli(), - properties("id", String.valueOf(id)), + new Schema("abc", false, ImmutableMap.of("id", String.valueOf(id))), ImmutableList.of(), ImmutableList.of(new InternalHiveBlock(0, fileSize.toBytes(), ImmutableList.of())), bucketNumber, bucketNumber, true, false, - TableToPartitionMapping.empty(), + ImmutableMap.of(), Optional.empty(), Optional.empty(), - false, Optional.empty(), partitionMatchSupplier); } diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHiveTaskFailureRecoveryTest.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHiveTaskFailureRecoveryTest.java index 7d6f119d99ac..7a653049d118 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHiveTaskFailureRecoveryTest.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHiveTaskFailureRecoveryTest.java @@ -17,6 +17,7 @@ import io.trino.operator.RetryPolicy; import io.trino.plugin.exchange.filesystem.FileSystemExchangePlugin; import io.trino.plugin.exchange.filesystem.containers.MinioStorage; +import io.trino.plugin.hive.containers.Hive3MinioDataLake; import io.trino.plugin.hive.containers.HiveMinioDataLake; import io.trino.plugin.hive.s3.S3HiveQueryRunner; import io.trino.testing.QueryRunner; @@ -48,7 +49,7 @@ protected QueryRunner createQueryRunner( throws Exception { String bucketName = "test-hive-insert-overwrite-" + randomNameSuffix(); // randomizing bucket name to ensure cached TrinoS3FileSystem objects are not reused - this.hiveMinioDataLake = new HiveMinioDataLake(bucketName); + this.hiveMinioDataLake = new Hive3MinioDataLake(bucketName); hiveMinioDataLake.start(); this.minioStorage = new MinioStorage("test-exchange-spooling-" + randomNameSuffix()); diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestNodeLocalDynamicSplitPruning.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestNodeLocalDynamicSplitPruning.java index 57327dba0880..ec03502d26d5 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestNodeLocalDynamicSplitPruning.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestNodeLocalDynamicSplitPruning.java @@ -17,6 +17,9 @@ import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; import io.airlift.testing.TempFile; +import io.trino.filesystem.Location; +import io.trino.filesystem.TrinoFileSystemFactory; +import io.trino.filesystem.memory.MemoryFileSystemFactory; import io.trino.metadata.TableHandle; import io.trino.plugin.hive.metastore.Column; import io.trino.plugin.hive.orc.OrcReaderConfig; @@ -30,6 +33,7 @@ import io.trino.spi.connector.EmptyPageSource; import io.trino.spi.predicate.Domain; import io.trino.spi.predicate.TupleDomain; +import io.trino.spi.security.ConnectorIdentity; import io.trino.testing.TestingConnectorSession; import org.testng.annotations.Test; @@ -38,23 +42,18 @@ import java.util.Map; import java.util.Optional; import java.util.OptionalInt; -import java.util.Properties; import java.util.Set; import java.util.concurrent.CompletableFuture; import static io.trino.plugin.hive.HiveColumnHandle.ColumnType.PARTITION_KEY; import static io.trino.plugin.hive.HiveColumnHandle.ColumnType.REGULAR; -import static io.trino.plugin.hive.HiveTestUtils.HDFS_ENVIRONMENT; import static io.trino.plugin.hive.HiveTestUtils.getDefaultHivePageSourceFactories; -import static io.trino.plugin.hive.HiveTestUtils.getDefaultHiveRecordCursorProviders; import static io.trino.plugin.hive.HiveType.HIVE_INT; import static io.trino.plugin.hive.util.HiveBucketing.BucketingVersion.BUCKETING_V1; import static io.trino.spi.type.IntegerType.INTEGER; import static io.trino.testing.TestingHandles.TEST_CATALOG_HANDLE; import static io.trino.type.InternalTypeManager.TESTING_TYPE_MANAGER; import static java.util.concurrent.CompletableFuture.completedFuture; -import static org.apache.hadoop.hive.metastore.api.hive_metastoreConstants.FILE_INPUT_FORMAT; -import static org.apache.hadoop.hive.serde.serdeConstants.SERIALIZATION_LIB; import static org.testng.Assert.assertEquals; public class TestNodeLocalDynamicSplitPruning @@ -115,27 +114,28 @@ public void testDynamicPartitionPruning() } private static ConnectorPageSource createTestingPageSource(HiveTransactionHandle transaction, HiveConfig hiveConfig, File outputFile, DynamicFilter dynamicFilter) + throws IOException { - Properties splitProperties = new Properties(); - splitProperties.setProperty(FILE_INPUT_FORMAT, hiveConfig.getHiveStorageFormat().getInputFormat()); - splitProperties.setProperty(SERIALIZATION_LIB, hiveConfig.getHiveStorageFormat().getSerde()); + Location location = Location.of("memory:///file"); + TrinoFileSystemFactory fileSystemFactory = new MemoryFileSystemFactory(); + fileSystemFactory.create(ConnectorIdentity.ofUser("test")).newOutputFile(location).create().close(); + HiveSplit split = new HiveSplit( "", - "file:///" + outputFile.getAbsolutePath(), + location.toString(), + 0, + 0, + 0, 0, - outputFile.length(), - outputFile.length(), - outputFile.lastModified(), - splitProperties, + new Schema(hiveConfig.getHiveStorageFormat().getSerde(), false, ImmutableMap.of()), ImmutableList.of(new HivePartitionKey(PARTITION_COLUMN.getName(), "42")), ImmutableList.of(), OptionalInt.of(1), OptionalInt.of(1), false, - TableToPartitionMapping.empty(), + ImmutableMap.of(), Optional.empty(), Optional.empty(), - false, Optional.empty(), SplitWeight.standard()); @@ -147,21 +147,20 @@ private static ConnectorPageSource createTestingPageSource(HiveTransactionHandle ImmutableMap.of(), ImmutableList.of(), ImmutableList.of(BUCKET_HIVE_COLUMN_HANDLE), - Optional.of(new HiveBucketHandle( - ImmutableList.of(BUCKET_HIVE_COLUMN_HANDLE), + Optional.of(new HiveTablePartitioning( + true, BUCKETING_V1, 20, - 20, - ImmutableList.of()))), + ImmutableList.of(BUCKET_HIVE_COLUMN_HANDLE), + false, + ImmutableList.of(), + true))), transaction); HivePageSourceProvider provider = new HivePageSourceProvider( TESTING_TYPE_MANAGER, - HDFS_ENVIRONMENT, hiveConfig, - getDefaultHivePageSourceFactories(HDFS_ENVIRONMENT, hiveConfig), - getDefaultHiveRecordCursorProviders(hiveConfig, HDFS_ENVIRONMENT), - new GenericHiveRecordCursorProvider(HDFS_ENVIRONMENT, hiveConfig)); + getDefaultHivePageSourceFactories(fileSystemFactory, hiveConfig)); return provider.createPageSource( transaction, @@ -209,7 +208,6 @@ private static TestingConnectorSession getSession(HiveConfig config) return TestingConnectorSession.builder() .setPropertyMetadata(new HiveSessionProperties( config, - new HiveFormatsConfig(), new OrcReaderConfig(), new OrcWriterConfig(), new ParquetReaderConfig(), diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestOrcPageSourceMemoryTracking.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestOrcPageSourceMemoryTracking.java index d3278c0345af..567d944afc4c 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestOrcPageSourceMemoryTracking.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestOrcPageSourceMemoryTracking.java @@ -14,6 +14,7 @@ package io.trino.plugin.hive; import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; import io.airlift.slice.Slice; import io.airlift.stats.Distribution; @@ -57,7 +58,6 @@ import org.apache.hadoop.hive.ql.exec.FileSinkOperator.RecordWriter; import org.apache.hadoop.hive.ql.io.orc.OrcFile; import org.apache.hadoop.hive.ql.io.orc.OrcFile.WriterOptions; -import org.apache.hadoop.hive.ql.io.orc.OrcInputFormat; import org.apache.hadoop.hive.ql.io.orc.OrcOutputFormat; import org.apache.hadoop.hive.ql.io.orc.OrcSerde; import org.apache.hadoop.hive.ql.io.orc.Writer; @@ -106,6 +106,8 @@ import static io.trino.plugin.hive.HiveTestUtils.HDFS_FILE_SYSTEM_FACTORY; import static io.trino.plugin.hive.HiveTestUtils.SESSION; import static io.trino.plugin.hive.acid.AcidTransaction.NO_ACID_TRANSACTION; +import static io.trino.plugin.hive.util.SerdeConstants.LIST_COLUMNS; +import static io.trino.plugin.hive.util.SerdeConstants.LIST_COLUMN_TYPES; import static io.trino.spi.type.VarcharType.createUnboundedVarcharType; import static io.trino.sql.relational.Expressions.field; import static io.trino.testing.TestingHandles.TEST_CATALOG_HANDLE; @@ -118,9 +120,7 @@ import static java.util.concurrent.Executors.newCachedThreadPool; import static java.util.concurrent.Executors.newScheduledThreadPool; import static java.util.stream.Collectors.toList; -import static org.apache.hadoop.hive.metastore.api.hive_metastoreConstants.FILE_INPUT_FORMAT; import static org.apache.hadoop.hive.ql.io.orc.CompressionKind.ZLIB; -import static org.apache.hadoop.hive.serde.serdeConstants.SERIALIZATION_LIB; import static org.apache.hadoop.hive.serde2.objectinspector.ObjectInspectorFactory.getStandardStructObjectInspector; import static org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorFactory.javaStringObjectInspector; import static org.apache.hadoop.mapreduce.lib.output.FileOutputFormat.COMPRESS_CODEC; @@ -325,7 +325,6 @@ public void testMaxReadBytes(int rowCount) int maxReadBytes = 1_000; HiveSessionProperties hiveSessionProperties = new HiveSessionProperties( new HiveConfig(), - new HiveFormatsConfig(), new OrcReaderConfig() .setMaxBlockSize(DataSize.ofBytes(maxReadBytes)), new OrcWriterConfig(), @@ -483,7 +482,7 @@ public void testScanFilterAndProjectOperator() private class TestPreparer { private final FileSplit fileSplit; - private final Properties schema; + private final Schema schema; private final List columns; private final List types; private final String partitionName; @@ -501,17 +500,13 @@ public TestPreparer(String tempFilePath, List testColumns, int numRo throws Exception { OrcSerde serde = new OrcSerde(); - schema = new Properties(); - schema.setProperty("columns", - testColumns.stream() - .map(TestColumn::getName) - .collect(Collectors.joining(","))); - schema.setProperty("columns.types", - testColumns.stream() - .map(TestColumn::getType) - .collect(Collectors.joining(","))); - schema.setProperty(FILE_INPUT_FORMAT, OrcInputFormat.class.getName()); - schema.setProperty(SERIALIZATION_LIB, serde.getClass().getName()); + schema = new Schema( + serde.getClass().getName(), + false, + ImmutableMap.builder() + .put(LIST_COLUMNS, testColumns.stream().map(TestColumn::getName).collect(Collectors.joining(","))) + .put(LIST_COLUMN_TYPES, testColumns.stream().map(TestColumn::getType).collect(Collectors.joining(","))) + .buildOrThrow()); partitionKeys = testColumns.stream() .filter(TestColumn::isPartitionKey) @@ -560,7 +555,7 @@ public ConnectorPageSource newPageSource(FileFormatDataSourceStats stats, Connec partitionKeys, columns, ImmutableList.of(), - TableToPartitionMapping.empty(), + ImmutableMap.of(), fileSplit.getPath().toString(), OptionalInt.empty(), fileSplit.getLength(), @@ -568,21 +563,18 @@ public ConnectorPageSource newPageSource(FileFormatDataSourceStats stats, Connec ConnectorPageSource connectorPageSource = HivePageSourceProvider.createHivePageSource( ImmutableSet.of(orcPageSourceFactory), - ImmutableSet.of(), - newEmptyConfiguration(), session, Location.of(fileSplit.getPath().toString()), OptionalInt.empty(), fileSplit.getStart(), fileSplit.getLength(), fileSplit.getLength(), + 12345, schema, TupleDomain.all(), - columns, TESTING_TYPE_MANAGER, Optional.empty(), Optional.empty(), - false, Optional.empty(), false, NO_ACID_TRANSACTION, diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestParquetPageSkipping.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestParquetPageSkipping.java index 255a89365d78..8eef2fe90c98 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestParquetPageSkipping.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestParquetPageSkipping.java @@ -14,8 +14,12 @@ package io.trino.plugin.hive; import com.google.common.collect.ImmutableMap; +import io.trino.filesystem.TrinoFileSystemFactory; +import io.trino.spi.security.ConnectorIdentity; import io.trino.testing.QueryRunner; +import static io.trino.plugin.hive.TestingHiveUtils.getConnectorService; + public class TestParquetPageSkipping extends AbstractTestParquetPageSkipping { @@ -23,12 +27,15 @@ public class TestParquetPageSkipping protected QueryRunner createQueryRunner() throws Exception { - return HiveQueryRunner.builder() + QueryRunner queryRunner = HiveQueryRunner.builder() .setHiveProperties( ImmutableMap.of( "parquet.use-column-index", "true", "parquet.max-buffer-size", "1MB", "parquet.optimized-reader.enabled", "false")) .build(); + fileSystem = getConnectorService(queryRunner, TrinoFileSystemFactory.class) + .create(ConnectorIdentity.ofUser("test")); + return queryRunner; } } diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestParquetPageSkippingWithOptimizedReader.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestParquetPageSkippingWithOptimizedReader.java index f68b673d9137..042e0508cd88 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestParquetPageSkippingWithOptimizedReader.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestParquetPageSkippingWithOptimizedReader.java @@ -14,8 +14,12 @@ package io.trino.plugin.hive; import com.google.common.collect.ImmutableMap; +import io.trino.filesystem.TrinoFileSystemFactory; +import io.trino.spi.security.ConnectorIdentity; import io.trino.testing.QueryRunner; +import static io.trino.plugin.hive.TestingHiveUtils.getConnectorService; + public class TestParquetPageSkippingWithOptimizedReader extends AbstractTestParquetPageSkipping { @@ -23,12 +27,15 @@ public class TestParquetPageSkippingWithOptimizedReader protected QueryRunner createQueryRunner() throws Exception { - return HiveQueryRunner.builder() + QueryRunner queryRunner = HiveQueryRunner.builder() .setHiveProperties( ImmutableMap.of( "parquet.use-column-index", "true", "parquet.max-buffer-size", "1MB", "parquet.optimized-reader.enabled", "true")) .build(); + fileSystem = getConnectorService(queryRunner, TrinoFileSystemFactory.class) + .create(ConnectorIdentity.ofUser("test")); + return queryRunner; } } diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestReaderProjectionsAdapter.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestReaderProjectionsAdapter.java index 819ff910d103..6cff487afce6 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestReaderProjectionsAdapter.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestReaderProjectionsAdapter.java @@ -19,6 +19,7 @@ import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; import io.trino.spi.block.ColumnarRow; +import io.trino.spi.block.DictionaryBlock; import io.trino.spi.block.LazyBlock; import io.trino.spi.block.RowBlock; import io.trino.spi.connector.ColumnHandle; @@ -115,8 +116,9 @@ public void testLazyDereferenceProjectionLoading() assertFalse(columnarRowLevel1.getField(1).isLoaded()); Block lazyBlockLevel2 = columnarRowLevel1.getField(0); - assertTrue(lazyBlockLevel2 instanceof LazyBlock); - RowBlock rowBlockLevel2 = ((RowBlock) (((LazyBlock) lazyBlockLevel2).getBlock())); + assertTrue(lazyBlockLevel2 instanceof DictionaryBlock); + assertTrue(((DictionaryBlock) lazyBlockLevel2).getDictionary() instanceof LazyBlock); + RowBlock rowBlockLevel2 = ((RowBlock) ((LazyBlock) ((DictionaryBlock) lazyBlockLevel2).getDictionary()).getBlock()); assertFalse(rowBlockLevel2.isLoaded()); ColumnarRow columnarRowLevel2 = toColumnarRow(rowBlockLevel2); // Assertion for "col.f_row_0.f_bigint_0" and "col.f_row_0.f_bigint_1" @@ -186,6 +188,9 @@ private static Block createRowBlockWithLazyNestedBlocks(List data, RowTy RowData row = (RowData) data.get(position); if (row == null) { isNull[position] = true; + for (int field = 0; field < fieldCount; field++) { + fieldsData.get(field).add(null); + } } else { for (int field = 0; field < fieldCount; field++) { diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestRegexTable.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestRegexTable.java index d7bd4d5a9d9a..d05e846e37d8 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestRegexTable.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestRegexTable.java @@ -14,6 +14,7 @@ package io.trino.plugin.hive; import com.google.common.collect.ImmutableMap; +import io.trino.filesystem.Location; import io.trino.testing.AbstractTestQueryFramework; import io.trino.testing.MaterializedResult; import io.trino.testing.QueryRunner; @@ -21,12 +22,9 @@ import org.testng.annotations.Test; import java.io.IOException; -import java.nio.file.Path; +import java.util.UUID; -import static com.google.common.io.MoreFiles.deleteRecursively; -import static com.google.common.io.RecursiveDeleteOption.ALLOW_INSECURE; import static io.trino.testing.QueryAssertions.assertEqualsIgnoreOrder; -import static java.nio.file.Files.createTempDirectory; public class TestRegexTable extends AbstractTestQueryFramework @@ -44,8 +42,7 @@ protected QueryRunner createQueryRunner() public void testCreateExternalTableWithData() throws IOException { - Path tempDir = createTempDirectory(null); - Path tableLocation = tempDir.resolve("data"); + Location tempDir = Location.of("local:///temp_" + UUID.randomUUID()); // REGEX format is read-only, so create data files using the text file format @Language("SQL") String createTableSql = """ @@ -55,7 +52,7 @@ public void testCreateExternalTableWithData() textfile_field_separator = 'x', external_location = '%s') AS SELECT nationkey, name FROM tpch.tiny.nation - """.formatted(tableLocation.toUri().toASCIIString()); + """.formatted(tempDir); assertUpdate(createTableSql, 25); MaterializedResult expected = computeActual("SELECT nationkey, name FROM tpch.tiny.nation"); @@ -71,7 +68,7 @@ CREATE TABLE test_regex ( format = 'regex', regex = '(\\d+)x(.+)', external_location = '%s') - """.formatted(tableLocation.toUri().toASCIIString()); + """.formatted(tempDir); assertUpdate(createTableSql); actual = computeActual("SELECT nationkey, name FROM test_regex"); @@ -91,7 +88,7 @@ CREATE TABLE test_regex ( regex = '(\\d+)X(.+)', regex_case_insensitive = true, external_location = '%s') - """.formatted(tableLocation.toUri().toASCIIString()); + """.formatted(tempDir); assertUpdate(createTableSql); actual = computeActual("SELECT nationkey, name FROM test_regex"); assertEqualsIgnoreOrder(actual.getMaterializedRows(), expected.getMaterializedRows()); @@ -106,14 +103,13 @@ CREATE TABLE test_regex ( format = 'regex', regex = '(\\d+)X(.+)', external_location = '%s') - """.formatted(tableLocation.toUri().toASCIIString()); + """.formatted(tempDir); assertUpdate(createTableSql); // when the pattern does not match all columns are null assertQueryReturnsEmptyResult("SELECT nationkey, name FROM test_regex WHERE nationkey IS NOT NULL AND name IS NOT NULL"); assertUpdate("DROP TABLE test_regex"); assertUpdate("DROP TABLE test_regex_data"); - deleteRecursively(tempDir, ALLOW_INSECURE); } @Test diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestingHiveConnectorFactory.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestingHiveConnectorFactory.java index 086472941d78..d81a3e5f1424 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestingHiveConnectorFactory.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestingHiveConnectorFactory.java @@ -13,18 +13,22 @@ */ package io.trino.plugin.hive; +import com.google.common.collect.ImmutableMap; import com.google.inject.Module; -import io.opentelemetry.api.OpenTelemetry; -import io.trino.plugin.hive.fs.DirectoryLister; +import io.trino.filesystem.TrinoFileSystemFactory; +import io.trino.filesystem.local.LocalFileSystemFactory; import io.trino.plugin.hive.metastore.HiveMetastore; +import io.trino.plugin.hive.metastore.file.FileHiveMetastoreConfig; import io.trino.spi.connector.Connector; import io.trino.spi.connector.ConnectorContext; import io.trino.spi.connector.ConnectorFactory; +import java.nio.file.Path; import java.util.Map; import java.util.Optional; -import static com.google.inject.util.Modules.EMPTY_MODULE; +import static com.google.inject.multibindings.MapBinder.newMapBinder; +import static io.airlift.configuration.ConfigBinder.configBinder; import static io.trino.plugin.hive.InternalHiveConnectorFactory.createConnector; import static java.util.Objects.requireNonNull; @@ -32,25 +36,24 @@ public class TestingHiveConnectorFactory implements ConnectorFactory { private final Optional metastore; - private final Optional openTelemetry; private final Module module; - private final Optional directoryLister; - public TestingHiveConnectorFactory(HiveMetastore metastore) + public TestingHiveConnectorFactory(Path localFileSystemRootPath) { - this(Optional.of(metastore), Optional.empty(), EMPTY_MODULE, Optional.empty()); + this(localFileSystemRootPath, Optional.empty()); } - public TestingHiveConnectorFactory( - Optional metastore, - Optional openTelemetry, - Module module, - Optional directoryLister) + @Deprecated + public TestingHiveConnectorFactory(Path localFileSystemRootPath, Optional metastore) { this.metastore = requireNonNull(metastore, "metastore is null"); - this.openTelemetry = requireNonNull(openTelemetry, "openTelemetry is null"); - this.module = requireNonNull(module, "module is null"); - this.directoryLister = requireNonNull(directoryLister, "directoryLister is null"); + + boolean ignored = localFileSystemRootPath.toFile().mkdirs(); + this.module = binder -> { + newMapBinder(binder, String.class, TrinoFileSystemFactory.class) + .addBinding("local").toInstance(new LocalFileSystemFactory(localFileSystemRootPath)); + configBinder(binder).bindConfigDefaults(FileHiveMetastoreConfig.class, config -> config.setCatalogDirectory("local:///")); + }; } @Override @@ -62,6 +65,12 @@ public String getName() @Override public Connector create(String catalogName, Map config, ConnectorContext context) { - return createConnector(catalogName, config, context, module, metastore, Optional.empty(), openTelemetry, directoryLister); + ImmutableMap.Builder configBuilder = ImmutableMap.builder() + .putAll(config) + .put("bootstrap.quiet", "true"); + if (metastore.isEmpty() && !config.containsKey("hive.metastore")) { + configBuilder.put("hive.metastore", "file"); + } + return createConnector(catalogName, configBuilder.buildOrThrow(), context, module, metastore, Optional.empty()); } } diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestingHivePlugin.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestingHivePlugin.java index 13975b1995b3..36573d2aa226 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestingHivePlugin.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestingHivePlugin.java @@ -14,47 +14,42 @@ package io.trino.plugin.hive; import com.google.common.collect.ImmutableList; -import com.google.inject.Module; -import io.opentelemetry.api.OpenTelemetry; -import io.trino.plugin.hive.fs.DirectoryLister; import io.trino.plugin.hive.metastore.HiveMetastore; import io.trino.spi.Plugin; import io.trino.spi.connector.ConnectorFactory; +import java.nio.file.Path; import java.util.Optional; -import static com.google.inject.util.Modules.EMPTY_MODULE; import static java.util.Objects.requireNonNull; public class TestingHivePlugin implements Plugin { + private final Path localFileSystemRootPath; private final Optional metastore; - private final Optional openTelemetry; - private final Module module; - private final Optional directoryLister; - public TestingHivePlugin() + public TestingHivePlugin(Path localFileSystemRootPath) { - this(Optional.empty(), Optional.empty(), EMPTY_MODULE, Optional.empty()); + this(localFileSystemRootPath, Optional.empty()); } - public TestingHivePlugin(HiveMetastore metastore) + @Deprecated + public TestingHivePlugin(Path localFileSystemRootPath, HiveMetastore metastore) { - this(Optional.of(metastore), Optional.empty(), EMPTY_MODULE, Optional.empty()); + this(localFileSystemRootPath, Optional.of(metastore)); } - public TestingHivePlugin(Optional metastore, Optional openTelemetry, Module module, Optional directoryLister) + @Deprecated + public TestingHivePlugin(Path localFileSystemRootPath, Optional metastore) { + this.localFileSystemRootPath = requireNonNull(localFileSystemRootPath, "localFileSystemRootPath is null"); this.metastore = requireNonNull(metastore, "metastore is null"); - this.openTelemetry = requireNonNull(openTelemetry, "openTelemetry is null"); - this.module = requireNonNull(module, "module is null"); - this.directoryLister = requireNonNull(directoryLister, "directoryLister is null"); } @Override public Iterable getConnectorFactories() { - return ImmutableList.of(new TestingHiveConnectorFactory(metastore, openTelemetry, module, directoryLister)); + return ImmutableList.of(new TestingHiveConnectorFactory(localFileSystemRootPath, metastore)); } } diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestingHiveUtils.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestingHiveUtils.java new file mode 100644 index 000000000000..33b98a002114 --- /dev/null +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestingHiveUtils.java @@ -0,0 +1,47 @@ +/* + * 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.hive; + +import com.google.inject.Injector; +import com.google.inject.Key; +import io.trino.testing.DistributedQueryRunner; +import io.trino.testing.PlanTester; +import io.trino.testing.QueryRunner; + +import static io.trino.plugin.hive.HiveQueryRunner.HIVE_CATALOG; + +public final class TestingHiveUtils +{ + private TestingHiveUtils() {} + + public static T getConnectorService(PlanTester planTester, Class clazz) + { + return ((HiveConnector) planTester.getConnector(HIVE_CATALOG)).getInjector().getInstance(clazz); + } + + public static T getConnectorService(QueryRunner queryRunner, Class clazz) + { + return getConnectorInjector(queryRunner).getInstance(clazz); + } + + public static T getConnectorService(QueryRunner queryRunner, Key key) + { + return getConnectorInjector(queryRunner).getInstance(key); + } + + private static Injector getConnectorInjector(QueryRunner queryRunner) + { + return ((HiveConnector) ((DistributedQueryRunner) queryRunner).getCoordinator().getConnector(HIVE_CATALOG)).getInjector(); + } +} diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestingThriftHiveMetastoreBuilder.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestingThriftHiveMetastoreBuilder.java index d0bbc9748dd3..c489446e2e96 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestingThriftHiveMetastoreBuilder.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestingThriftHiveMetastoreBuilder.java @@ -14,8 +14,9 @@ package io.trino.plugin.hive; import com.google.common.collect.ImmutableSet; -import com.google.common.net.HostAndPort; import io.airlift.units.Duration; +import io.trino.filesystem.TrinoFileSystemFactory; +import io.trino.filesystem.hdfs.HdfsFileSystemFactory; import io.trino.hdfs.DynamicHdfsConfiguration; import io.trino.hdfs.HdfsConfig; import io.trino.hdfs.HdfsConfigurationInitializer; @@ -36,10 +37,14 @@ import io.trino.plugin.hive.metastore.thrift.TokenAwareMetastoreClientFactory; import io.trino.plugin.hive.metastore.thrift.UgiBasedMetastoreClientFactory; +import java.net.URI; import java.util.Optional; +import java.util.concurrent.ExecutorService; +import java.util.function.Consumer; import static com.google.common.base.Preconditions.checkState; import static io.trino.plugin.base.security.UserNameProvider.SIMPLE_USER_NAME_PROVIDER; +import static io.trino.plugin.hive.HiveTestUtils.HDFS_FILE_SYSTEM_STATS; import static java.util.Objects.requireNonNull; import static java.util.concurrent.Executors.newFixedThreadPool; @@ -59,9 +64,8 @@ public final class TestingThriftHiveMetastoreBuilder new NoHdfsAuthentication()); private TokenAwareMetastoreClientFactory tokenAwareMetastoreClientFactory; - private HiveConfig hiveConfig = new HiveConfig(); private ThriftMetastoreConfig thriftMetastoreConfig = new ThriftMetastoreConfig(); - private HdfsEnvironment hdfsEnvironment = HDFS_ENVIRONMENT; + private TrinoFileSystemFactory fileSystemFactory = new HdfsFileSystemFactory(HDFS_ENVIRONMENT, HDFS_FILE_SYSTEM_STATS); public static TestingThriftHiveMetastoreBuilder testingThriftHiveMetastoreBuilder() { @@ -70,7 +74,7 @@ public static TestingThriftHiveMetastoreBuilder testingThriftHiveMetastoreBuilde private TestingThriftHiveMetastoreBuilder() {} - public TestingThriftHiveMetastoreBuilder metastoreClient(HostAndPort address, Duration timeout) + public TestingThriftHiveMetastoreBuilder metastoreClient(URI address, Duration timeout) { requireNonNull(address, "address is null"); requireNonNull(timeout, "timeout is null"); @@ -79,7 +83,7 @@ public TestingThriftHiveMetastoreBuilder metastoreClient(HostAndPort address, Du return this; } - public TestingThriftHiveMetastoreBuilder metastoreClient(HostAndPort address) + public TestingThriftHiveMetastoreBuilder metastoreClient(URI address) { requireNonNull(address, "address is null"); checkState(tokenAwareMetastoreClientFactory == null, "Metastore client already set"); @@ -95,34 +99,23 @@ public TestingThriftHiveMetastoreBuilder metastoreClient(ThriftMetastoreClient c return this; } - public TestingThriftHiveMetastoreBuilder hiveConfig(HiveConfig hiveConfig) - { - this.hiveConfig = requireNonNull(hiveConfig, "hiveConfig is null"); - return this; - } - public TestingThriftHiveMetastoreBuilder thriftMetastoreConfig(ThriftMetastoreConfig thriftMetastoreConfig) { this.thriftMetastoreConfig = requireNonNull(thriftMetastoreConfig, "thriftMetastoreConfig is null"); return this; } - public TestingThriftHiveMetastoreBuilder hdfsEnvironment(HdfsEnvironment hdfsEnvironment) - { - this.hdfsEnvironment = requireNonNull(hdfsEnvironment, "hdfsEnvironment is null"); - return this; - } - - public ThriftMetastore build() + public ThriftMetastore build(Consumer registerResource) { checkState(tokenAwareMetastoreClientFactory != null, "metastore client not set"); + ExecutorService executorService = newFixedThreadPool(thriftMetastoreConfig.getWriteStatisticsThreads()); + registerResource.accept(executorService::shutdown); ThriftHiveMetastoreFactory metastoreFactory = new ThriftHiveMetastoreFactory( new UgiBasedMetastoreClientFactory(tokenAwareMetastoreClientFactory, SIMPLE_USER_NAME_PROVIDER, thriftMetastoreConfig), new HiveMetastoreConfig().isHideDeltaLakeTables(), - hiveConfig.isTranslateHiveViews(), thriftMetastoreConfig, - hdfsEnvironment, - newFixedThreadPool(thriftMetastoreConfig.getWriteStatisticsThreads())); + fileSystemFactory, + executorService); return metastoreFactory.createMetastore(Optional.empty()); } } diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/benchmark/AbstractFileFormat.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/benchmark/AbstractFileFormat.java index d6afdddd1a54..43ecbc14be76 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/benchmark/AbstractFileFormat.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/benchmark/AbstractFileFormat.java @@ -16,9 +16,9 @@ import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; +import com.google.common.collect.Maps; import io.trino.filesystem.Location; import io.trino.hdfs.HdfsEnvironment; -import io.trino.plugin.hive.GenericHiveRecordCursorProvider; import io.trino.plugin.hive.HiveColumnHandle; import io.trino.plugin.hive.HiveConfig; import io.trino.plugin.hive.HivePageSourceFactory; @@ -31,7 +31,7 @@ import io.trino.plugin.hive.HiveType; import io.trino.plugin.hive.HiveTypeName; import io.trino.plugin.hive.ReaderPageSource; -import io.trino.plugin.hive.TableToPartitionMapping; +import io.trino.plugin.hive.Schema; import io.trino.spi.SplitWeight; import io.trino.spi.connector.ColumnHandle; import io.trino.spi.connector.ConnectorPageSource; @@ -125,13 +125,10 @@ public ConnectorPageSource createGenericReader( { HivePageSourceProvider factory = new HivePageSourceProvider( TESTING_TYPE_MANAGER, - hdfsEnvironment, new HiveConfig(), - getHivePageSourceFactory(hdfsEnvironment).map(ImmutableSet::of).orElse(ImmutableSet.of()), - getHiveRecordCursorProvider(hdfsEnvironment).map(ImmutableSet::of).orElse(ImmutableSet.of()), - new GenericHiveRecordCursorProvider(hdfsEnvironment, new HiveConfig())); + getHivePageSourceFactory(hdfsEnvironment).map(ImmutableSet::of).orElse(ImmutableSet.of())); - Properties schema = createSchema(getFormat(), schemaColumnNames, schemaColumnTypes); + Schema schema = createSchema(getFormat(), schemaColumnNames, schemaColumnTypes); HiveSplit split = new HiveSplit( "", @@ -142,14 +139,12 @@ public ConnectorPageSource createGenericReader( targetFile.lastModified(), schema, ImmutableList.of(), - ImmutableList.of(), OptionalInt.empty(), OptionalInt.empty(), false, - TableToPartitionMapping.empty(), + ImmutableMap.of(), Optional.empty(), Optional.empty(), - false, Optional.empty(), SplitWeight.standard()); @@ -186,7 +181,7 @@ static ConnectorPageSource createPageSource( 0, targetFile.length(), targetFile.length(), - createSchema(format, columnNames, columnTypes), + createSchema(format, columnNames, columnTypes).serdeProperties(), readColumns, TupleDomain.all(), TESTING_TYPE_MANAGER, @@ -209,7 +204,7 @@ static ConnectorPageSource createPageSource( List readColumns = getBaseColumns(columnNames, columnTypes); - Properties schema = createSchema(format, columnNames, columnTypes); + Schema schema = createSchema(format, columnNames, columnTypes); Optional readerPageSourceWithProjections = pageSourceFactory .createPageSource( session, @@ -217,6 +212,7 @@ static ConnectorPageSource createPageSource( 0, targetFile.length(), targetFile.length(), + targetFile.lastModified(), schema, readColumns, TupleDomain.all(), @@ -244,7 +240,7 @@ static List getBaseColumns(List columnNames, List columnNames, List columnTypes) + static Schema createSchema(HiveStorageFormat format, List columnNames, List columnTypes) { Properties schema = new Properties(); schema.setProperty(SERIALIZATION_LIB, format.getSerde()); @@ -255,6 +251,6 @@ static Properties createSchema(HiveStorageFormat format, List columnName .map(HiveType::getHiveTypeName) .map(HiveTypeName::toString) .collect(joining(":"))); - return schema; + return new Schema(format.getSerde(), false, Maps.fromProperties(schema)); } } diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/benchmark/BenchmarkHiveFileFormat.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/benchmark/BenchmarkHiveFileFormat.java index 4670320d0e55..8421d01757c8 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/benchmark/BenchmarkHiveFileFormat.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/benchmark/BenchmarkHiveFileFormat.java @@ -81,10 +81,10 @@ public class BenchmarkHiveFileFormat { private static final ConnectorSession SESSION = getHiveSession( - new HiveConfig(), new ParquetReaderConfig().setOptimizedReaderEnabled(false)); + new HiveConfig(), new ParquetReaderConfig()); private static final ConnectorSession SESSION_OPTIMIZED_PARQUET_READER = getHiveSession( - new HiveConfig(), new ParquetReaderConfig().setOptimizedReaderEnabled(true)); + new HiveConfig(), new ParquetReaderConfig()); static { HadoopNative.requireHadoopNative(); diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/benchmark/StandardFileFormats.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/benchmark/StandardFileFormats.java index ee450abc9663..24d01fe8288e 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/benchmark/StandardFileFormats.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/benchmark/StandardFileFormats.java @@ -72,7 +72,7 @@ public HiveStorageFormat getFormat() @Override public Optional getHivePageSourceFactory(HdfsEnvironment hdfsEnvironment) { - return Optional.of(new RcFilePageSourceFactory(HDFS_FILE_SYSTEM_FACTORY, new FileFormatDataSourceStats(), new HiveConfig().setRcfileTimeZone("UTC"))); + return Optional.of(new RcFilePageSourceFactory(HDFS_FILE_SYSTEM_FACTORY, new HiveConfig().setRcfileTimeZone("UTC"))); } @Override @@ -103,7 +103,7 @@ public HiveStorageFormat getFormat() @Override public Optional getHivePageSourceFactory(HdfsEnvironment hdfsEnvironment) { - return Optional.of(new RcFilePageSourceFactory(HDFS_FILE_SYSTEM_FACTORY, new FileFormatDataSourceStats(), new HiveConfig().setRcfileTimeZone("UTC"))); + return Optional.of(new RcFilePageSourceFactory(HDFS_FILE_SYSTEM_FACTORY, new HiveConfig().setRcfileTimeZone("UTC"))); } @Override @@ -204,9 +204,8 @@ public PrestoParquetFormatWriter(File targetFile, List columnNames, List schemaConverter.getMessageType(), schemaConverter.getPrimitiveTypes(), ParquetWriterOptions.builder().build(), - compressionCodec.getParquetCompressionCodec(), + compressionCodec.getParquetCompressionCodec().orElseThrow(), "test-version", - false, Optional.of(DateTimeZone.getDefault()), Optional.empty()); } diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/containers/Hive3MinioDataLake.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/containers/Hive3MinioDataLake.java new file mode 100644 index 000000000000..fdd37cbe0ff8 --- /dev/null +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/containers/Hive3MinioDataLake.java @@ -0,0 +1,47 @@ +/* + * 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.hive.containers; + +import com.google.common.collect.ImmutableMap; + +import java.util.Map; + +import static io.trino.plugin.hive.containers.HiveMinioDataLake.State.STARTED; +import static io.trino.testing.containers.TestContainers.getPathFromClassPathResource; + +public class Hive3MinioDataLake + extends HiveMinioDataLake +{ + public Hive3MinioDataLake(String bucketName) + { + this(bucketName, HiveHadoop.HIVE3_IMAGE); + } + + public Hive3MinioDataLake(String bucketName, String hiveHadoopImage) + { + this(bucketName, ImmutableMap.of("/etc/hadoop/conf/core-site.xml", getPathFromClassPathResource("hive_minio_datalake/hive-core-site.xml")), hiveHadoopImage); + } + + public Hive3MinioDataLake(String bucketName, Map hiveHadoopFilesToMount, String hiveHadoopImage) + { + super(bucketName, hiveHadoopFilesToMount, hiveHadoopImage); + } + + @Override + public void start() + { + super.start(); + state = STARTED; + } +} diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/containers/HiveHadoop.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/containers/HiveHadoop.java index 6682ba610d68..0fe95c158bd7 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/containers/HiveHadoop.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/containers/HiveHadoop.java @@ -21,6 +21,7 @@ import io.trino.testing.containers.BaseTestContainer; import org.testcontainers.containers.Network; +import java.net.URI; import java.util.Map; import java.util.Optional; import java.util.Set; @@ -32,7 +33,6 @@ public class HiveHadoop { private static final Logger log = Logger.get(HiveHadoop.class); - public static final String DEFAULT_IMAGE = "ghcr.io/trinodb/testing/hdp2.6-hive:" + TestingProperties.getDockerImagesVersion(); public static final String HIVE3_IMAGE = "ghcr.io/trinodb/testing/hdp3.1-hive:" + TestingProperties.getDockerImagesVersion(); public static final String HOST_NAME = "hadoop-master"; @@ -93,9 +93,10 @@ public String runOnMetastore(String query) return executeInContainerFailOnError("mysql", "-D", "metastore", "-uroot", "-proot", "--batch", "--column-names=false", "-e", query).replaceAll("\n$", ""); } - public HostAndPort getHiveMetastoreEndpoint() + public URI getHiveMetastoreEndpoint() { - return getMappedHostAndPortForExposedPort(HIVE_METASTORE_PORT); + HostAndPort address = getMappedHostAndPortForExposedPort(HIVE_METASTORE_PORT); + return URI.create("thrift://" + address.getHost() + ":" + address.getPort()); } public static class Builder @@ -103,7 +104,7 @@ public static class Builder { private Builder() { - this.image = DEFAULT_IMAGE; + this.image = HIVE3_IMAGE; this.hostName = HOST_NAME; this.exposePorts = ImmutableSet.of(HIVE_METASTORE_PORT); } diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/containers/HiveMinioDataLake.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/containers/HiveMinioDataLake.java index 15e7f0178450..4610c074df92 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/containers/HiveMinioDataLake.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/containers/HiveMinioDataLake.java @@ -20,6 +20,7 @@ import io.trino.util.AutoCloseableCloser; import org.testcontainers.containers.Network; +import java.net.URI; import java.util.List; import java.util.Map; @@ -44,14 +45,14 @@ public class HiveMinioDataLake private final Minio minio; private final HiveHadoop hiveHadoop; - private final AutoCloseableCloser closer = AutoCloseableCloser.create(); - - private State state = State.INITIAL; + protected final AutoCloseableCloser closer = AutoCloseableCloser.create(); + protected final Network network; + protected State state = State.INITIAL; private MinioClient minioClient; public HiveMinioDataLake(String bucketName) { - this(bucketName, HiveHadoop.DEFAULT_IMAGE); + this(bucketName, HiveHadoop.HIVE3_IMAGE); } public HiveMinioDataLake(String bucketName, String hiveHadoopImage) @@ -62,7 +63,7 @@ public HiveMinioDataLake(String bucketName, String hiveHadoopImage) public HiveMinioDataLake(String bucketName, Map hiveHadoopFilesToMount, String hiveHadoopImage) { this.bucketName = requireNonNull(bucketName, "bucketName is null"); - Network network = closer.register(newNetwork()); + this.network = closer.register(newNetwork()); this.minio = closer.register( Minio.builder() .withNetwork(network) @@ -135,6 +136,11 @@ public HiveHadoop getHiveHadoop() return hiveHadoop; } + public URI getHiveMetastoreEndpoint() + { + return hiveHadoop.getHiveMetastoreEndpoint(); + } + public String getBucketName() { return bucketName; @@ -147,7 +153,7 @@ public void close() stop(); } - private enum State + protected enum State { INITIAL, STARTING, diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/fs/BaseCachingDirectoryListerTest.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/fs/BaseCachingDirectoryListerTest.java index ee26f071ecc7..e9590822c34b 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/fs/BaseCachingDirectoryListerTest.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/fs/BaseCachingDirectoryListerTest.java @@ -17,10 +17,12 @@ import com.google.common.collect.ImmutableMap; import io.trino.filesystem.Location; import io.trino.plugin.hive.HiveQueryRunner; +import io.trino.plugin.hive.metastore.HiveMetastore; +import io.trino.plugin.hive.metastore.HiveMetastoreFactory; import io.trino.plugin.hive.metastore.PrincipalPrivileges; import io.trino.plugin.hive.metastore.Table; -import io.trino.plugin.hive.metastore.file.FileHiveMetastore; import io.trino.testing.AbstractTestQueryFramework; +import io.trino.testing.DistributedQueryRunner; import io.trino.testing.MaterializedRow; import io.trino.testing.QueryRunner; import org.testng.annotations.Test; @@ -34,16 +36,16 @@ import static com.google.common.io.MoreFiles.deleteRecursively; import static com.google.common.io.RecursiveDeleteOption.ALLOW_INSECURE; import static io.trino.plugin.hive.HiveQueryRunner.TPCH_SCHEMA; -import static io.trino.plugin.hive.metastore.file.TestingFileHiveMetastore.createTestingFileHiveMetastore; +import static io.trino.plugin.hive.TestingHiveUtils.getConnectorService; import static java.lang.String.format; import static java.nio.file.Files.createTempDirectory; import static org.assertj.core.api.Assertions.assertThat; -public abstract class BaseCachingDirectoryListerTest +public abstract class BaseCachingDirectoryListerTest extends AbstractTestQueryFramework { - private C directoryLister; - private FileHiveMetastore fileHiveMetastore; + private CachingDirectoryLister directoryLister; + private HiveMetastore fileHiveMetastore; @Override protected QueryRunner createQueryRunner() @@ -57,17 +59,16 @@ protected QueryRunner createQueryRunner(Map properties) { Path temporaryMetastoreDirectory = createTempDirectory(null); closeAfterClass(() -> deleteRecursively(temporaryMetastoreDirectory, ALLOW_INSECURE)); - directoryLister = createDirectoryLister(); - return HiveQueryRunner.builder() + DistributedQueryRunner queryRunner = HiveQueryRunner.builder() .setHiveProperties(properties) - .setMetastore(distributedQueryRunner -> fileHiveMetastore = createTestingFileHiveMetastore(temporaryMetastoreDirectory.toFile())) - .setDirectoryLister(directoryLister) .build(); - } - protected abstract C createDirectoryLister(); + directoryLister = getConnectorService(queryRunner, CachingDirectoryLister.class); - protected abstract boolean isCached(C directoryLister, Location location); + fileHiveMetastore = getConnectorService(queryRunner, HiveMetastoreFactory.class) + .createMetastore(Optional.empty()); + return queryRunner; + } @Test public void testCacheInvalidationIsAppliedSpecificallyOnTheNonPartitionedTableBeingChanged() @@ -367,6 +368,6 @@ protected String getPartitionLocation(String schemaName, String tableName, List< protected boolean isCached(String path) { - return isCached(directoryLister, Location.of(path)); + return directoryLister.isCached(Location.of(path)); } } diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/fs/TestCachingDirectoryLister.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/fs/TestCachingDirectoryLister.java index 3d37a7e05fe5..bfc1c8dbae12 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/fs/TestCachingDirectoryLister.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/fs/TestCachingDirectoryLister.java @@ -13,30 +13,26 @@ */ package io.trino.plugin.hive.fs; -import io.airlift.units.DataSize; -import io.airlift.units.Duration; -import io.trino.filesystem.Location; +import com.google.common.collect.ImmutableMap; +import io.trino.testing.QueryRunner; import org.testng.annotations.Test; -import java.util.List; - -import static io.airlift.units.DataSize.Unit.MEGABYTE; - // some tests may invalidate the whole cache affecting therefore other concurrent tests @Test(singleThreaded = true) public class TestCachingDirectoryLister - extends BaseCachingDirectoryListerTest + extends BaseCachingDirectoryListerTest { @Override - protected CachingDirectoryLister createDirectoryLister() - { - return new CachingDirectoryLister(Duration.valueOf("5m"), DataSize.of(1, MEGABYTE), List.of("tpch.*")); - } - - @Override - protected boolean isCached(CachingDirectoryLister directoryLister, Location location) + protected QueryRunner createQueryRunner() + throws Exception { - return directoryLister.isCached(location); + return createQueryRunner(ImmutableMap.builder() + .put("hive.allow-register-partition-procedure", "true") + .put("hive.recursive-directories", "true") + .put("hive.file-status-cache-expire-time", "5m") + .put("hive.file-status-cache.max-retained-size", "1MB") + .put("hive.file-status-cache-tables", "tpch.*") + .buildOrThrow()); } @Test diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/fs/TestCachingDirectoryListerRecursiveFilesOnly.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/fs/TestCachingDirectoryListerRecursiveFilesOnly.java index 010b8330145c..e6e0dd7ef9b1 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/fs/TestCachingDirectoryListerRecursiveFilesOnly.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/fs/TestCachingDirectoryListerRecursiveFilesOnly.java @@ -15,18 +15,13 @@ import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; -import io.airlift.units.DataSize; -import io.airlift.units.Duration; -import io.trino.filesystem.Location; import io.trino.plugin.hive.metastore.MetastoreUtil; import io.trino.plugin.hive.metastore.Table; import io.trino.testing.QueryRunner; import org.testng.annotations.Test; -import java.util.List; import java.util.NoSuchElementException; -import static io.airlift.units.DataSize.Unit.MEGABYTE; import static io.trino.plugin.hive.HiveQueryRunner.TPCH_SCHEMA; import static io.trino.plugin.hive.metastore.PrincipalPrivileges.NO_PRIVILEGES; import static java.lang.String.format; @@ -36,27 +31,19 @@ // some tests may invalidate the whole cache affecting therefore other concurrent tests @Test(singleThreaded = true) public class TestCachingDirectoryListerRecursiveFilesOnly - extends BaseCachingDirectoryListerTest + extends BaseCachingDirectoryListerTest { - @Override - protected CachingDirectoryLister createDirectoryLister() - { - return new CachingDirectoryLister(Duration.valueOf("5m"), DataSize.of(1, MEGABYTE), List.of("tpch.*")); - } - @Override protected QueryRunner createQueryRunner() throws Exception { - return createQueryRunner(ImmutableMap.of( - "hive.allow-register-partition-procedure", "true", - "hive.recursive-directories", "true")); - } - - @Override - protected boolean isCached(CachingDirectoryLister directoryLister, Location location) - { - return directoryLister.isCached(location); + return createQueryRunner(ImmutableMap.builder() + .put("hive.allow-register-partition-procedure", "true") + .put("hive.recursive-directories", "true") + .put("hive.file-status-cache-expire-time", "5m") + .put("hive.file-status-cache.max-retained-size", "1MB") + .put("hive.file-status-cache-tables", "tpch.*") + .buildOrThrow()); } @Test diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/fs/TestTransactionScopeCachingDirectoryLister.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/fs/TestTransactionScopeCachingDirectoryLister.java index a67bb89bf514..6704fb69f081 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/fs/TestTransactionScopeCachingDirectoryLister.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/fs/TestTransactionScopeCachingDirectoryLister.java @@ -26,6 +26,7 @@ import io.trino.plugin.hive.metastore.Storage; import io.trino.plugin.hive.metastore.StorageFormat; import io.trino.plugin.hive.metastore.Table; +import io.trino.testing.QueryRunner; import org.testng.annotations.Test; import java.io.IOException; @@ -35,8 +36,6 @@ import java.util.Optional; import java.util.OptionalLong; -import static io.airlift.units.DataSize.Unit.MEGABYTE; -import static io.trino.plugin.hive.util.HiveBucketing.BucketingVersion.BUCKETING_V1; import static java.util.Objects.requireNonNull; import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatThrownBy; @@ -44,8 +43,21 @@ // some tests may invalidate the whole cache affecting therefore other concurrent tests @Test(singleThreaded = true) public class TestTransactionScopeCachingDirectoryLister - extends BaseCachingDirectoryListerTest + extends BaseCachingDirectoryListerTest { + @Override + protected QueryRunner createQueryRunner() + throws Exception + { + return createQueryRunner(ImmutableMap.builder() + .put("hive.allow-register-partition-procedure", "true") + .put("hive.recursive-directories", "true") + .put("hive.file-status-cache-expire-time", "5m") + .put("hive.file-status-cache.max-retained-size", "1MB") + .put("hive.file-status-cache-tables", "tpch.*") + .buildOrThrow()); + } + private static final Column TABLE_COLUMN = new Column( "column", HiveType.HIVE_INT, @@ -53,7 +65,7 @@ public class TestTransactionScopeCachingDirectoryLister private static final Storage TABLE_STORAGE = new Storage( StorageFormat.create("serde", "input", "output"), Optional.of("location"), - Optional.of(new HiveBucketProperty(ImmutableList.of("column"), BUCKETING_V1, 10, ImmutableList.of(new SortingColumn("column", SortingColumn.Order.ASCENDING)))), + Optional.of(new HiveBucketProperty(ImmutableList.of("column"), 10, ImmutableList.of(new SortingColumn("column", SortingColumn.Order.ASCENDING)))), true, ImmutableMap.of("param", "value2")); private static final Table TABLE = new Table( @@ -69,18 +81,6 @@ public class TestTransactionScopeCachingDirectoryLister Optional.of("expanded_text"), OptionalLong.empty()); - @Override - protected TransactionScopeCachingDirectoryLister createDirectoryLister() - { - return (TransactionScopeCachingDirectoryLister) new TransactionScopeCachingDirectoryListerFactory(DataSize.of(1, MEGABYTE), Optional.empty()).get(new FileSystemDirectoryLister()); - } - - @Override - protected boolean isCached(TransactionScopeCachingDirectoryLister directoryLister, Location location) - { - return directoryLister.isCached(location); - } - @Test public void testConcurrentDirectoryListing() throws IOException diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/CountingAccessHiveMetastore.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/CountingAccessHiveMetastore.java index 5f4a7cd914d9..ef06dc88c097 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/CountingAccessHiveMetastore.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/CountingAccessHiveMetastore.java @@ -15,64 +15,34 @@ import com.google.common.collect.ConcurrentHashMultiset; import com.google.common.collect.ImmutableMultiset; +import com.google.common.collect.ImmutableSet; import com.google.common.collect.Multiset; import com.google.errorprone.annotations.ThreadSafe; -import io.trino.plugin.hive.HiveColumnStatisticType; import io.trino.plugin.hive.HiveType; import io.trino.plugin.hive.PartitionStatistics; -import io.trino.plugin.hive.acid.AcidTransaction; import io.trino.plugin.hive.metastore.HivePrivilegeInfo.HivePrivilege; -import io.trino.spi.connector.SchemaTableName; import io.trino.spi.predicate.TupleDomain; import io.trino.spi.security.RoleGrant; -import io.trino.spi.type.Type; import java.util.List; import java.util.Map; import java.util.Optional; +import java.util.OptionalLong; import java.util.Set; -import java.util.function.Function; - -import static io.trino.plugin.hive.metastore.CountingAccessHiveMetastore.Method.GET_ALL_TABLES; -import static io.trino.plugin.hive.metastore.CountingAccessHiveMetastore.Method.GET_ALL_VIEWS; @ThreadSafe public class CountingAccessHiveMetastore implements HiveMetastore { - public enum Method - { - CREATE_DATABASE, - CREATE_TABLE, - GET_ALL_DATABASES, - GET_DATABASE, - GET_TABLE, - GET_ALL_TABLES, - GET_ALL_TABLES_FROM_DATABASE, - GET_TABLE_WITH_PARAMETER, - GET_TABLE_STATISTICS, - GET_ALL_VIEWS, - GET_ALL_VIEWS_FROM_DATABASE, - UPDATE_TABLE_STATISTICS, - ADD_PARTITIONS, - GET_PARTITION_NAMES_BY_FILTER, - GET_PARTITIONS_BY_NAMES, - GET_PARTITION, - GET_PARTITION_STATISTICS, - UPDATE_PARTITION_STATISTICS, - REPLACE_TABLE, - DROP_TABLE, - } - private final HiveMetastore delegate; - private final ConcurrentHashMultiset methodInvocations = ConcurrentHashMultiset.create(); + private final ConcurrentHashMultiset methodInvocations = ConcurrentHashMultiset.create(); public CountingAccessHiveMetastore(HiveMetastore delegate) { this.delegate = delegate; } - public Multiset getMethodInvocations() + public Multiset getMethodInvocations() { return ImmutableMultiset.copyOf(methodInvocations); } @@ -85,66 +55,77 @@ public void resetCounters() @Override public Optional
getTable(String databaseName, String tableName) { - methodInvocations.add(Method.GET_TABLE); + methodInvocations.add(MetastoreMethod.GET_TABLE); return delegate.getTable(databaseName, tableName); } @Override - public Set getSupportedColumnStatistics(Type type) + public Map getTableColumnStatistics(String databaseName, String tableName, Set columnNames) { - // No need to count that, since it's a pure local operation. - return delegate.getSupportedColumnStatistics(type); + methodInvocations.add(MetastoreMethod.GET_TABLE_STATISTICS); + return delegate.getTableColumnStatistics(databaseName, tableName, columnNames); } @Override - public List getAllDatabases() + public Map> getPartitionColumnStatistics(String databaseName, String tableName, Set partitionNames, Set columnNames) { - methodInvocations.add(Method.GET_ALL_DATABASES); - return delegate.getAllDatabases(); + methodInvocations.add(MetastoreMethod.GET_PARTITION_STATISTICS); + return delegate.getPartitionColumnStatistics(databaseName, tableName, partitionNames, columnNames); } @Override - public Optional getDatabase(String databaseName) + public void updateTableStatistics(String databaseName, String tableName, OptionalLong acidWriteId, StatisticsUpdateMode mode, PartitionStatistics statisticsUpdate) { - methodInvocations.add(Method.GET_DATABASE); - return delegate.getDatabase(databaseName); + methodInvocations.add(MetastoreMethod.UPDATE_TABLE_STATISTICS); + delegate.updateTableStatistics(databaseName, tableName, acidWriteId, mode, statisticsUpdate); + } + + @Override + public void updatePartitionStatistics(Table table, StatisticsUpdateMode mode, Map partitionUpdates) + { + methodInvocations.add(MetastoreMethod.UPDATE_PARTITION_STATISTICS); + delegate.updatePartitionStatistics(table, mode, partitionUpdates); + } + + @Override + public List getTables(String databaseName) + { + methodInvocations.add(MetastoreMethod.GET_ALL_TABLES_FROM_DATABASE); + return delegate.getTables(databaseName); } @Override - public List getTablesWithParameter(String databaseName, String parameterKey, String parameterValue) + public List getTableNamesWithParameters(String databaseName, String parameterKey, ImmutableSet parameterValues) { - methodInvocations.add(Method.GET_TABLE_WITH_PARAMETER); - return delegate.getTablesWithParameter(databaseName, parameterKey, parameterValue); + methodInvocations.add(MetastoreMethod.GET_TABLE_WITH_PARAMETER); + return delegate.getTableNamesWithParameters(databaseName, parameterKey, parameterValues); } @Override - public List getAllViews(String databaseName) + public List getAllDatabases() { - methodInvocations.add(Method.GET_ALL_VIEWS_FROM_DATABASE); - return delegate.getAllViews(databaseName); + methodInvocations.add(MetastoreMethod.GET_ALL_DATABASES); + return delegate.getAllDatabases(); } @Override - public Optional> getAllViews() + public Optional getDatabase(String databaseName) { - Optional> allViews = delegate.getAllViews(); - if (allViews.isPresent()) { - methodInvocations.add(GET_ALL_VIEWS); - } - return allViews; + methodInvocations.add(MetastoreMethod.GET_DATABASE); + return delegate.getDatabase(databaseName); } @Override public void createDatabase(Database database) { - methodInvocations.add(Method.CREATE_DATABASE); + methodInvocations.add(MetastoreMethod.CREATE_DATABASE); delegate.createDatabase(database); } @Override public void dropDatabase(String databaseName, boolean deleteData) { - throw new UnsupportedOperationException(); + delegate.dropDatabase(databaseName, deleteData); } @Override @@ -162,22 +143,22 @@ public void setDatabaseOwner(String databaseName, HivePrincipal principal) @Override public void createTable(Table table, PrincipalPrivileges principalPrivileges) { - methodInvocations.add(Method.CREATE_TABLE); + methodInvocations.add(MetastoreMethod.CREATE_TABLE); delegate.createTable(table, principalPrivileges); } @Override public void dropTable(String databaseName, String tableName, boolean deleteData) { - methodInvocations.add(Method.DROP_TABLE); + methodInvocations.add(MetastoreMethod.DROP_TABLE); delegate.dropTable(databaseName, tableName, deleteData); } @Override - public void replaceTable(String databaseName, String tableName, Table newTable, PrincipalPrivileges principalPrivileges) + public void replaceTable(String databaseName, String tableName, Table newTable, PrincipalPrivileges principalPrivileges, Map environmentContext) { - methodInvocations.add(Method.REPLACE_TABLE); - delegate.replaceTable(databaseName, tableName, newTable, principalPrivileges); + methodInvocations.add(MetastoreMethod.REPLACE_TABLE); + delegate.replaceTable(databaseName, tableName, newTable, principalPrivileges, environmentContext); } @Override @@ -225,7 +206,7 @@ public void dropColumn(String databaseName, String tableName, String columnName) @Override public Optional getPartition(Table table, List partitionValues) { - methodInvocations.add(Method.GET_PARTITION); + methodInvocations.add(MetastoreMethod.GET_PARTITION); return delegate.getPartition(table, partitionValues); } @@ -235,21 +216,21 @@ public Optional> getPartitionNamesByFilter(String databaseName, List columnNames, TupleDomain partitionKeysFilter) { - methodInvocations.add(Method.GET_PARTITION_NAMES_BY_FILTER); + methodInvocations.add(MetastoreMethod.GET_PARTITION_NAMES_BY_FILTER); return delegate.getPartitionNamesByFilter(databaseName, tableName, columnNames, partitionKeysFilter); } @Override public Map> getPartitionsByNames(Table table, List partitionNames) { - methodInvocations.add(Method.GET_PARTITIONS_BY_NAMES); + methodInvocations.add(MetastoreMethod.GET_PARTITIONS_BY_NAMES); return delegate.getPartitionsByNames(table, partitionNames); } @Override public void addPartitions(String databaseName, String tableName, List partitions) { - methodInvocations.add(Method.ADD_PARTITIONS); + methodInvocations.add(MetastoreMethod.ADD_PARTITIONS); delegate.addPartitions(databaseName, tableName, partitions); } @@ -295,16 +276,10 @@ public void revokeRoles(Set roles, Set grantees, boolean throw new UnsupportedOperationException(); } - @Override - public Set listGrantedPrincipals(String role) - { - throw new UnsupportedOperationException(); - } - @Override public Set listRoleGrants(HivePrincipal principal) { - throw new UnsupportedOperationException(); + return Set.of(); } @Override @@ -322,54 +297,6 @@ public void revokeTablePrivileges(String databaseName, String tableName, String @Override public Set listTablePrivileges(String databaseName, String tableName, Optional tableOwner, Optional principal) { - throw new UnsupportedOperationException(); - } - - @Override - public PartitionStatistics getTableStatistics(Table table) - { - methodInvocations.add(Method.GET_TABLE_STATISTICS); - return delegate.getTableStatistics(table); - } - - @Override - public Map getPartitionStatistics(Table table, List partitions) - { - methodInvocations.add(Method.GET_PARTITION_STATISTICS); - return delegate.getPartitionStatistics(table, partitions); - } - - @Override - public void updateTableStatistics(String databaseName, - String tableName, - AcidTransaction transaction, - Function update) - { - methodInvocations.add(Method.UPDATE_TABLE_STATISTICS); - delegate.updateTableStatistics(databaseName, tableName, transaction, update); - } - - @Override - public void updatePartitionStatistics(Table table, Map> updates) - { - methodInvocations.add(Method.UPDATE_PARTITION_STATISTICS); - delegate.updatePartitionStatistics(table, updates); - } - - @Override - public List getAllTables(String databaseName) - { - methodInvocations.add(Method.GET_ALL_TABLES_FROM_DATABASE); - return delegate.getAllTables(databaseName); - } - - @Override - public Optional> getAllTables() - { - Optional> allTables = delegate.getAllTables(); - if (allTables.isPresent()) { - methodInvocations.add(GET_ALL_TABLES); - } - return allTables; + return Set.of(); } } diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/TestMetastoreUtil.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/TestMetastoreUtil.java index 94cf2db5ff5a..d446f9eb3afa 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/TestMetastoreUtil.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/TestMetastoreUtil.java @@ -32,7 +32,6 @@ import java.util.List; import java.util.Map; import java.util.Optional; -import java.util.Properties; import static io.airlift.slice.Slices.utf8Slice; import static io.trino.plugin.hive.HiveColumnHandle.ColumnType.PARTITION_KEY; @@ -174,14 +173,14 @@ public void testPartitionRoundTrip() @Test public void testHiveSchemaTable() { - Properties actual = MetastoreUtil.getHiveSchema(ThriftMetastoreUtil.fromMetastoreApiTable(TEST_TABLE_WITH_UNSUPPORTED_FIELDS, TEST_SCHEMA)); + Map actual = MetastoreUtil.getHiveSchema(ThriftMetastoreUtil.fromMetastoreApiTable(TEST_TABLE_WITH_UNSUPPORTED_FIELDS, TEST_SCHEMA)); assertEquals(actual, TEST_TABLE_METADATA); } @Test public void testHiveSchemaPartition() { - Properties actual = MetastoreUtil.getHiveSchema(ThriftMetastoreUtil.fromMetastoreApiPartition(TEST_PARTITION_WITH_UNSUPPORTED_FIELDS), ThriftMetastoreUtil.fromMetastoreApiTable(TEST_TABLE_WITH_UNSUPPORTED_FIELDS, TEST_SCHEMA)); + Map actual = MetastoreUtil.getHiveSchema(ThriftMetastoreUtil.fromMetastoreApiPartition(TEST_PARTITION_WITH_UNSUPPORTED_FIELDS), ThriftMetastoreUtil.fromMetastoreApiTable(TEST_TABLE_WITH_UNSUPPORTED_FIELDS, TEST_SCHEMA)); assertEquals(actual, TEST_TABLE_METADATA); } diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/TestSemiTransactionalHiveMetastore.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/TestSemiTransactionalHiveMetastore.java deleted file mode 100644 index 6270c025aa65..000000000000 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/TestSemiTransactionalHiveMetastore.java +++ /dev/null @@ -1,190 +0,0 @@ -/* - * 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.hive.metastore; - -import com.google.common.collect.ImmutableList; -import com.google.common.collect.ImmutableMap; -import io.trino.filesystem.Location; -import io.trino.plugin.hive.HiveBucketProperty; -import io.trino.plugin.hive.HiveMetastoreClosure; -import io.trino.plugin.hive.HiveType; -import io.trino.plugin.hive.PartitionStatistics; -import io.trino.plugin.hive.acid.AcidTransaction; -import io.trino.plugin.hive.fs.FileSystemDirectoryLister; -import org.testng.annotations.Test; - -import java.util.List; -import java.util.Optional; -import java.util.OptionalLong; -import java.util.concurrent.CountDownLatch; -import java.util.concurrent.Executor; -import java.util.concurrent.TimeUnit; -import java.util.function.Function; -import java.util.stream.IntStream; - -import static com.google.common.util.concurrent.MoreExecutors.directExecutor; -import static io.trino.plugin.hive.HiveBasicStatistics.createEmptyStatistics; -import static io.trino.plugin.hive.HiveTestUtils.HDFS_ENVIRONMENT; -import static io.trino.plugin.hive.acid.AcidOperation.INSERT; -import static io.trino.plugin.hive.util.HiveBucketing.BucketingVersion.BUCKETING_V1; -import static io.trino.testing.TestingConnectorSession.SESSION; -import static java.util.concurrent.Executors.newFixedThreadPool; -import static java.util.concurrent.Executors.newScheduledThreadPool; -import static org.testng.Assert.assertTrue; - -// countDownLatch field is shared between tests -@Test(singleThreaded = true) -public class TestSemiTransactionalHiveMetastore -{ - private static final Column TABLE_COLUMN = new Column( - "column", - HiveType.HIVE_INT, - Optional.of("comment")); - private static final Storage TABLE_STORAGE = new Storage( - StorageFormat.create("serde", "input", "output"), - Optional.of("/test"), - Optional.of(new HiveBucketProperty(ImmutableList.of("column"), BUCKETING_V1, 10, ImmutableList.of(new SortingColumn("column", SortingColumn.Order.ASCENDING)))), - true, - ImmutableMap.of("param", "value2")); - - private CountDownLatch countDownLatch; - - @Test - public void testParallelPartitionDrops() - { - int partitionsToDrop = 5; - IntStream dropThreadsConfig = IntStream.of(1, 2); - dropThreadsConfig.forEach(dropThreads -> { - countDownLatch = new CountDownLatch(dropThreads); - SemiTransactionalHiveMetastore semiTransactionalHiveMetastore = getSemiTransactionalHiveMetastoreWithDropExecutor(newFixedThreadPool(dropThreads)); - IntStream.range(0, partitionsToDrop).forEach(i -> semiTransactionalHiveMetastore.dropPartition(SESSION, - "test", - "test", - ImmutableList.of(String.valueOf(i)), - true)); - semiTransactionalHiveMetastore.commit(); - }); - } - - private SemiTransactionalHiveMetastore getSemiTransactionalHiveMetastoreWithDropExecutor(Executor dropExecutor) - { - return new SemiTransactionalHiveMetastore(HDFS_ENVIRONMENT, - new HiveMetastoreClosure(new TestingHiveMetastore()), - directExecutor(), - dropExecutor, - directExecutor(), - false, - false, - true, - Optional.empty(), - newScheduledThreadPool(1), - new FileSystemDirectoryLister()); - } - - @Test - public void testParallelUpdateStatisticsOperations() - { - int tablesToUpdate = 5; - IntStream updateThreadsConfig = IntStream.of(1, 2); - updateThreadsConfig.forEach(updateThreads -> { - countDownLatch = new CountDownLatch(updateThreads); - SemiTransactionalHiveMetastore semiTransactionalHiveMetastore; - if (updateThreads == 1) { - semiTransactionalHiveMetastore = getSemiTransactionalHiveMetastoreWithUpdateExecutor(directExecutor()); - } - else { - semiTransactionalHiveMetastore = getSemiTransactionalHiveMetastoreWithUpdateExecutor(newFixedThreadPool(updateThreads)); - } - IntStream.range(0, tablesToUpdate).forEach(i -> semiTransactionalHiveMetastore.finishChangingExistingTable(INSERT, SESSION, - "database", - "table_" + i, - Location.of(TABLE_STORAGE.getLocation()), - ImmutableList.of(), - PartitionStatistics.empty(), - false)); - semiTransactionalHiveMetastore.commit(); - }); - } - - private SemiTransactionalHiveMetastore getSemiTransactionalHiveMetastoreWithUpdateExecutor(Executor updateExecutor) - { - return new SemiTransactionalHiveMetastore(HDFS_ENVIRONMENT, - new HiveMetastoreClosure(new TestingHiveMetastore()), - directExecutor(), - directExecutor(), - updateExecutor, - false, - false, - true, - Optional.empty(), - newScheduledThreadPool(1), - new FileSystemDirectoryLister()); - } - - private class TestingHiveMetastore - extends UnimplementedHiveMetastore - { - @Override - public Optional
getTable(String databaseName, String tableName) - { - if (databaseName.equals("database")) { - return Optional.of(new Table( - "database", - tableName, - Optional.of("owner"), - "table_type", - TABLE_STORAGE, - ImmutableList.of(TABLE_COLUMN), - ImmutableList.of(TABLE_COLUMN), - ImmutableMap.of("param", "value3"), - Optional.of("original_text"), - Optional.of("expanded_text"), - OptionalLong.empty())); - } - return Optional.empty(); - } - - @Override - public PartitionStatistics getTableStatistics(Table table) - { - return new PartitionStatistics(createEmptyStatistics(), ImmutableMap.of()); - } - - @Override - public void dropPartition(String databaseName, String tableName, List parts, boolean deleteData) - { - assertCountDownLatch(); - } - - @Override - public void updateTableStatistics(String databaseName, - String tableName, - AcidTransaction transaction, - Function update) - { - assertCountDownLatch(); - } - - private void assertCountDownLatch() - { - try { - countDownLatch.countDown(); - assertTrue(countDownLatch.await(10, TimeUnit.SECONDS)); //all other threads launched should count down within 10 seconds - } - catch (InterruptedException ie) { - Thread.currentThread().interrupt(); - } - } - } -} diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/UnimplementedHiveMetastore.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/UnimplementedHiveMetastore.java index 8a14aeb6a838..7dc3ce59f5e6 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/UnimplementedHiveMetastore.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/UnimplementedHiveMetastore.java @@ -13,21 +13,18 @@ */ package io.trino.plugin.hive.metastore; -import io.trino.plugin.hive.HiveColumnStatisticType; +import com.google.common.collect.ImmutableSet; import io.trino.plugin.hive.HiveType; import io.trino.plugin.hive.PartitionStatistics; -import io.trino.plugin.hive.acid.AcidTransaction; import io.trino.plugin.hive.metastore.HivePrivilegeInfo.HivePrivilege; -import io.trino.spi.connector.SchemaTableName; import io.trino.spi.predicate.TupleDomain; import io.trino.spi.security.RoleGrant; -import io.trino.spi.type.Type; import java.util.List; import java.util.Map; import java.util.Optional; +import java.util.OptionalLong; import java.util.Set; -import java.util.function.Function; public class UnimplementedHiveMetastore implements HiveMetastore @@ -51,64 +48,37 @@ public Optional
getTable(String databaseName, String tableName) } @Override - public Set getSupportedColumnStatistics(Type type) + public Map getTableColumnStatistics(String databaseName, String tableName, Set columnNames) { throw new UnsupportedOperationException(); } @Override - public PartitionStatistics getTableStatistics(Table table) + public Map> getPartitionColumnStatistics(String databaseName, String tableName, Set partitionNames, Set columnNames) { throw new UnsupportedOperationException(); } @Override - public Map getPartitionStatistics(Table table, List partitions) + public void updateTableStatistics(String databaseName, String tableName, OptionalLong acidWriteId, StatisticsUpdateMode mode, PartitionStatistics statisticsUpdate) { throw new UnsupportedOperationException(); } @Override - public void updateTableStatistics(String databaseName, - String tableName, - AcidTransaction transaction, - Function update) - { - throw new UnsupportedOperationException(); - } - - @Override - public void updatePartitionStatistics(Table table, Map> updates) - { - throw new UnsupportedOperationException(); - } - - @Override - public List getAllTables(String databaseName) - { - throw new UnsupportedOperationException(); - } - - @Override - public Optional> getAllTables() + public void updatePartitionStatistics(Table table, StatisticsUpdateMode mode, Map partitionUpdates) { throw new UnsupportedOperationException(); } @Override - public List getTablesWithParameter(String databaseName, String parameterKey, String parameterValue) + public List getTables(String databaseName) { throw new UnsupportedOperationException(); } @Override - public List getAllViews(String databaseName) - { - throw new UnsupportedOperationException(); - } - - @Override - public Optional> getAllViews() + public List getTableNamesWithParameters(String databaseName, String parameterKey, ImmutableSet parameterValues) { throw new UnsupportedOperationException(); } @@ -156,7 +126,7 @@ public void dropTable(String databaseName, String tableName, boolean deleteData) } @Override - public void replaceTable(String databaseName, String tableName, Table newTable, PrincipalPrivileges principalPrivileges) + public void replaceTable(String databaseName, String tableName, Table newTable, PrincipalPrivileges principalPrivileges, Map environmentContext) { throw new UnsupportedOperationException(); } @@ -284,12 +254,6 @@ public void revokeRoles(Set roles, Set grantees, boolean throw new UnsupportedOperationException(); } - @Override - public Set listGrantedPrincipals(String role) - { - throw new UnsupportedOperationException(); - } - @Override public Set listRoleGrants(HivePrincipal principal) { diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/cache/TestCachingHiveMetastore.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/cache/TestCachingHiveMetastore.java index f8486585c11c..6e2e86636919 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/cache/TestCachingHiveMetastore.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/cache/TestCachingHiveMetastore.java @@ -15,29 +15,25 @@ import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; -import com.google.common.collect.ImmutableMultimap; +import com.google.common.collect.ImmutableSet; import com.google.common.collect.Iterables; import com.google.common.util.concurrent.ListeningExecutorService; -import com.google.common.util.concurrent.ThreadFactoryBuilder; import io.airlift.log.Logger; import io.airlift.units.Duration; import io.trino.hive.thrift.metastore.ColumnStatisticsData; import io.trino.hive.thrift.metastore.ColumnStatisticsObj; import io.trino.hive.thrift.metastore.LongColumnStatsData; +import io.trino.plugin.base.util.AutoCloseableCloser; import io.trino.plugin.hive.HiveBasicStatistics; import io.trino.plugin.hive.HiveColumnHandle; -import io.trino.plugin.hive.HiveMetastoreClosure; import io.trino.plugin.hive.PartitionStatistics; -import io.trino.plugin.hive.metastore.Column; import io.trino.plugin.hive.metastore.Database; import io.trino.plugin.hive.metastore.HiveColumnStatistics; import io.trino.plugin.hive.metastore.HiveMetastore; import io.trino.plugin.hive.metastore.HivePrincipal; import io.trino.plugin.hive.metastore.Partition; -import io.trino.plugin.hive.metastore.PrincipalPrivileges; import io.trino.plugin.hive.metastore.Table; -import io.trino.plugin.hive.metastore.UnimplementedHiveMetastore; -import io.trino.plugin.hive.metastore.cache.CachingHiveMetastore.CachingHiveMetastoreBuilder; +import io.trino.plugin.hive.metastore.TableInfo; import io.trino.plugin.hive.metastore.thrift.BridgingHiveMetastore; import io.trino.plugin.hive.metastore.thrift.MockThriftMetastoreClient; import io.trino.plugin.hive.metastore.thrift.ThriftHiveMetastore; @@ -49,23 +45,18 @@ import io.trino.spi.predicate.Range; import io.trino.spi.predicate.TupleDomain; import io.trino.spi.predicate.ValueSet; -import io.trino.testing.DataProviders; import org.apache.thrift.TException; import org.testng.annotations.AfterClass; import org.testng.annotations.BeforeMethod; import org.testng.annotations.Test; -import java.util.ArrayList; -import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.Optional; import java.util.OptionalLong; -import java.util.concurrent.ConcurrentHashMap; -import java.util.concurrent.ConcurrentMap; -import java.util.concurrent.CopyOnWriteArrayList; import java.util.concurrent.CountDownLatch; import java.util.concurrent.CyclicBarrier; +import java.util.concurrent.Executor; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.Future; @@ -73,25 +64,18 @@ import java.util.function.Consumer; import static com.google.common.base.Preconditions.checkState; -import static com.google.common.util.concurrent.MoreExecutors.directExecutor; import static com.google.common.util.concurrent.MoreExecutors.listeningDecorator; import static io.airlift.concurrent.Threads.daemonThreadsNamed; import static io.airlift.slice.Slices.utf8Slice; import static io.trino.plugin.hive.HiveColumnHandle.ColumnType.PARTITION_KEY; import static io.trino.plugin.hive.HiveColumnHandle.createBaseColumn; -import static io.trino.plugin.hive.HiveStorageFormat.TEXTFILE; -import static io.trino.plugin.hive.HiveType.HIVE_LONG; import static io.trino.plugin.hive.HiveType.HIVE_STRING; -import static io.trino.plugin.hive.HiveType.toHiveType; -import static io.trino.plugin.hive.TableType.VIRTUAL_VIEW; import static io.trino.plugin.hive.TestingThriftHiveMetastoreBuilder.testingThriftHiveMetastoreBuilder; import static io.trino.plugin.hive.metastore.HiveColumnStatistics.createIntegerColumnStatistics; import static io.trino.plugin.hive.metastore.MetastoreUtil.computePartitionKeyFilter; import static io.trino.plugin.hive.metastore.MetastoreUtil.makePartitionName; -import static io.trino.plugin.hive.metastore.StorageFormat.VIEW_STORAGE_FORMAT; -import static io.trino.plugin.hive.metastore.StorageFormat.fromHiveStorageFormat; -import static io.trino.plugin.hive.metastore.cache.CachingHiveMetastore.memoizeMetastore; -import static io.trino.plugin.hive.metastore.cache.TestCachingHiveMetastore.PartitionCachingAssertions.assertThatCachingWithDisabledPartitionCache; +import static io.trino.plugin.hive.metastore.StatisticsUpdateMode.MERGE_INCREMENTAL; +import static io.trino.plugin.hive.metastore.cache.CachingHiveMetastore.createPerTransactionCache; import static io.trino.plugin.hive.metastore.thrift.MockThriftMetastoreClient.BAD_DATABASE; import static io.trino.plugin.hive.metastore.thrift.MockThriftMetastoreClient.BAD_PARTITION; import static io.trino.plugin.hive.metastore.thrift.MockThriftMetastoreClient.PARTITION_COLUMN_NAMES; @@ -112,9 +96,8 @@ import static java.lang.String.format; import static java.util.concurrent.Executors.newCachedThreadPool; import static java.util.concurrent.TimeUnit.SECONDS; -import static java.util.function.Function.identity; -import static org.apache.hadoop.hive.metastore.TableType.EXTERNAL_TABLE; import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertFalse; import static org.testng.Assert.assertNotNull; @@ -125,15 +108,20 @@ public class TestCachingHiveMetastore { private static final Logger log = Logger.get(TestCachingHiveMetastore.class); + private static final HiveBasicStatistics TEST_BASIC_STATS = new HiveBasicStatistics(OptionalLong.empty(), OptionalLong.of(2398040535435L), OptionalLong.empty(), OptionalLong.empty()); + private static final ImmutableMap TEST_COLUMN_STATS = ImmutableMap.of(TEST_COLUMN, createIntegerColumnStatistics(OptionalLong.empty(), OptionalLong.empty(), OptionalLong.empty(), OptionalLong.empty())); private static final PartitionStatistics TEST_STATS = PartitionStatistics.builder() .setBasicStatistics(new HiveBasicStatistics(OptionalLong.empty(), OptionalLong.of(2398040535435L), OptionalLong.empty(), OptionalLong.empty())) .setColumnStatistics(ImmutableMap.of(TEST_COLUMN, createIntegerColumnStatistics(OptionalLong.empty(), OptionalLong.empty(), OptionalLong.empty(), OptionalLong.empty()))) .build(); private static final SchemaTableName TEST_SCHEMA_TABLE = new SchemaTableName(TEST_DATABASE, TEST_TABLE); + private static final TableInfo TEST_TABLE_INFO = new TableInfo(TEST_SCHEMA_TABLE, TableInfo.ExtendedRelationType.TABLE); + private static final Duration CACHE_TTL = new Duration(5, TimeUnit.MINUTES); + private AutoCloseableCloser closer; private MockThriftMetastoreClient mockClient; + private ThriftMetastore thriftHiveMetastore; private ListeningExecutorService executor; - private CachingHiveMetastoreBuilder metastoreBuilder; private CachingHiveMetastore metastore; private CachingHiveMetastore statsOnlyCacheMetastore; private ThriftMetastoreStats stats; @@ -141,36 +129,26 @@ public class TestCachingHiveMetastore @BeforeMethod public void setUp() { + closer = AutoCloseableCloser.create(); mockClient = new MockThriftMetastoreClient(); - ThriftMetastore thriftHiveMetastore = createThriftHiveMetastore(); + thriftHiveMetastore = createThriftHiveMetastore(); executor = listeningDecorator(newCachedThreadPool(daemonThreadsNamed(getClass().getSimpleName() + "-%s"))); - metastoreBuilder = CachingHiveMetastore.builder() - .delegate(new BridgingHiveMetastore(thriftHiveMetastore)) - .executor(executor) - .metadataCacheEnabled(true) - .statsCacheEnabled(true) - .cacheTtl(new Duration(5, TimeUnit.MINUTES)) - .refreshInterval(new Duration(1, TimeUnit.MINUTES)) - .maximumSize(1000) - .cacheMissing(new CachingHiveMetastoreConfig().isCacheMissing()) - .partitionCacheEnabled(true); - - metastore = metastoreBuilder.build(); - statsOnlyCacheMetastore = CachingHiveMetastore.builder(metastoreBuilder) - .metadataCacheEnabled(false) - .statsCacheEnabled(true) // only cache stats - .build(); + metastore = createCachingHiveMetastore(new BridgingHiveMetastore(thriftHiveMetastore), CACHE_TTL, true, true, executor); + statsOnlyCacheMetastore = createCachingHiveMetastore(new BridgingHiveMetastore(thriftHiveMetastore), Duration.ZERO, true, true, executor); stats = ((ThriftHiveMetastore) thriftHiveMetastore).getStats(); } @AfterClass(alwaysRun = true) public void tearDown() + throws Exception { executor.shutdownNow(); executor = null; metastore = null; + closer.close(); + closer = null; } private ThriftMetastore createThriftHiveMetastore() @@ -178,11 +156,11 @@ private ThriftMetastore createThriftHiveMetastore() return createThriftHiveMetastore(mockClient); } - private static ThriftMetastore createThriftHiveMetastore(ThriftMetastoreClient client) + private ThriftMetastore createThriftHiveMetastore(ThriftMetastoreClient client) { return testingThriftHiveMetastoreBuilder() .metastoreClient(client) - .build(); + .build(closer::register); } @Test @@ -193,7 +171,7 @@ public void testCachingWithOnlyPartitionsCacheEnabled() .usesCache(); assertThatCachingWithDisabledPartitionCache() - .whenExecuting(testedMetastore -> testedMetastore.getAllTables(TEST_DATABASE)) + .whenExecuting(testedMetastore -> testedMetastore.getTables(TEST_DATABASE)) .usesCache(); assertThatCachingWithDisabledPartitionCache() @@ -241,47 +219,26 @@ public void testGetAllDatabases() @Test public void testGetAllTable() { - assertEquals(mockClient.getAccessCount(), 0); - assertEquals(metastore.getAllTables(TEST_DATABASE), ImmutableList.of(TEST_TABLE)); - assertEquals(mockClient.getAccessCount(), 1); - assertEquals(metastore.getAllTables(TEST_DATABASE), ImmutableList.of(TEST_TABLE)); - assertEquals(mockClient.getAccessCount(), 1); - assertEquals(metastore.getTableNamesStats().getRequestCount(), 2); - assertEquals(metastore.getTableNamesStats().getHitRate(), 0.5); - - metastore.flushCache(); - - assertEquals(metastore.getAllTables(TEST_DATABASE), ImmutableList.of(TEST_TABLE)); - assertEquals(mockClient.getAccessCount(), 2); - assertEquals(metastore.getTableNamesStats().getRequestCount(), 3); - assertEquals(metastore.getTableNamesStats().getHitRate(), 1.0 / 3); - } - - @Test - public void testBatchGetAllTable() - { - assertEquals(mockClient.getAccessCount(), 0); - assertEquals(metastore.getAllTables(), Optional.of(ImmutableList.of(TEST_SCHEMA_TABLE))); - assertEquals(mockClient.getAccessCount(), 1); - assertEquals(metastore.getAllTables(), Optional.of(ImmutableList.of(TEST_SCHEMA_TABLE))); - assertEquals(mockClient.getAccessCount(), 1); - assertEquals(metastore.getAllTables(TEST_DATABASE), ImmutableList.of(TEST_TABLE)); - assertEquals(mockClient.getAccessCount(), 2); - assertEquals(metastore.getAllTableNamesStats().getRequestCount(), 2); - assertEquals(metastore.getAllTableNamesStats().getHitRate(), .5); + assertThat(mockClient.getAccessCount()).isEqualTo(0); + assertThat(metastore.getTables(TEST_DATABASE)).isEqualTo(ImmutableList.of(TEST_TABLE_INFO)); + assertThat(mockClient.getAccessCount()).isEqualTo(1); + assertThat(metastore.getTables(TEST_DATABASE)).isEqualTo(ImmutableList.of(TEST_TABLE_INFO)); + assertThat(mockClient.getAccessCount()).isEqualTo(1); + assertThat(metastore.getTableNamesStats().getRequestCount()).isEqualTo(2); + assertThat(metastore.getTableNamesStats().getHitRate()).isEqualTo(0.5); metastore.flushCache(); - assertEquals(metastore.getAllTables(), Optional.of(ImmutableList.of(TEST_SCHEMA_TABLE))); - assertEquals(mockClient.getAccessCount(), 3); - assertEquals(metastore.getAllTableNamesStats().getRequestCount(), 3); - assertEquals(metastore.getAllTableNamesStats().getHitRate(), 1. / 3); + assertThat(metastore.getTables(TEST_DATABASE)).isEqualTo(ImmutableList.of(TEST_TABLE_INFO)); + assertThat(mockClient.getAccessCount()).isEqualTo(2); + assertThat(metastore.getTableNamesStats().getRequestCount()).isEqualTo(3); + assertThat(metastore.getTableNamesStats().getHitRate()).isEqualTo(1.0 / 3); } @Test public void testInvalidDbGetAllTAbles() { - assertTrue(metastore.getAllTables(BAD_DATABASE).isEmpty()); + assertThat(metastore.getTables(BAD_DATABASE)).isEmpty(); } @Test @@ -538,80 +495,73 @@ public void testListRoles() @Test public void testGetTableStatistics() { - assertEquals(mockClient.getAccessCount(), 0); + assertThat(mockClient.getAccessCount()).isEqualTo(0); - Table table = metastore.getTable(TEST_DATABASE, TEST_TABLE).orElseThrow(); - assertEquals(mockClient.getAccessCount(), 1); + assertThat(metastore.getTable(TEST_DATABASE, TEST_TABLE)).isPresent(); + assertThat(mockClient.getAccessCount()).isEqualTo(1); - assertEquals(metastore.getTableStatistics(table), TEST_STATS); - assertEquals(mockClient.getAccessCount(), 2); + assertThat(metastore.getTableColumnStatistics(TEST_DATABASE, TEST_TABLE, TEST_COLUMN_STATS.keySet())).isEqualTo(TEST_COLUMN_STATS); + assertThat(mockClient.getAccessCount()).isEqualTo(2); - assertEquals(metastore.getTableStatistics(table), TEST_STATS); - assertEquals(mockClient.getAccessCount(), 2); + assertThat(metastore.getTableColumnStatistics(TEST_DATABASE, TEST_TABLE, TEST_COLUMN_STATS.keySet())).isEqualTo(TEST_COLUMN_STATS); + assertThat(mockClient.getAccessCount()).isEqualTo(2); - assertEquals(metastore.getTableStatisticsStats().getRequestCount(), 2); - assertEquals(metastore.getTableStatisticsStats().getHitRate(), 0.5); + assertThat(metastore.getTableColumnStatisticsStats().getRequestCount()).isEqualTo(2); + assertThat(metastore.getTableColumnStatisticsStats().getHitRate()).isEqualTo(0.5); - assertEquals(metastore.getTableStats().getRequestCount(), 1); - assertEquals(metastore.getTableStats().getHitRate(), 0.0); + assertThat(metastore.getTableStats().getRequestCount()).isEqualTo(1); + assertThat(metastore.getTableStats().getHitRate()).isEqualTo(0.0); // check empty column list does not trigger the call - Table emptyColumnListTable = Table.builder(table).setDataColumns(ImmutableList.of()).build(); - assertThat(metastore.getTableStatistics(emptyColumnListTable).getBasicStatistics()).isEqualTo(TEST_STATS.getBasicStatistics()); - assertEquals(metastore.getTableStatisticsStats().getRequestCount(), 3); - assertEquals(metastore.getTableStatisticsStats().getHitRate(), 2.0 / 3); + assertThatThrownBy(() -> metastore.getTableColumnStatistics(TEST_DATABASE, TEST_TABLE, ImmutableSet.of())) + .isInstanceOf(IllegalArgumentException.class); + assertThat(metastore.getTableColumnStatisticsStats().getRequestCount()).isEqualTo(2); + assertThat(metastore.getTableColumnStatisticsStats().getHitRate()).isEqualTo(0.5); mockClient.mockColumnStats(TEST_DATABASE, TEST_TABLE, ImmutableMap.of( "col1", ColumnStatisticsData.longStats(new LongColumnStatsData().setNumNulls(1)), "col2", ColumnStatisticsData.longStats(new LongColumnStatsData().setNumNulls(2)), "col3", ColumnStatisticsData.longStats(new LongColumnStatsData().setNumNulls(3)))); - Table tableCol1 = Table.builder(table).setDataColumns(ImmutableList.of(new Column("col1", HIVE_LONG, Optional.empty()))).build(); - assertThat(metastore.getTableStatistics(tableCol1).getColumnStatistics()).containsEntry("col1", intColumnStats(1)); - Table tableCol2 = Table.builder(table).setDataColumns(ImmutableList.of(new Column("col2", HIVE_LONG, Optional.empty()))).build(); - assertThat(metastore.getTableStatistics(tableCol2).getColumnStatistics()).containsEntry("col2", intColumnStats(2)); - Table tableCol23 = Table.builder(table) - .setDataColumns(ImmutableList.of(new Column("col2", HIVE_LONG, Optional.empty()), new Column("col3", HIVE_LONG, Optional.empty()))) - .build(); - assertThat(metastore.getTableStatistics(tableCol23).getColumnStatistics()) + assertThat(metastore.getTableColumnStatistics(TEST_DATABASE, TEST_TABLE, ImmutableSet.of("col1"))).containsEntry("col1", intColumnStats(1)); + assertThat(metastore.getTableColumnStatistics(TEST_DATABASE, TEST_TABLE, ImmutableSet.of("col2"))).containsEntry("col2", intColumnStats(2)); + assertThat(metastore.getTableColumnStatistics(TEST_DATABASE, TEST_TABLE, ImmutableSet.of("col2", "col3"))) .containsEntry("col2", intColumnStats(2)) .containsEntry("col3", intColumnStats(3)); - metastore.getTableStatistics(table); // ensure cached - assertEquals(mockClient.getAccessCount(), 5); + metastore.getTableColumnStatistics(TEST_DATABASE, TEST_TABLE, ImmutableSet.of(TEST_COLUMN)); // ensure cached + assertThat(mockClient.getAccessCount()).isEqualTo(5); ColumnStatisticsData newStats = new ColumnStatisticsData(); newStats.setLongStats(new LongColumnStatsData(327843, 4324)); mockClient.mockColumnStats(TEST_DATABASE, TEST_TABLE, ImmutableMap.of(TEST_COLUMN, newStats)); metastore.invalidateTable(TEST_DATABASE, TEST_TABLE); - assertEquals(metastore.getTableStatistics(table), PartitionStatistics.builder() - .setBasicStatistics(TEST_STATS.getBasicStatistics()) - .setColumnStatistics(ImmutableMap.of(TEST_COLUMN, createIntegerColumnStatistics( + assertThat(metastore.getTableColumnStatistics(TEST_DATABASE, TEST_TABLE, ImmutableSet.of(TEST_COLUMN))) + .isEqualTo(ImmutableMap.of(TEST_COLUMN, createIntegerColumnStatistics( OptionalLong.empty(), OptionalLong.empty(), OptionalLong.of(newStats.getLongStats().getNumNulls()), - OptionalLong.of(newStats.getLongStats().getNumDVs() - 1)))) - .build()); - assertEquals(mockClient.getAccessCount(), 6); + OptionalLong.of(newStats.getLongStats().getNumDVs())))); + assertThat(mockClient.getAccessCount()).isEqualTo(6); } @Test public void testGetTableStatisticsWithoutMetadataCache() { - assertEquals(mockClient.getAccessCount(), 0); + assertThat(mockClient.getAccessCount()).isEqualTo(0); - Table table = statsOnlyCacheMetastore.getTable(TEST_DATABASE, TEST_TABLE).orElseThrow(); - assertEquals(mockClient.getAccessCount(), 1); + assertThat(statsOnlyCacheMetastore.getTable(TEST_DATABASE, TEST_TABLE)).isPresent(); + assertThat(mockClient.getAccessCount()).isEqualTo(1); - assertEquals(statsOnlyCacheMetastore.getTableStatistics(table), TEST_STATS); - assertEquals(mockClient.getAccessCount(), 2); + assertThat(statsOnlyCacheMetastore.getTableColumnStatistics(TEST_DATABASE, TEST_TABLE, ImmutableSet.of(TEST_COLUMN))).isEqualTo(TEST_COLUMN_STATS); + assertThat(mockClient.getAccessCount()).isEqualTo(2); - assertEquals(statsOnlyCacheMetastore.getTableStatistics(table), TEST_STATS); - assertEquals(mockClient.getAccessCount(), 2); + assertThat(statsOnlyCacheMetastore.getTableColumnStatistics(TEST_DATABASE, TEST_TABLE, ImmutableSet.of(TEST_COLUMN))).isEqualTo(TEST_COLUMN_STATS); + assertThat(mockClient.getAccessCount()).isEqualTo(2); - assertEquals(statsOnlyCacheMetastore.getTableStatisticsStats().getRequestCount(), 2); - assertEquals(statsOnlyCacheMetastore.getTableStatisticsStats().getHitRate(), 0.5); + assertThat(statsOnlyCacheMetastore.getTableColumnStatisticsStats().getRequestCount()).isEqualTo(2); + assertThat(statsOnlyCacheMetastore.getTableColumnStatisticsStats().getHitRate()).isEqualTo(0.5); - assertEquals(statsOnlyCacheMetastore.getTableStats().getRequestCount(), 0); - assertEquals(statsOnlyCacheMetastore.getTableStats().getHitRate(), 1.0); + assertThat(statsOnlyCacheMetastore.getTableStats().getRequestCount()).isEqualTo(0); + assertThat(statsOnlyCacheMetastore.getTableStats().getHitRate()).isEqualTo(1.0); } @Test @@ -638,11 +588,11 @@ public List getTableColumnStatistics(String databaseName, S return result; } }; - CachingHiveMetastore metastore = createMetastore(mockClient); + CachingHiveMetastore metastore = createCachingHiveMetastore(new BridgingHiveMetastore(createThriftHiveMetastore(mockClient)), CACHE_TTL, true, true, executor); - Table table = metastore.getTable(TEST_DATABASE, TEST_TABLE).orElseThrow(); + assertThat(metastore.getTable(TEST_DATABASE, TEST_TABLE)).isPresent(); - ExecutorService executorService = Executors.newFixedThreadPool(1, new ThreadFactoryBuilder().setNameFormat("invalidation-%d").build()); + ExecutorService executorService = Executors.newFixedThreadPool(1, daemonThreadsNamed("invalidation-%d")); try { // invalidate thread Future invalidateFuture = executorService.submit( @@ -658,12 +608,12 @@ public List getTableColumnStatistics(String databaseName, S }); // start get stats before the invalidation, it will wait until invalidation is done to finish - assertEquals(metastore.getTableStatistics(table), TEST_STATS); - assertEquals(mockClient.getAccessCount(), 2); + assertThat(metastore.getTableColumnStatistics(TEST_DATABASE, TEST_TABLE, TEST_COLUMN_STATS.keySet())).isEqualTo(TEST_COLUMN_STATS); + assertThat(mockClient.getAccessCount()).isEqualTo(2); // get stats after invalidate - assertEquals(metastore.getTableStatistics(table), TEST_STATS); + assertThat(metastore.getTableColumnStatistics(TEST_DATABASE, TEST_TABLE, TEST_COLUMN_STATS.keySet())).isEqualTo(TEST_COLUMN_STATS); // the value was not cached - assertEquals(mockClient.getAccessCount(), 3); + assertThat(mockClient.getAccessCount()).isEqualTo(3); // make sure invalidateFuture is done invalidateFuture.get(1, SECONDS); } @@ -675,10 +625,10 @@ public List getTableColumnStatistics(String databaseName, S @Test public void testGetPartitionStatistics() { - assertEquals(mockClient.getAccessCount(), 0); + assertThat(mockClient.getAccessCount()).isEqualTo(0); Table table = metastore.getTable(TEST_DATABASE, TEST_TABLE).orElseThrow(); - assertEquals(mockClient.getAccessCount(), 1); + assertThat(mockClient.getAccessCount()).isEqualTo(1); Partition partition = metastore.getPartition(table, TEST_PARTITION_VALUES1).orElseThrow(); String partitionName = makePartitionName(table, partition); @@ -686,51 +636,46 @@ public void testGetPartitionStatistics() String partition2Name = makePartitionName(table, partition2); Partition partition3 = metastore.getPartition(table, TEST_PARTITION_VALUES3).orElseThrow(); String partition3Name = makePartitionName(table, partition3); - assertEquals(mockClient.getAccessCount(), 4); + assertThat(mockClient.getAccessCount()).isEqualTo(4); - assertEquals(metastore.getPartitionStatistics(table, ImmutableList.of(partition)), ImmutableMap.of(TEST_PARTITION1, TEST_STATS)); - assertEquals(mockClient.getAccessCount(), 5); + assertThat(metastore.getPartitionColumnStatistics(TEST_DATABASE, TEST_TABLE, ImmutableSet.of(TEST_PARTITION1), ImmutableSet.of(TEST_COLUMN))) + .isEqualTo(ImmutableMap.of(TEST_PARTITION1, TEST_COLUMN_STATS)); + assertThat(mockClient.getAccessCount()).isEqualTo(5); - assertEquals(metastore.getPartitionStatisticsStats().getRequestCount(), 1); - assertEquals(metastore.getPartitionStatistics(table, ImmutableList.of(partition)), ImmutableMap.of(TEST_PARTITION1, TEST_STATS)); - assertEquals(mockClient.getAccessCount(), 5); + assertThat(metastore.getPartitionStatisticsStats().getRequestCount()).isEqualTo(1); + assertThat(metastore.getPartitionColumnStatistics(TEST_DATABASE, TEST_TABLE, ImmutableSet.of(TEST_PARTITION1), ImmutableSet.of(TEST_COLUMN))) + .isEqualTo(ImmutableMap.of(TEST_PARTITION1, TEST_COLUMN_STATS)); + assertThat(mockClient.getAccessCount()).isEqualTo(5); - assertEquals(metastore.getPartitionStatisticsStats().getRequestCount(), 2); - assertEquals(metastore.getPartitionStatisticsStats().getHitRate(), 1.0 / 2); + assertThat(metastore.getPartitionStatisticsStats().getRequestCount()).isEqualTo(2); + assertThat(metastore.getPartitionStatisticsStats().getHitRate()).isEqualTo(1.0 / 2); - assertEquals(metastore.getTableStats().getRequestCount(), 1); - assertEquals(metastore.getTableStats().getHitRate(), 0.0); + assertThat(metastore.getTableStats().getRequestCount()).isEqualTo(1); + assertThat(metastore.getTableStats().getHitRate()).isEqualTo(0.0); - assertEquals(metastore.getPartitionStats().getRequestCount(), 3); - assertEquals(metastore.getPartitionStats().getHitRate(), 0.0); + assertThat(metastore.getPartitionStats().getRequestCount()).isEqualTo(3); + assertThat(metastore.getPartitionStats().getHitRate()).isEqualTo(0.0); // check empty column list does not trigger the call - Table emptyColumnListTable = Table.builder(table).setDataColumns(ImmutableList.of()).build(); - Map partitionStatistics = metastore.getPartitionStatistics(emptyColumnListTable, ImmutableList.of(partition)); - assertThat(partitionStatistics).containsOnlyKeys(TEST_PARTITION1); - assertThat(partitionStatistics.get(TEST_PARTITION1).getBasicStatistics()).isEqualTo(TEST_STATS.getBasicStatistics()); - assertEquals(metastore.getPartitionStatisticsStats().getRequestCount(), 3); - assertEquals(metastore.getPartitionStatisticsStats().getHitRate(), 2.0 / 3); + assertThatThrownBy(() -> metastore.getPartitionColumnStatistics(TEST_DATABASE, TEST_TABLE, ImmutableSet.of(TEST_PARTITION1), ImmutableSet.of())) + .isInstanceOf(IllegalArgumentException.class); + assertThat(metastore.getPartitionStatisticsStats().getRequestCount()).isEqualTo(2); + assertThat(metastore.getPartitionStatisticsStats().getHitRate()).isEqualTo(0.5); mockClient.mockPartitionColumnStats(TEST_DATABASE, TEST_TABLE, TEST_PARTITION1, ImmutableMap.of( "col1", ColumnStatisticsData.longStats(new LongColumnStatsData().setNumNulls(1)), "col2", ColumnStatisticsData.longStats(new LongColumnStatsData().setNumNulls(2)), "col3", ColumnStatisticsData.longStats(new LongColumnStatsData().setNumNulls(3)))); - Table tableCol1 = Table.builder(table).setDataColumns(ImmutableList.of(new Column("col1", HIVE_LONG, Optional.empty()))).build(); - Map tableCol1PartitionStatistics = metastore.getPartitionStatistics(tableCol1, ImmutableList.of(partition)); + var tableCol1PartitionStatistics = metastore.getPartitionColumnStatistics(TEST_DATABASE, TEST_TABLE, ImmutableSet.of(partitionName), ImmutableSet.of("col1")); assertThat(tableCol1PartitionStatistics).containsOnlyKeys(partitionName); - assertThat(tableCol1PartitionStatistics.get(partitionName).getColumnStatistics()).containsEntry("col1", intColumnStats(1)); - Table tableCol2 = Table.builder(table).setDataColumns(ImmutableList.of(new Column("col2", HIVE_LONG, Optional.empty()))).build(); - Map tableCol2PartitionStatistics = metastore.getPartitionStatistics(tableCol2, ImmutableList.of(partition)); + assertThat(tableCol1PartitionStatistics.get(partitionName)).containsEntry("col1", intColumnStats(1)); + var tableCol2PartitionStatistics = metastore.getPartitionColumnStatistics(TEST_DATABASE, TEST_TABLE, ImmutableSet.of(partitionName), ImmutableSet.of("col2")); assertThat(tableCol2PartitionStatistics).containsOnlyKeys(partitionName); - assertThat(tableCol2PartitionStatistics.get(partitionName).getColumnStatistics()).containsEntry("col2", intColumnStats(2)); - Table tableCol23 = Table.builder(table) - .setDataColumns(ImmutableList.of(new Column("col2", HIVE_LONG, Optional.empty()), new Column("col3", HIVE_LONG, Optional.empty()))) - .build(); - Map tableCol23PartitionStatistics = metastore.getPartitionStatistics(tableCol23, ImmutableList.of(partition)); + assertThat(tableCol2PartitionStatistics.get(partitionName)).containsEntry("col2", intColumnStats(2)); + var tableCol23PartitionStatistics = metastore.getPartitionColumnStatistics(TEST_DATABASE, TEST_TABLE, ImmutableSet.of(partitionName), ImmutableSet.of("col2", "col3")); assertThat(tableCol23PartitionStatistics).containsOnlyKeys(partitionName); - assertThat(tableCol23PartitionStatistics.get(partitionName).getColumnStatistics()) + assertThat(tableCol23PartitionStatistics.get(partitionName)) .containsEntry("col2", intColumnStats(2)) .containsEntry("col3", intColumnStats(3)); @@ -744,19 +689,23 @@ public void testGetPartitionStatistics() "col2", ColumnStatisticsData.longStats(new LongColumnStatsData().setNumNulls(32)), "col3", ColumnStatisticsData.longStats(new LongColumnStatsData().setNumNulls(33)))); - Map tableCol2Partition2Statistics = metastore.getPartitionStatistics(tableCol2, ImmutableList.of(partition2)); + var tableCol2Partition2Statistics = metastore.getPartitionColumnStatistics(TEST_DATABASE, TEST_TABLE, ImmutableSet.of(partition2Name), ImmutableSet.of("col2")); assertThat(tableCol2Partition2Statistics).containsOnlyKeys(partition2Name); - assertThat(tableCol2Partition2Statistics.get(partition2Name).getColumnStatistics()).containsEntry("col2", intColumnStats(22)); + assertThat(tableCol2Partition2Statistics.get(partition2Name)).containsEntry("col2", intColumnStats(22)); - Map tableCol23Partition123Statistics = metastore.getPartitionStatistics(tableCol23, ImmutableList.of(partition, partition2, partition3)); + var tableCol23Partition123Statistics = metastore.getPartitionColumnStatistics( + TEST_DATABASE, + TEST_TABLE, + ImmutableSet.of(partitionName, partition2Name, partition3Name), + ImmutableSet.of("col2", "col3")); assertThat(tableCol23Partition123Statistics).containsOnlyKeys(partitionName, partition2Name, partition3Name); - assertThat(tableCol23Partition123Statistics.get(partitionName).getColumnStatistics()) + assertThat(tableCol23Partition123Statistics.get(partitionName)) .containsEntry("col2", intColumnStats(2)) .containsEntry("col3", intColumnStats(3)); - assertThat(tableCol23Partition123Statistics.get(partition2Name).getColumnStatistics()) + assertThat(tableCol23Partition123Statistics.get(partition2Name)) .containsEntry("col2", intColumnStats(22)) .containsEntry("col3", intColumnStats(23)); - assertThat(tableCol23Partition123Statistics.get(partition3Name).getColumnStatistics()) + assertThat(tableCol23Partition123Statistics.get(partition3Name)) .containsEntry("col2", intColumnStats(32)) .containsEntry("col3", intColumnStats(33)); } @@ -764,28 +713,30 @@ public void testGetPartitionStatistics() @Test public void testGetPartitionStatisticsWithoutMetadataCache() { - assertEquals(mockClient.getAccessCount(), 0); + assertThat(mockClient.getAccessCount()).isEqualTo(0); Table table = statsOnlyCacheMetastore.getTable(TEST_DATABASE, TEST_TABLE).orElseThrow(); - assertEquals(mockClient.getAccessCount(), 1); + assertThat(mockClient.getAccessCount()).isEqualTo(1); - Partition partition = statsOnlyCacheMetastore.getPartition(table, TEST_PARTITION_VALUES1).orElseThrow(); - assertEquals(mockClient.getAccessCount(), 2); + assertThat(statsOnlyCacheMetastore.getPartition(table, TEST_PARTITION_VALUES1)).isPresent(); + assertThat(mockClient.getAccessCount()).isEqualTo(2); - assertEquals(statsOnlyCacheMetastore.getPartitionStatistics(table, ImmutableList.of(partition)), ImmutableMap.of(TEST_PARTITION1, TEST_STATS)); - assertEquals(mockClient.getAccessCount(), 3); + assertThat(statsOnlyCacheMetastore.getPartitionColumnStatistics(TEST_DATABASE, TEST_TABLE, ImmutableSet.of(TEST_PARTITION1), ImmutableSet.of(TEST_COLUMN))) + .isEqualTo(ImmutableMap.of(TEST_PARTITION1, TEST_COLUMN_STATS)); + assertThat(mockClient.getAccessCount()).isEqualTo(3); - assertEquals(statsOnlyCacheMetastore.getPartitionStatistics(table, ImmutableList.of(partition)), ImmutableMap.of(TEST_PARTITION1, TEST_STATS)); - assertEquals(mockClient.getAccessCount(), 3); + assertThat(statsOnlyCacheMetastore.getPartitionColumnStatistics(TEST_DATABASE, TEST_TABLE, ImmutableSet.of(TEST_PARTITION1), ImmutableSet.of(TEST_COLUMN))) + .isEqualTo(ImmutableMap.of(TEST_PARTITION1, TEST_COLUMN_STATS)); + assertThat(mockClient.getAccessCount()).isEqualTo(3); - assertEquals(statsOnlyCacheMetastore.getPartitionStatisticsStats().getRequestCount(), 2); - assertEquals(statsOnlyCacheMetastore.getPartitionStatisticsStats().getHitRate(), 1.0 / 2); + assertThat(statsOnlyCacheMetastore.getPartitionStatisticsStats().getRequestCount()).isEqualTo(2); + assertThat(statsOnlyCacheMetastore.getPartitionStatisticsStats().getHitRate()).isEqualTo(1.0 / 2); - assertEquals(statsOnlyCacheMetastore.getTableStats().getRequestCount(), 0); - assertEquals(statsOnlyCacheMetastore.getTableStats().getHitRate(), 1.0); + assertThat(statsOnlyCacheMetastore.getTableStats().getRequestCount()).isEqualTo(0); + assertThat(statsOnlyCacheMetastore.getTableStats().getHitRate()).isEqualTo(1.0); - assertEquals(statsOnlyCacheMetastore.getPartitionStats().getRequestCount(), 0); - assertEquals(statsOnlyCacheMetastore.getPartitionStats().getHitRate(), 1.0); + assertThat(statsOnlyCacheMetastore.getPartitionStats().getRequestCount()).isEqualTo(0); + assertThat(statsOnlyCacheMetastore.getPartitionStats().getHitRate()).isEqualTo(1.0); } @Test @@ -812,13 +763,13 @@ public Map> getPartitionColumnStatistics(Strin return result; } }; - CachingHiveMetastore metastore = createMetastore(mockClient); + CachingHiveMetastore metastore = createCachingHiveMetastore(new BridgingHiveMetastore(createThriftHiveMetastore(mockClient)), CACHE_TTL, true, true, executor); Table table = metastore.getTable(TEST_DATABASE, TEST_TABLE).orElseThrow(); - Partition partition = metastore.getPartition(table, TEST_PARTITION_VALUES1).orElseThrow(); + assertThat(metastore.getPartition(table, TEST_PARTITION_VALUES1)).isPresent(); - ExecutorService executorService = Executors.newFixedThreadPool(1, new ThreadFactoryBuilder().setNameFormat("invalidation-%d").build()); + ExecutorService executorService = Executors.newFixedThreadPool(1, daemonThreadsNamed("invalidation-%d")); try { // invalidate thread Future invalidateFuture = executorService.submit( @@ -834,12 +785,14 @@ public Map> getPartitionColumnStatistics(Strin }); // start get stats before the invalidation, it will wait until invalidation is done to finish - assertEquals(metastore.getPartitionStatistics(table, ImmutableList.of(partition)), ImmutableMap.of(TEST_PARTITION1, TEST_STATS)); - assertEquals(mockClient.getAccessCount(), 3); + assertThat(metastore.getPartitionColumnStatistics(TEST_DATABASE, TEST_TABLE, ImmutableSet.of(TEST_PARTITION1), ImmutableSet.of(TEST_COLUMN))) + .isEqualTo(ImmutableMap.of(TEST_PARTITION1, TEST_COLUMN_STATS)); + assertThat(mockClient.getAccessCount()).isEqualTo(3); // get stats after invalidate - assertEquals(metastore.getPartitionStatistics(table, ImmutableList.of(partition)), ImmutableMap.of(TEST_PARTITION1, TEST_STATS)); + assertThat(metastore.getPartitionColumnStatistics(TEST_DATABASE, TEST_TABLE, ImmutableSet.of(TEST_PARTITION1), ImmutableSet.of(TEST_COLUMN))) + .isEqualTo(ImmutableMap.of(TEST_PARTITION1, TEST_COLUMN_STATS)); // the value was not cached - assertEquals(mockClient.getAccessCount(), 4); + assertThat(mockClient.getAccessCount()).isEqualTo(4); // make sure invalidateFuture is done invalidateFuture.get(1, SECONDS); } @@ -848,33 +801,16 @@ public Map> getPartitionColumnStatistics(Strin } } - private CachingHiveMetastore createMetastore(MockThriftMetastoreClient mockClient) - { - return CachingHiveMetastore.builder() - .delegate(new BridgingHiveMetastore(createThriftHiveMetastore(mockClient))) - .executor(executor) - .metadataCacheEnabled(true) - .statsCacheEnabled(true) - .cacheTtl(new Duration(5, TimeUnit.MINUTES)) - .refreshInterval(new Duration(1, TimeUnit.MINUTES)) - .maximumSize(1000) - .cacheMissing(new CachingHiveMetastoreConfig().isCacheMissing()) - .partitionCacheEnabled(true) - .build(); - } - @Test public void testUpdatePartitionStatistics() { - assertEquals(mockClient.getAccessCount(), 0); - - HiveMetastoreClosure hiveMetastoreClosure = new HiveMetastoreClosure(metastore); + assertThat(mockClient.getAccessCount()).isEqualTo(0); - Table table = hiveMetastoreClosure.getTable(TEST_DATABASE, TEST_TABLE).orElseThrow(); - assertEquals(mockClient.getAccessCount(), 1); + Table table = metastore.getTable(TEST_DATABASE, TEST_TABLE).orElseThrow(); + assertThat(mockClient.getAccessCount()).isEqualTo(1); - hiveMetastoreClosure.updatePartitionStatistics(table.getDatabaseName(), table.getTableName(), TEST_PARTITION1, identity()); - assertEquals(mockClient.getAccessCount(), 5); + metastore.updatePartitionStatistics(table, MERGE_INCREMENTAL, Map.of(TEST_PARTITION1, TEST_STATS)); + assertThat(mockClient.getAccessCount()).isEqualTo(5); } @Test @@ -911,230 +847,51 @@ public void testNoCacheExceptions() @Test public void testNoCacheMissing() { - CachingHiveMetastore metastore = CachingHiveMetastore.builder(metastoreBuilder) - .cacheMissing(false) - .build(); + CachingHiveMetastore metastore = createCachingHiveMetastore(new BridgingHiveMetastore(thriftHiveMetastore), CACHE_TTL, false, true, executor); mockClient.setReturnTable(false); - assertEquals(mockClient.getAccessCount(), 0); + assertThat(mockClient.getAccessCount()).isEqualTo(0); // First access assertThat(metastore.getTable(TEST_DATABASE, TEST_TABLE)).isEmpty(); - assertEquals(mockClient.getAccessCount(), 1); + assertThat(mockClient.getAccessCount()).isEqualTo(1); // Second access, second load assertThat(metastore.getTable(TEST_DATABASE, TEST_TABLE)).isEmpty(); - assertEquals(mockClient.getAccessCount(), 2); + assertThat(mockClient.getAccessCount()).isEqualTo(2); // Table get be accessed once it exists mockClient.setReturnTable(true); assertThat(metastore.getTable(TEST_DATABASE, TEST_TABLE)).isPresent(); - assertEquals(mockClient.getAccessCount(), 3); + assertThat(mockClient.getAccessCount()).isEqualTo(3); // Table existence is cached mockClient.setReturnTable(true); assertThat(metastore.getTable(TEST_DATABASE, TEST_TABLE)).isPresent(); - assertEquals(mockClient.getAccessCount(), 3); + assertThat(mockClient.getAccessCount()).isEqualTo(3); // Table is returned even if no longer exists mockClient.setReturnTable(false); assertThat(metastore.getTable(TEST_DATABASE, TEST_TABLE)).isPresent(); - assertEquals(mockClient.getAccessCount(), 3); + assertThat(mockClient.getAccessCount()).isEqualTo(3); // After cache invalidation, table absence is apparent metastore.invalidateTable(TEST_DATABASE, TEST_TABLE); assertThat(metastore.getTable(TEST_DATABASE, TEST_TABLE)).isEmpty(); - assertEquals(mockClient.getAccessCount(), 4); - } - - @Test - public void testCachingHiveMetastoreCreationWithTtlOnly() - { - CachingHiveMetastoreConfig config = new CachingHiveMetastoreConfig(); - config.setMetastoreCacheTtl(new Duration(10, TimeUnit.MILLISECONDS)); - - CachingHiveMetastore metastore = createMetastoreWithDirectExecutor(config); - - assertThat(metastore).isNotNull(); + assertThat(mockClient.getAccessCount()).isEqualTo(4); } @Test public void testCachingHiveMetastoreCreationViaMemoize() { - ThriftMetastore thriftHiveMetastore = createThriftHiveMetastore(); - metastore = memoizeMetastore( - new BridgingHiveMetastore(thriftHiveMetastore), - 1000); - - assertEquals(mockClient.getAccessCount(), 0); - assertEquals(metastore.getAllDatabases(), ImmutableList.of(TEST_DATABASE)); - assertEquals(mockClient.getAccessCount(), 1); - assertEquals(metastore.getAllDatabases(), ImmutableList.of(TEST_DATABASE)); - assertEquals(mockClient.getAccessCount(), 1); - assertEquals(metastore.getDatabaseNamesStats().getRequestCount(), 0); - } - - @Test(timeOut = 60_000, dataProviderClass = DataProviders.class, dataProvider = "trueFalse") - public void testLoadAfterInvalidate(boolean invalidateAll) - throws Exception - { - // State - CopyOnWriteArrayList tableColumns = new CopyOnWriteArrayList<>(); - ConcurrentMap tablePartitionsByName = new ConcurrentHashMap<>(); - Map tableParameters = new ConcurrentHashMap<>(); - tableParameters.put("frequent-changing-table-parameter", "parameter initial value"); - - // Initialize data - String databaseName = "my_database"; - String tableName = "my_table_name"; - - tableColumns.add(new Column("value", toHiveType(VARCHAR), Optional.empty() /* comment */)); - tableColumns.add(new Column("pk", toHiveType(VARCHAR), Optional.empty() /* comment */)); - - List partitionNames = new ArrayList<>(); - for (int i = 0; i < 10; i++) { - String partitionName = "pk=" + i; - tablePartitionsByName.put( - partitionName, - Partition.builder() - .setDatabaseName(databaseName) - .setTableName(tableName) - .setColumns(ImmutableList.copyOf(tableColumns)) - .setValues(List.of(Integer.toString(i))) - .withStorage(storage -> storage.setStorageFormat(fromHiveStorageFormat(TEXTFILE))) - .setParameters(Map.of("frequent-changing-partition-parameter", "parameter initial value")) - .build()); - partitionNames.add(partitionName); - } - - // Mock metastore - CountDownLatch getTableEnteredLatch = new CountDownLatch(1); - CountDownLatch getTableReturnLatch = new CountDownLatch(1); - CountDownLatch getTableFinishedLatch = new CountDownLatch(1); - CountDownLatch getPartitionsByNamesEnteredLatch = new CountDownLatch(1); - CountDownLatch getPartitionsByNamesReturnLatch = new CountDownLatch(1); - CountDownLatch getPartitionsByNamesFinishedLatch = new CountDownLatch(1); - - HiveMetastore mockMetastore = new UnimplementedHiveMetastore() - { - @Override - public Optional
getTable(String databaseName, String tableName) - { - Optional
table = Optional.of(Table.builder() - .setDatabaseName(databaseName) - .setTableName(tableName) - .setTableType(EXTERNAL_TABLE.name()) - .setDataColumns(tableColumns) - .setParameters(ImmutableMap.copyOf(tableParameters)) - // Required by 'Table', but not used by view translation. - .withStorage(storage -> storage.setStorageFormat(fromHiveStorageFormat(TEXTFILE))) - .setOwner(Optional.empty()) - .build()); - - getTableEnteredLatch.countDown(); // 1 - await(getTableReturnLatch, 10, SECONDS); // 2 - - return table; - } - - @Override - public Map> getPartitionsByNames(Table table, List partitionNames) - { - Map> result = new HashMap<>(); - for (String partitionName : partitionNames) { - result.put(partitionName, Optional.ofNullable(tablePartitionsByName.get(partitionName))); - } - - getPartitionsByNamesEnteredLatch.countDown(); // loader#1 - await(getPartitionsByNamesReturnLatch, 10, SECONDS); // loader#2 - - return result; - } - }; - - // Caching metastore - metastore = CachingHiveMetastore.builder() - .delegate(mockMetastore) - .executor(executor) - .metadataCacheEnabled(true) - .statsCacheEnabled(true) - .cacheTtl(new Duration(5, TimeUnit.MINUTES)) - .refreshInterval(new Duration(1, TimeUnit.MINUTES)) - .maximumSize(1000) - .cacheMissing(new CachingHiveMetastoreConfig().isCacheMissing()) - .partitionCacheEnabled(true) - .build(); - - // The test. Main thread does modifications and verifies subsequent load sees them. Background thread loads the state into the cache. - ExecutorService executor = Executors.newFixedThreadPool(1); - try { - Future future = executor.submit(() -> { - try { - Table table; - - table = metastore.getTable(databaseName, tableName).orElseThrow(); - getTableFinishedLatch.countDown(); // 3 + metastore = createPerTransactionCache(new BridgingHiveMetastore(createThriftHiveMetastore()), 1000); - metastore.getPartitionsByNames(table, partitionNames); - getPartitionsByNamesFinishedLatch.countDown(); // 6 - - return null; - } - catch (Throwable e) { - log.error(e); - throw e; - } - }); - - await(getTableEnteredLatch, 10, SECONDS); // 21 - tableParameters.put("frequent-changing-table-parameter", "main-thread-put-xyz"); - if (invalidateAll) { - metastore.flushCache(); - } - else { - metastore.invalidateTable(databaseName, tableName); - } - getTableReturnLatch.countDown(); // 2 - await(getTableFinishedLatch, 10, SECONDS); // 3 - Table table = metastore.getTable(databaseName, tableName).orElseThrow(); - assertThat(table.getParameters()) - .isEqualTo(Map.of("frequent-changing-table-parameter", "main-thread-put-xyz")); - - await(getPartitionsByNamesEnteredLatch, 10, SECONDS); // 4 - String partitionName = partitionNames.get(2); - Map newPartitionParameters = Map.of("frequent-changing-partition-parameter", "main-thread-put-alice"); - tablePartitionsByName.put(partitionName, - Partition.builder(tablePartitionsByName.get(partitionName)) - .setParameters(newPartitionParameters) - .build()); - if (invalidateAll) { - metastore.flushCache(); - } - else { - metastore.invalidateTable(databaseName, tableName); - } - getPartitionsByNamesReturnLatch.countDown(); // 5 - await(getPartitionsByNamesFinishedLatch, 10, SECONDS); // 6 - Map> loadedPartitions = metastore.getPartitionsByNames(table, partitionNames); - assertThat(loadedPartitions.get(partitionName)) - .isNotNull() - .isPresent() - .hasValueSatisfying(partition -> assertThat(partition.getParameters()).isEqualTo(newPartitionParameters)); - - // verify no failure in the background thread - future.get(10, SECONDS); - } - finally { - getTableEnteredLatch.countDown(); - getTableReturnLatch.countDown(); - getTableFinishedLatch.countDown(); - getPartitionsByNamesEnteredLatch.countDown(); - getPartitionsByNamesReturnLatch.countDown(); - getPartitionsByNamesFinishedLatch.countDown(); - - executor.shutdownNow(); - assertTrue(executor.awaitTermination(10, SECONDS)); - } + assertThat(mockClient.getAccessCount()).isEqualTo(0); + assertThat(metastore.getAllDatabases()).isEqualTo(ImmutableList.of(TEST_DATABASE)); + assertThat(mockClient.getAccessCount()).isEqualTo(1); + assertThat(metastore.getAllDatabases()).isEqualTo(ImmutableList.of(TEST_DATABASE)); + assertThat(mockClient.getAccessCount()).isEqualTo(1); + assertThat(metastore.getDatabaseNamesStats().getRequestCount()).isEqualTo(0); } @Test @@ -1193,40 +950,6 @@ public void testAllDatabases() assertThat(mockClient.getAccessCount()).isEqualTo(3); // should read it from cache } - @Test - public void testAllTables() - { - assertThat(mockClient.getAccessCount()).isEqualTo(0); - - assertThat(metastore.getAllTables()).contains(ImmutableList.of(TEST_SCHEMA_TABLE)); - assertThat(mockClient.getAccessCount()).isEqualTo(1); - assertThat(metastore.getAllTables()).contains(ImmutableList.of(TEST_SCHEMA_TABLE)); - assertThat(mockClient.getAccessCount()).isEqualTo(1); // should read it from cache - - metastore.dropTable(TEST_DATABASE, TEST_TABLE, false); - assertThat(mockClient.getAccessCount()).isEqualTo(2); // dropTable check if the table exists - - assertThat(metastore.getAllTables()).contains(ImmutableList.of(TEST_SCHEMA_TABLE)); - assertThat(mockClient.getAccessCount()).isEqualTo(3); - assertThat(metastore.getAllTables()).contains(ImmutableList.of(TEST_SCHEMA_TABLE)); - assertThat(mockClient.getAccessCount()).isEqualTo(3); // should read it from cache - - metastore.createTable( - Table.builder() - .setDatabaseName(TEST_DATABASE) - .setTableName(TEST_TABLE) - .setOwner(Optional.empty()) - .setTableType(VIRTUAL_VIEW.name()) - .withStorage(storage -> storage.setStorageFormat(VIEW_STORAGE_FORMAT)) - .build(), - new PrincipalPrivileges(ImmutableMultimap.of(), ImmutableMultimap.of())); - - assertThat(metastore.getAllTables()).contains(ImmutableList.of(TEST_SCHEMA_TABLE)); - assertThat(mockClient.getAccessCount()).isEqualTo(4); - assertThat(metastore.getAllTables()).contains(ImmutableList.of(TEST_SCHEMA_TABLE)); - assertThat(mockClient.getAccessCount()).isEqualTo(4); // should read it from cache - } - private static HiveColumnStatistics intColumnStats(int nullsCount) { return createIntegerColumnStatistics(OptionalLong.empty(), OptionalLong.empty(), OptionalLong.of(nullsCount), OptionalLong.empty()); @@ -1244,31 +967,21 @@ private static void await(CountDownLatch latch, long timeout, TimeUnit unit) } } - static class PartitionCachingAssertions + private PartitionCachingAssertions assertThatCachingWithDisabledPartitionCache() + { + return new PartitionCachingAssertions(executor); + } + + class PartitionCachingAssertions { private final CachingHiveMetastore cachingHiveMetastore; private final MockThriftMetastoreClient thriftClient; private Consumer metastoreInteractions = hiveMetastore -> {}; - static PartitionCachingAssertions assertThatCachingWithDisabledPartitionCache() - { - return new PartitionCachingAssertions(); - } - - private PartitionCachingAssertions() + private PartitionCachingAssertions(Executor refreshExecutor) { thriftClient = new MockThriftMetastoreClient(); - cachingHiveMetastore = CachingHiveMetastore.builder() - .delegate(new BridgingHiveMetastore(createThriftHiveMetastore(thriftClient))) - .executor(listeningDecorator(newCachedThreadPool(daemonThreadsNamed("test-%s")))) - .metadataCacheEnabled(true) - .statsCacheEnabled(true) - .cacheTtl(new Duration(5, TimeUnit.MINUTES)) - .refreshInterval(new Duration(1, TimeUnit.MINUTES)) - .maximumSize(1000) - .cacheMissing(true) - .partitionCacheEnabled(false) - .build(); + cachingHiveMetastore = createCachingHiveMetastore(new BridgingHiveMetastore(createThriftHiveMetastore(thriftClient)), CACHE_TTL, true, false, refreshExecutor); } PartitionCachingAssertions whenExecuting(Consumer interactions) @@ -1309,18 +1022,17 @@ void omitsCacheForNumberOfOperations(int expectedCacheOmittingOperations) } } - private CachingHiveMetastore createMetastoreWithDirectExecutor(CachingHiveMetastoreConfig config) + private static CachingHiveMetastore createCachingHiveMetastore(HiveMetastore hiveMetastore, Duration cacheTtl, boolean cacheMissing, boolean partitionCacheEnabled, Executor executor) { - return CachingHiveMetastore.builder() - .delegate(new BridgingHiveMetastore(createThriftHiveMetastore())) - .executor(directExecutor()) - .metadataCacheEnabled(true) - .statsCacheEnabled(true) - .cacheTtl(config.getMetastoreCacheTtl()) - .refreshInterval(config.getMetastoreRefreshInterval()) - .maximumSize(config.getMetastoreCacheMaximumSize()) - .cacheMissing(config.isCacheMissing()) - .partitionCacheEnabled(config.isPartitionCacheEnabled()) - .build(); + return CachingHiveMetastore.createCachingHiveMetastore( + hiveMetastore, + cacheTtl, + CACHE_TTL, + Optional.of(new Duration(1, TimeUnit.MINUTES)), + executor, + 1000, + CachingHiveMetastore.StatsRecording.ENABLED, + partitionCacheEnabled, + cacheMissing ? ImmutableSet.copyOf(CachingHiveMetastore.ObjectType.values()) : ImmutableSet.of()); } } diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/cache/TestCachingHiveMetastoreConfig.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/cache/TestCachingHiveMetastoreConfig.java index cdad25ced405..e1cc0ac55b8b 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/cache/TestCachingHiveMetastoreConfig.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/cache/TestCachingHiveMetastoreConfig.java @@ -36,7 +36,9 @@ public void testDefaults() .setMetastoreCacheMaximumSize(10000) .setMaxMetastoreRefreshThreads(10) .setCacheMissing(true) - .setPartitionCacheEnabled(true)); + .setPartitionCacheEnabled(true) + .setCacheMissingPartitions(true) + .setCacheMissingStats(true)); } @Test @@ -50,6 +52,8 @@ public void testExplicitPropertyMappings() .put("hive.metastore-refresh-max-threads", "2500") .put("hive.metastore-cache.cache-partitions", "false") .put("hive.metastore-cache.cache-missing", "false") + .put("hive.metastore-cache.cache-missing-partitions", "false") + .put("hive.metastore-cache.cache-missing-stats", "false") .buildOrThrow(); CachingHiveMetastoreConfig expected = new CachingHiveMetastoreConfig() @@ -59,7 +63,9 @@ public void testExplicitPropertyMappings() .setMetastoreCacheMaximumSize(5000) .setMaxMetastoreRefreshThreads(2500) .setCacheMissing(false) - .setPartitionCacheEnabled(false); + .setPartitionCacheEnabled(false) + .setCacheMissingPartitions(false) + .setCacheMissingStats(false); assertFullMapping(properties, expected); } diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/cache/TestCachingHiveMetastoreWithQueryRunner.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/cache/TestCachingHiveMetastoreWithQueryRunner.java index a5d520820803..1217688ade5d 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/cache/TestCachingHiveMetastoreWithQueryRunner.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/cache/TestCachingHiveMetastoreWithQueryRunner.java @@ -16,29 +16,27 @@ import com.google.common.base.Joiner; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; +import com.google.inject.Key; import io.trino.Session; import io.trino.plugin.hive.HiveQueryRunner; -import io.trino.plugin.hive.metastore.file.FileHiveMetastore; +import io.trino.plugin.hive.metastore.HiveMetastore; +import io.trino.plugin.hive.metastore.HiveMetastoreFactory; +import io.trino.plugin.hive.metastore.RawHiveMetastoreFactory; import io.trino.spi.security.Identity; import io.trino.spi.security.SelectedRole; import io.trino.testing.AbstractTestQueryFramework; -import io.trino.testing.DistributedQueryRunner; import io.trino.testing.QueryRunner; import org.testng.annotations.DataProvider; import org.testng.annotations.Test; -import java.nio.file.Path; import java.util.List; import java.util.Optional; import static com.google.common.base.Verify.verify; import static com.google.common.collect.Lists.cartesianProduct; -import static com.google.common.io.MoreFiles.deleteRecursively; -import static com.google.common.io.RecursiveDeleteOption.ALLOW_INSECURE; -import static io.trino.plugin.hive.metastore.file.TestingFileHiveMetastore.createTestingFileHiveMetastore; +import static io.trino.plugin.hive.TestingHiveUtils.getConnectorService; import static io.trino.spi.security.SelectedRole.Type.ROLE; import static io.trino.testing.TestingSession.testSessionBuilder; -import static java.nio.file.Files.createTempDirectory; import static java.util.Collections.nCopies; import static org.assertj.core.api.Assertions.assertThatThrownBy; @@ -54,28 +52,26 @@ public class TestCachingHiveMetastoreWithQueryRunner private static final String ALICE_NAME = "alice"; private static final Session ALICE = getTestSession(new Identity.Builder(ALICE_NAME).build()); - private FileHiveMetastore fileHiveMetastore; + private HiveMetastore fileHiveMetastore; @Override protected QueryRunner createQueryRunner() throws Exception { - Path temporaryMetastoreDirectory = createTempDirectory(null); - closeAfterClass(() -> deleteRecursively(temporaryMetastoreDirectory, ALLOW_INSECURE)); - - DistributedQueryRunner queryRunner = HiveQueryRunner.builder(ADMIN) - .setNodeCount(3) + QueryRunner queryRunner = HiveQueryRunner.builder(ADMIN) // Required by testPartitionAppend test. // Coordinator needs to be excluded from workers to deterministically reproduce the original problem // https://github.com/trinodb/trino/pull/6853 .setCoordinatorProperties(ImmutableMap.of("node-scheduler.include-coordinator", "false")) - .setMetastore(distributedQueryRunner -> fileHiveMetastore = createTestingFileHiveMetastore(temporaryMetastoreDirectory.toFile())) .setHiveProperties(ImmutableMap.of( "hive.security", "sql-standard", "hive.metastore-cache-ttl", "60m", "hive.metastore-refresh-interval", "10m")) .build(); + fileHiveMetastore = getConnectorService(queryRunner, Key.get(HiveMetastoreFactory.class, RawHiveMetastoreFactory.class)) + .createMetastore(Optional.empty()); + queryRunner.execute(ADMIN, "CREATE SCHEMA " + SCHEMA); queryRunner.execute("CREATE TABLE test (test INT)"); diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/file/TestFileHiveMetastore.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/file/TestFileHiveMetastore.java index b596ea33d4e6..d2aba891e04f 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/file/TestFileHiveMetastore.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/file/TestFileHiveMetastore.java @@ -14,10 +14,10 @@ package io.trino.plugin.hive.metastore.file; import com.google.common.collect.ImmutableMap; +import io.trino.filesystem.local.LocalFileSystemFactory; import io.trino.plugin.hive.NodeVersion; import io.trino.plugin.hive.metastore.Column; import io.trino.plugin.hive.metastore.Database; -import io.trino.plugin.hive.metastore.HiveMetastoreConfig; import io.trino.plugin.hive.metastore.StorageFormat; import io.trino.plugin.hive.metastore.Table; import org.apache.hadoop.hive.metastore.TableType; @@ -29,65 +29,66 @@ import java.io.IOException; import java.nio.file.Path; +import java.util.Map; import java.util.Optional; import static com.google.common.io.MoreFiles.deleteRecursively; import static com.google.common.io.RecursiveDeleteOption.ALLOW_INSECURE; -import static io.trino.plugin.hive.HiveTestUtils.HDFS_ENVIRONMENT; +import static io.trino.plugin.hive.HiveMetadata.TRINO_QUERY_ID_NAME; import static io.trino.plugin.hive.HiveType.HIVE_INT; import static io.trino.plugin.hive.metastore.PrincipalPrivileges.NO_PRIVILEGES; import static io.trino.plugin.hive.util.HiveClassNames.HUDI_PARQUET_INPUT_FORMAT; -import static io.trino.spi.security.PrincipalType.USER; import static io.trino.testing.TestingNames.randomNameSuffix; import static java.nio.file.Files.createTempDirectory; import static org.assertj.core.api.Assertions.assertThat; public class TestFileHiveMetastore { - private Path tmpDir; + private Path tempDir; private FileHiveMetastore metastore; @BeforeClass public void setUp() throws IOException { - tmpDir = createTempDirectory(getClass().getSimpleName()); + tempDir = createTempDirectory(getClass().getSimpleName()); + LocalFileSystemFactory fileSystemFactory = new LocalFileSystemFactory(tempDir); metastore = new FileHiveMetastore( new NodeVersion("testversion"), - HDFS_ENVIRONMENT, - new HiveMetastoreConfig().isHideDeltaLakeTables(), + fileSystemFactory, + false, new FileHiveMetastoreConfig() - .setCatalogDirectory(tmpDir.toString()) - .setDisableLocationChecks(true) - /*.setMetastoreUser("test")*/); - - metastore.createDatabase(Database.builder() - .setDatabaseName("default") - .setOwnerName(Optional.of("test")) - .setOwnerType(Optional.of(USER)) - .build()); + .setCatalogDirectory("local:///") + .setMetastoreUser("test") + .setDisableLocationChecks(true)); } @AfterClass(alwaysRun = true) public void tearDown() throws IOException { - deleteRecursively(tmpDir, ALLOW_INSECURE); - metastore = null; - tmpDir = null; + deleteRecursively(tempDir, ALLOW_INSECURE); } @Test public void testPreserveHudiInputFormat() { + String databaseName = "test_database_" + randomNameSuffix(); + Database.Builder database = Database.builder() + .setDatabaseName(databaseName) + .setParameters(Map.of(TRINO_QUERY_ID_NAME, "query_id")) + .setOwnerName(Optional.empty()) + .setOwnerType(Optional.empty()); + metastore.createDatabase(database.build()); + StorageFormat storageFormat = StorageFormat.create( ParquetHiveSerDe.class.getName(), HUDI_PARQUET_INPUT_FORMAT, MapredParquetOutputFormat.class.getName()); Table table = Table.builder() - .setDatabaseName("default") + .setDatabaseName(databaseName) .setTableName("some_table_name" + randomNameSuffix()) .setTableType(TableType.EXTERNAL_TABLE.name()) .setOwner(Optional.of("public")) diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/file/TestingFileHiveMetastore.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/file/TestingFileHiveMetastore.java index 9f677fddf382..0a5a5adda443 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/file/TestingFileHiveMetastore.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/file/TestingFileHiveMetastore.java @@ -13,33 +13,38 @@ */ package io.trino.plugin.hive.metastore.file; -import com.google.common.collect.ImmutableSet; -import io.trino.hdfs.DynamicHdfsConfiguration; -import io.trino.hdfs.HdfsConfig; -import io.trino.hdfs.HdfsConfiguration; -import io.trino.hdfs.HdfsConfigurationInitializer; -import io.trino.hdfs.HdfsEnvironment; -import io.trino.hdfs.authentication.NoHdfsAuthentication; +import io.trino.filesystem.Location; +import io.trino.filesystem.TrinoFileSystemFactory; +import io.trino.filesystem.local.LocalFileSystemFactory; import io.trino.plugin.hive.NodeVersion; import io.trino.plugin.hive.metastore.HiveMetastoreConfig; import java.io.File; +import static com.google.common.base.Verify.verify; + public final class TestingFileHiveMetastore { private TestingFileHiveMetastore() {} public static FileHiveMetastore createTestingFileHiveMetastore(File catalogDirectory) { - HdfsConfig hdfsConfig = new HdfsConfig(); - HdfsConfiguration hdfsConfiguration = new DynamicHdfsConfiguration(new HdfsConfigurationInitializer(hdfsConfig), ImmutableSet.of()); - HdfsEnvironment hdfsEnvironment = new HdfsEnvironment(hdfsConfiguration, hdfsConfig, new NoHdfsAuthentication()); + if (!catalogDirectory.exists()) { + verify(catalogDirectory.mkdirs()); + } + return createTestingFileHiveMetastore( + new LocalFileSystemFactory(catalogDirectory.toPath()), + Location.of("local:///")); + } + + public static FileHiveMetastore createTestingFileHiveMetastore(TrinoFileSystemFactory fileSystemFactory, Location catalogDirectory) + { return new FileHiveMetastore( new NodeVersion("testversion"), - hdfsEnvironment, + fileSystemFactory, new HiveMetastoreConfig().isHideDeltaLakeTables(), new FileHiveMetastoreConfig() - .setCatalogDirectory(catalogDirectory.toURI().toString()) + .setCatalogDirectory(catalogDirectory.toString()) .setMetastoreUser("test")); } } diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/glue/TestGlueConverter.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/glue/TestGlueConverter.java new file mode 100644 index 000000000000..2280e45ad3fd --- /dev/null +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/glue/TestGlueConverter.java @@ -0,0 +1,428 @@ +/* + * 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.hive.metastore.glue; + +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import io.trino.plugin.hive.HiveBucketProperty; +import io.trino.plugin.hive.TableType; +import io.trino.plugin.hive.metastore.Column; +import io.trino.plugin.hive.metastore.Database; +import io.trino.plugin.hive.metastore.Partition; +import io.trino.plugin.hive.metastore.Storage; +import io.trino.plugin.hive.metastore.StorageFormat; +import io.trino.plugin.hive.metastore.Table; +import io.trino.spi.security.PrincipalType; +import org.junit.jupiter.api.Test; +import software.amazon.awssdk.services.glue.model.DatabaseInput; +import software.amazon.awssdk.services.glue.model.PartitionInput; +import software.amazon.awssdk.services.glue.model.SerDeInfo; +import software.amazon.awssdk.services.glue.model.StorageDescriptor; +import software.amazon.awssdk.services.glue.model.TableInput; + +import java.util.List; +import java.util.Map; +import java.util.Optional; + +import static io.trino.plugin.hive.HiveType.HIVE_STRING; +import static io.trino.plugin.hive.TableType.EXTERNAL_TABLE; +import static io.trino.plugin.hive.ViewReaderUtil.PRESTO_VIEW_FLAG; +import static io.trino.plugin.hive.metastore.Table.TABLE_COMMENT; +import static io.trino.plugin.hive.metastore.TableInfo.ICEBERG_MATERIALIZED_VIEW_COMMENT; +import static io.trino.plugin.hive.metastore.glue.GlueConverter.PUBLIC_OWNER; +import static io.trino.plugin.hive.metastore.glue.GlueConverter.getTableTypeNullable; +import static io.trino.plugin.hive.util.HiveUtil.DELTA_LAKE_PROVIDER; +import static io.trino.plugin.hive.util.HiveUtil.ICEBERG_TABLE_TYPE_NAME; +import static io.trino.plugin.hive.util.HiveUtil.ICEBERG_TABLE_TYPE_VALUE; +import static io.trino.plugin.hive.util.HiveUtil.SPARK_TABLE_PROVIDER_KEY; +import static org.assertj.core.api.Assertions.assertThat; + +class TestGlueConverter +{ + private final Database trinoDatabase = Database.builder() + .setDatabaseName("test-database") + .setComment(Optional.of("database desc")) + .setLocation(Optional.of("/database")) + .setParameters(ImmutableMap.of()) + .setOwnerName(Optional.of("PUBLIC")) + .setOwnerType(Optional.of(PrincipalType.ROLE)) + .build(); + + private final Table trinoTable = Table.builder() + .setDatabaseName(trinoDatabase.getDatabaseName()) + .setTableName("test-table") + .setOwner(Optional.of("owner")) + .setParameters(ImmutableMap.of()) + .setTableType(EXTERNAL_TABLE.name()) + .setDataColumns(ImmutableList.of(new Column("table-data", HIVE_STRING, Optional.of("table data column comment"), Map.of()))) + .setPartitionColumns(ImmutableList.of(new Column("table-partition", HIVE_STRING, Optional.of("table partition column comment"), Map.of()))) + .setViewOriginalText(Optional.of("originalText")) + .setViewExpandedText(Optional.of("expandedText")) + .withStorage(storage -> storage.setStorageFormat(StorageFormat.create("TableSerdeLib", "TableInputFormat", "TableOutputFormat")) + .setLocation("/test-table") + .setBucketProperty(Optional.empty()) + .setSerdeParameters(ImmutableMap.of())).build(); + + private final Partition trinoPartition = Partition.builder() + .setDatabaseName(trinoDatabase.getDatabaseName()) + .setTableName(trinoTable.getTableName()) + .setValues(ImmutableList.of("val1")) + .setColumns(ImmutableList.of(new Column("partition-data", HIVE_STRING, Optional.of("partition data column comment"), Map.of()))) + .setParameters(ImmutableMap.of()) + .withStorage(storage -> storage.setStorageFormat(StorageFormat.create("PartitionSerdeLib", "PartitionInputFormat", "PartitionOutputFormat")) + .setLocation("/test-table/partition") + .setBucketProperty(Optional.empty()) + .setSerdeParameters(ImmutableMap.of())).build(); + + private final software.amazon.awssdk.services.glue.model.Database glueDatabase = software.amazon.awssdk.services.glue.model.Database.builder() + .name("test-database") + .description("database desc") + .locationUri("/database") + .parameters(ImmutableMap.of("key", "database-value")) + .build(); + + private final software.amazon.awssdk.services.glue.model.Table glueMaterializedView = software.amazon.awssdk.services.glue.model.Table.builder() + .databaseName(glueDatabase.name()) + .name("test-materialized-view") + .owner("owner") + .parameters(ImmutableMap.builder() + .put(PRESTO_VIEW_FLAG, "true") + .put(TABLE_COMMENT, ICEBERG_MATERIALIZED_VIEW_COMMENT) + .buildOrThrow()) + .tableType(TableType.VIRTUAL_VIEW.name()) + .viewOriginalText("/* %s: base64encodedquery */".formatted(ICEBERG_MATERIALIZED_VIEW_COMMENT)) + .viewExpandedText(ICEBERG_MATERIALIZED_VIEW_COMMENT) + .build(); + + private final software.amazon.awssdk.services.glue.model.Table glueTable = software.amazon.awssdk.services.glue.model.Table.builder() + .databaseName(glueDatabase.name()) + .name("test-table") + .owner("owner") + .parameters(ImmutableMap.of()) + .partitionKeys(software.amazon.awssdk.services.glue.model.Column.builder() + .name("table-partition") + .type("string") + .comment("table partition column comment") + .build()) + .storageDescriptor(StorageDescriptor.builder() + .bucketColumns(ImmutableList.of("test-bucket-col")) + .columns(software.amazon.awssdk.services.glue.model.Column.builder() + .name("table-data") + .type("string") + .comment("table data column comment") + .build()) + .parameters(ImmutableMap.of()) + .serdeInfo(SerDeInfo.builder() + .serializationLibrary("SerdeLib") + .parameters(ImmutableMap.of()) + .build()) + .inputFormat("InputFormat") + .outputFormat("OutputFormat") + .location("/test-table") + .numberOfBuckets(1) + .build()) + .tableType(EXTERNAL_TABLE.name()) + .viewOriginalText("originalText") + .viewExpandedText("expandedText") + .build(); + + private final software.amazon.awssdk.services.glue.model.Partition gluePartition = software.amazon.awssdk.services.glue.model.Partition.builder() + .databaseName(glueDatabase.name()) + .tableName(glueTable.name()) + .values(ImmutableList.of("val1")) + .parameters(ImmutableMap.of()) + .storageDescriptor(StorageDescriptor.builder() + .bucketColumns(ImmutableList.of("partition-bucket-col")) + .columns(software.amazon.awssdk.services.glue.model.Column.builder() + .name("partition-data") + .type("string") + .comment("partition data column comment") + .build()) + .parameters(ImmutableMap.of()) + .serdeInfo(SerDeInfo.builder() + .serializationLibrary("SerdeLib") + .parameters(ImmutableMap.of()) + .build()) + .inputFormat("InputFormat") + .outputFormat("OutputFormat") + .location("/test-table") + .numberOfBuckets(1) + .build()) + .build(); + + @Test + void testToGlueDatabaseInput() + { + DatabaseInput databaseInput = GlueConverter.toGlueDatabaseInput(trinoDatabase); + + assertThat(databaseInput.name()).isEqualTo(trinoDatabase.getDatabaseName()); + assertThat(databaseInput.description()).isEqualTo(trinoDatabase.getComment().orElse(null)); + assertThat(databaseInput.locationUri()).isEqualTo(trinoDatabase.getLocation().orElse(null)); + assertThat(databaseInput.parameters()).isEqualTo(trinoDatabase.getParameters()); + } + + @Test + void testToGlueTableInput() + { + TableInput tableInput = GlueConverter.toGlueTableInput(trinoTable); + + assertThat(tableInput.name()).isEqualTo(trinoTable.getTableName()); + assertThat(tableInput.owner()).isEqualTo(trinoTable.getOwner().orElse(null)); + assertThat(tableInput.tableType()).isEqualTo(trinoTable.getTableType()); + assertThat(tableInput.parameters()).isEqualTo(trinoTable.getParameters()); + assertColumnList(tableInput.storageDescriptor().columns(), trinoTable.getDataColumns()); + assertColumnList(tableInput.partitionKeys(), trinoTable.getPartitionColumns()); + assertStorage(tableInput.storageDescriptor(), trinoTable.getStorage()); + assertThat(tableInput.viewExpandedText()).isEqualTo(trinoTable.getViewExpandedText().orElse(null)); + assertThat(tableInput.viewOriginalText()).isEqualTo(trinoTable.getViewOriginalText().orElse(null)); + } + + @Test + void testToGluePartitionInput() + { + PartitionInput partitionInput = GlueConverter.toGluePartitionInput(trinoPartition); + + assertThat(partitionInput.parameters()).isEqualTo(trinoPartition.getParameters()); + assertStorage(partitionInput.storageDescriptor(), trinoPartition.getStorage()); + assertThat(partitionInput.values()).isEqualTo(trinoPartition.getValues()); + } + + @Test + void testConvertDatabase() + { + io.trino.plugin.hive.metastore.Database trinoDatabase = GlueConverter.fromGlueDatabase(glueDatabase); + assertThat(trinoDatabase.getDatabaseName()).isEqualTo(glueDatabase.name()); + assertThat(trinoDatabase.getLocation()).hasValue(glueDatabase.locationUri()); + assertThat(trinoDatabase.getComment()).hasValue(glueDatabase.description()); + assertThat(trinoDatabase.getParameters()).isEqualTo(glueDatabase.parameters()); + assertThat(trinoDatabase.getOwnerName()).isEqualTo(Optional.of(PUBLIC_OWNER)); + assertThat(trinoDatabase.getOwnerType()).isEqualTo(Optional.of(PrincipalType.ROLE)); + } + + @Test + void testConvertTable() + { + io.trino.plugin.hive.metastore.Table trinoTable = GlueConverter.fromGlueTable(glueTable, glueDatabase.name()); + assertThat(trinoTable.getTableName()).isEqualTo(glueTable.name()); + assertThat(trinoTable.getDatabaseName()).isEqualTo(glueDatabase.name()); + assertThat(trinoTable.getTableType()).isEqualTo(getTableTypeNullable(glueTable)); + assertThat(trinoTable.getOwner().orElse(null)).isEqualTo(glueTable.owner()); + assertThat(trinoTable.getParameters()).isEqualTo(glueTable.parameters()); + assertColumnList(glueTable.storageDescriptor().columns(), trinoTable.getDataColumns()); + assertColumnList(glueTable.partitionKeys(), trinoTable.getPartitionColumns()); + assertStorage(glueTable.storageDescriptor(), trinoTable.getStorage()); + assertThat(trinoTable.getViewOriginalText()).hasValue(glueTable.viewOriginalText()); + assertThat(trinoTable.getViewExpandedText()).hasValue(glueTable.viewExpandedText()); + } + + @Test + void testConvertTableWithOpenCSVSerDe() + { + software.amazon.awssdk.services.glue.model.Table glueTable = this.glueTable.toBuilder() + .storageDescriptor(this.glueTable.storageDescriptor().toBuilder() + .columns(ImmutableList.of(software.amazon.awssdk.services.glue.model.Column.builder() + .name("int_column") + .type("int") + .comment("int column") + .build())) + .inputFormat("org.apache.hadoop.hive.ql.io.parquet.MapredParquetInputFormat") + .outputFormat("org.apache.hadoop.hive.ql.io.parquet.MapredParquetOutputFormat") + .serdeInfo(SerDeInfo.builder() + .serializationLibrary("org.apache.hadoop.hive.serde2.OpenCSVSerde") + .parameters(ImmutableMap.of()) + .build()) + .build()) + .build(); + io.trino.plugin.hive.metastore.Table trinoTable = GlueConverter.fromGlueTable(glueTable, glueTable.databaseName()); + + assertThat(trinoTable.getTableName()).isEqualTo(glueTable.name()); + assertThat(trinoTable.getDatabaseName()).isEqualTo(glueDatabase.name()); + assertThat(trinoTable.getTableType()).isEqualTo(getTableTypeNullable(glueTable)); + assertThat(trinoTable.getOwner().orElse(null)).isEqualTo(glueTable.owner()); + assertThat(trinoTable.getParameters()).isEqualTo(glueTable.parameters()); + assertThat(trinoTable.getDataColumns()).hasSize(1); + assertThat(trinoTable.getDataColumns().get(0).getType()).isEqualTo(HIVE_STRING); + + assertColumnList(glueTable.partitionKeys(), trinoTable.getPartitionColumns()); + assertStorage(glueTable.storageDescriptor(), trinoTable.getStorage()); + assertThat(trinoTable.getViewOriginalText()).hasValue(glueTable.viewOriginalText()); + assertThat(trinoTable.getViewExpandedText()).hasValue(glueTable.viewExpandedText()); + } + + @Test + void testConvertTableWithoutTableType() + { + software.amazon.awssdk.services.glue.model.Table table = glueTable.toBuilder() + .tableType(null) + .build(); + io.trino.plugin.hive.metastore.Table trinoTable = GlueConverter.fromGlueTable(table, table.databaseName()); + assertThat(trinoTable.getTableType()).isEqualTo(EXTERNAL_TABLE.name()); + } + + @Test + void testConvertTableNullPartitions() + { + software.amazon.awssdk.services.glue.model.Table table = glueTable.toBuilder() + .partitionKeys(ImmutableList.of()) + .build(); + io.trino.plugin.hive.metastore.Table trinoTable = GlueConverter.fromGlueTable(table, table.databaseName()); + assertThat(trinoTable.getPartitionColumns()).isEmpty(); + } + + @Test + void testConvertTableUppercaseColumnType() + { + software.amazon.awssdk.services.glue.model.Table table = glueTable.toBuilder() + .partitionKeys(software.amazon.awssdk.services.glue.model.Column.builder() + .name("table-partition") + .type("String") + .comment("table partition column comment") + .build()) + .build(); + + assertThat(GlueConverter.fromGlueTable(table, table.databaseName()).getPartitionColumns().get(0).getType()).isEqualTo(HIVE_STRING); + } + + @Test + void testConvertPartition() + { + io.trino.plugin.hive.metastore.Partition trinoPartition = GlueConverter.fromGluePartition(gluePartition.databaseName(), gluePartition.tableName(), gluePartition); + assertThat(trinoPartition.getDatabaseName()).isEqualTo(gluePartition.databaseName()); + assertThat(trinoPartition.getTableName()).isEqualTo(gluePartition.tableName()); + assertColumnList(gluePartition.storageDescriptor().columns(), trinoPartition.getColumns()); + assertThat(trinoPartition.getValues()).isEqualTo(gluePartition.values()); + assertStorage(gluePartition.storageDescriptor(), trinoPartition.getStorage()); + assertThat(trinoPartition.getParameters()).isEqualTo(gluePartition.parameters()); + } + + @Test + void testDatabaseNullParameters() + { + software.amazon.awssdk.services.glue.model.Database database = glueDatabase.toBuilder() + .parameters(null) + .build(); + assertThat(GlueConverter.fromGlueDatabase(database).getParameters()).isNotNull(); + } + + @Test + void testTableNullParameters() + { + software.amazon.awssdk.services.glue.model.Table table = glueTable.toBuilder() + .parameters(null) + .storageDescriptor(glueTable.storageDescriptor().toBuilder() + .serdeInfo(glueTable.storageDescriptor().serdeInfo().toBuilder() + .parameters(null) + .build()) + .build()) + .build(); + io.trino.plugin.hive.metastore.Table trinoTable = GlueConverter.fromGlueTable(table, glueTable.databaseName()); + assertThat(trinoTable.getParameters()).isNotNull(); + assertThat(trinoTable.getStorage().getSerdeParameters()).isNotNull(); + } + + @Test + void testIcebergTableNullStorageDescriptor() + { + software.amazon.awssdk.services.glue.model.Table table = glueTable.toBuilder() + .parameters(ImmutableMap.of(ICEBERG_TABLE_TYPE_NAME, ICEBERG_TABLE_TYPE_VALUE)) + .storageDescriptor((StorageDescriptor) null) + .build(); + io.trino.plugin.hive.metastore.Table trinoTable = GlueConverter.fromGlueTable(table, table.databaseName()); + assertThat(trinoTable.getDataColumns()).hasSize(1); + } + + @Test + void testIcebergTableNonNullStorageDescriptor() + { + software.amazon.awssdk.services.glue.model.Table table = glueTable.toBuilder() + .parameters(ImmutableMap.of(ICEBERG_TABLE_TYPE_NAME, ICEBERG_TABLE_TYPE_VALUE)) + .build(); + assertThat(table.storageDescriptor()).isNotNull(); + io.trino.plugin.hive.metastore.Table trinoTable = GlueConverter.fromGlueTable(table, table.databaseName()); + assertThat(trinoTable.getDataColumns()).hasSize(1); + } + + @Test + void testDeltaTableNullStorageDescriptor() + { + software.amazon.awssdk.services.glue.model.Table table = glueTable.toBuilder() + .parameters(ImmutableMap.of(SPARK_TABLE_PROVIDER_KEY, DELTA_LAKE_PROVIDER)) + .storageDescriptor((StorageDescriptor) null) + .build(); + io.trino.plugin.hive.metastore.Table trinoTable = GlueConverter.fromGlueTable(table, table.databaseName()); + assertThat(trinoTable.getDataColumns()).hasSize(1); + } + + @Test + void testDeltaTableNonNullStorageDescriptor() + { + software.amazon.awssdk.services.glue.model.Table table = glueTable.toBuilder() + .parameters(ImmutableMap.of(SPARK_TABLE_PROVIDER_KEY, DELTA_LAKE_PROVIDER)) + .build(); + assertThat(table.storageDescriptor()).isNotNull(); + io.trino.plugin.hive.metastore.Table trinoTable = GlueConverter.fromGlueTable(table, table.databaseName()); + assertThat(trinoTable.getDataColumns()).hasSize(1); + } + + @Test + public void testIcebergMaterializedViewNullStorageDescriptor() + { + assertThat(glueMaterializedView.storageDescriptor()).isNull(); + Table trinoTable = GlueConverter.fromGlueTable(glueMaterializedView, glueMaterializedView.databaseName()); + assertThat(trinoTable.getDataColumns()).hasSize(1); + } + + @Test + void testPartitionNullParameters() + { + software.amazon.awssdk.services.glue.model.Partition partition = gluePartition.toBuilder() + .parameters(null) + .build(); + assertThat(GlueConverter.fromGluePartition(partition.databaseName(), partition.tableName(), partition).getParameters()).isNotNull(); + } + + private static void assertColumnList(List glueColumns, List trinoColumns) + { + if (trinoColumns == null) { + assertThat(glueColumns).isNull(); + } + assertThat(glueColumns).isNotNull(); + assertThat(glueColumns).hasSize(trinoColumns.size()); + + for (int i = 0; i < trinoColumns.size(); i++) { + assertColumn(glueColumns.get(i), trinoColumns.get(i)); + } + } + + private static void assertColumn(software.amazon.awssdk.services.glue.model.Column glueColumn, Column trinoColumn) + { + assertThat(glueColumn.name()).isEqualTo(trinoColumn.getName()); + assertThat(glueColumn.type()).isEqualTo(trinoColumn.getType().getHiveTypeName().toString()); + assertThat(glueColumn.comment()).isEqualTo(trinoColumn.getComment().orElse(null)); + } + + private static void assertStorage(StorageDescriptor glueStorage, Storage trinoStorage) + { + assertThat(glueStorage.location()).isEqualTo(trinoStorage.getLocation()); + assertThat(glueStorage.serdeInfo().serializationLibrary()).isEqualTo(trinoStorage.getStorageFormat().getSerde()); + assertThat(glueStorage.inputFormat()).isEqualTo(trinoStorage.getStorageFormat().getInputFormat()); + assertThat(glueStorage.outputFormat()).isEqualTo(trinoStorage.getStorageFormat().getOutputFormat()); + + if (trinoStorage.getBucketProperty().isPresent()) { + HiveBucketProperty bucketProperty = trinoStorage.getBucketProperty().get(); + assertThat(glueStorage.bucketColumns()).isEqualTo(bucketProperty.getBucketedBy()); + assertThat(glueStorage.numberOfBuckets().intValue()).isEqualTo(bucketProperty.getBucketCount()); + } + } +} diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/glue/TestGlueHiveMetastoreConfig.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/glue/TestGlueHiveMetastoreConfig.java index a5c6a4472fa8..93b95aa3a92a 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/glue/TestGlueHiveMetastoreConfig.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/glue/TestGlueHiveMetastoreConfig.java @@ -42,12 +42,12 @@ public void testDefaults() .setAwsAccessKey(null) .setAwsSecretKey(null) .setAwsCredentialsProvider(null) + .setUseWebIdentityTokenCredentialsProvider(false) .setCatalogId(null) .setPartitionSegments(5) - .setGetPartitionThreads(20) + .setThreads(40) .setAssumeCanonicalPartitionKeys(false) - .setReadStatisticsThreads(5) - .setWriteStatisticsThreads(20)); + .setSkipArchive(false)); } @Test @@ -69,11 +69,11 @@ public void testExplicitPropertyMapping() .put("hive.metastore.glue.aws-secret-key", "DEF") .put("hive.metastore.glue.aws-credentials-provider", "custom") .put("hive.metastore.glue.catalogid", "0123456789") + .put("hive.metastore.glue.use-web-identity-token-credentials-provider", "true") .put("hive.metastore.glue.partitions-segments", "10") .put("hive.metastore.glue.get-partition-threads", "42") .put("hive.metastore.glue.assume-canonical-partition-keys", "true") - .put("hive.metastore.glue.read-statistics-threads", "42") - .put("hive.metastore.glue.write-statistics-threads", "43") + .put("hive.metastore.glue.skip-archive", "true") .buildOrThrow(); GlueHiveMetastoreConfig expected = new GlueHiveMetastoreConfig() @@ -91,12 +91,12 @@ public void testExplicitPropertyMapping() .setAwsAccessKey("ABC") .setAwsSecretKey("DEF") .setAwsCredentialsProvider("custom") + .setUseWebIdentityTokenCredentialsProvider(true) .setCatalogId("0123456789") .setPartitionSegments(10) - .setGetPartitionThreads(42) + .setThreads(42) .setAssumeCanonicalPartitionKeys(true) - .setReadStatisticsThreads(42) - .setWriteStatisticsThreads(43); + .setSkipArchive(true); assertFullMapping(properties, expected); } diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/glue/TestGlueInputConverter.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/glue/TestGlueInputConverter.java deleted file mode 100644 index 9b812db8b8cc..000000000000 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/glue/TestGlueInputConverter.java +++ /dev/null @@ -1,113 +0,0 @@ -/* - * 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.hive.metastore.glue; - -import com.amazonaws.services.glue.model.DatabaseInput; -import com.amazonaws.services.glue.model.PartitionInput; -import com.amazonaws.services.glue.model.StorageDescriptor; -import com.amazonaws.services.glue.model.TableInput; -import com.google.common.collect.ImmutableList; -import io.trino.plugin.hive.HiveBucketProperty; -import io.trino.plugin.hive.metastore.Column; -import io.trino.plugin.hive.metastore.Database; -import io.trino.plugin.hive.metastore.Partition; -import io.trino.plugin.hive.metastore.Storage; -import io.trino.plugin.hive.metastore.Table; -import io.trino.plugin.hive.metastore.glue.converter.GlueInputConverter; -import org.testng.annotations.Test; - -import java.util.List; - -import static io.trino.plugin.hive.metastore.glue.TestingMetastoreObjects.getPrestoTestDatabase; -import static io.trino.plugin.hive.metastore.glue.TestingMetastoreObjects.getPrestoTestPartition; -import static io.trino.plugin.hive.metastore.glue.TestingMetastoreObjects.getPrestoTestTable; -import static org.testng.Assert.assertEquals; -import static org.testng.Assert.assertNull; - -public class TestGlueInputConverter -{ - private final Database testDb = getPrestoTestDatabase(); - private final Table testTbl = getPrestoTestTable(testDb.getDatabaseName()); - private final Partition testPartition = getPrestoTestPartition(testDb.getDatabaseName(), testTbl.getTableName(), ImmutableList.of("val1")); - - @Test - public void testConvertDatabase() - { - DatabaseInput dbInput = GlueInputConverter.convertDatabase(testDb); - - assertEquals(dbInput.getName(), testDb.getDatabaseName()); - assertEquals(dbInput.getDescription(), testDb.getComment().get()); - assertEquals(dbInput.getLocationUri(), testDb.getLocation().get()); - assertEquals(dbInput.getParameters(), testDb.getParameters()); - } - - @Test - public void testConvertTable() - { - TableInput tblInput = GlueInputConverter.convertTable(testTbl); - - assertEquals(tblInput.getName(), testTbl.getTableName()); - assertEquals(tblInput.getOwner(), testTbl.getOwner().orElse(null)); - assertEquals(tblInput.getTableType(), testTbl.getTableType()); - assertEquals(tblInput.getParameters(), testTbl.getParameters()); - assertColumnList(tblInput.getStorageDescriptor().getColumns(), testTbl.getDataColumns()); - assertColumnList(tblInput.getPartitionKeys(), testTbl.getPartitionColumns()); - assertStorage(tblInput.getStorageDescriptor(), testTbl.getStorage()); - assertEquals(tblInput.getViewExpandedText(), testTbl.getViewExpandedText().get()); - assertEquals(tblInput.getViewOriginalText(), testTbl.getViewOriginalText().get()); - } - - @Test - public void testConvertPartition() - { - PartitionInput partitionInput = GlueInputConverter.convertPartition(testPartition); - - assertEquals(partitionInput.getParameters(), testPartition.getParameters()); - assertStorage(partitionInput.getStorageDescriptor(), testPartition.getStorage()); - assertEquals(partitionInput.getValues(), testPartition.getValues()); - } - - private static void assertColumnList(List actual, List expected) - { - if (expected == null) { - assertNull(actual); - } - assertEquals(actual.size(), expected.size()); - - for (int i = 0; i < expected.size(); i++) { - assertColumn(actual.get(i), expected.get(i)); - } - } - - private static void assertColumn(com.amazonaws.services.glue.model.Column actual, Column expected) - { - assertEquals(actual.getName(), expected.getName()); - assertEquals(actual.getType(), expected.getType().getHiveTypeName().toString()); - assertEquals(actual.getComment(), expected.getComment().get()); - } - - private static void assertStorage(StorageDescriptor actual, Storage expected) - { - assertEquals(actual.getLocation(), expected.getLocation()); - assertEquals(actual.getSerdeInfo().getSerializationLibrary(), expected.getStorageFormat().getSerde()); - assertEquals(actual.getInputFormat(), expected.getStorageFormat().getInputFormat()); - assertEquals(actual.getOutputFormat(), expected.getStorageFormat().getOutputFormat()); - - if (expected.getBucketProperty().isPresent()) { - HiveBucketProperty bucketProperty = expected.getBucketProperty().get(); - assertEquals(actual.getBucketColumns(), bucketProperty.getBucketedBy()); - assertEquals(actual.getNumberOfBuckets().intValue(), bucketProperty.getBucketCount()); - } - } -} diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/glue/TestGlueToTrinoConverter.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/glue/TestGlueToTrinoConverter.java deleted file mode 100644 index 5f391aa7c7ca..000000000000 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/glue/TestGlueToTrinoConverter.java +++ /dev/null @@ -1,312 +0,0 @@ -/* - * 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.hive.metastore.glue; - -import com.amazonaws.services.glue.model.Database; -import com.amazonaws.services.glue.model.Partition; -import com.amazonaws.services.glue.model.StorageDescriptor; -import com.amazonaws.services.glue.model.Table; -import com.google.common.collect.ImmutableList; -import com.google.common.collect.ImmutableMap; -import io.trino.plugin.hive.HiveBucketProperty; -import io.trino.plugin.hive.metastore.Column; -import io.trino.plugin.hive.metastore.Storage; -import io.trino.plugin.hive.metastore.glue.converter.GlueToTrinoConverter; -import io.trino.plugin.hive.metastore.glue.converter.GlueToTrinoConverter.GluePartitionConverter; -import io.trino.spi.security.PrincipalType; -import org.testng.annotations.BeforeMethod; -import org.testng.annotations.Test; - -import java.util.ArrayList; -import java.util.HashMap; -import java.util.List; -import java.util.Optional; - -import static com.amazonaws.util.CollectionUtils.isNullOrEmpty; -import static com.google.common.collect.ImmutableSet.toImmutableSet; -import static io.trino.plugin.hive.HiveType.HIVE_STRING; -import static io.trino.plugin.hive.metastore.glue.TestingMetastoreObjects.getGlueTestColumn; -import static io.trino.plugin.hive.metastore.glue.TestingMetastoreObjects.getGlueTestDatabase; -import static io.trino.plugin.hive.metastore.glue.TestingMetastoreObjects.getGlueTestPartition; -import static io.trino.plugin.hive.metastore.glue.TestingMetastoreObjects.getGlueTestStorageDescriptor; -import static io.trino.plugin.hive.metastore.glue.TestingMetastoreObjects.getGlueTestTable; -import static io.trino.plugin.hive.metastore.glue.TestingMetastoreObjects.getGlueTestTrinoMaterializedView; -import static io.trino.plugin.hive.metastore.glue.converter.GlueToTrinoConverter.getPartitionParameters; -import static io.trino.plugin.hive.metastore.glue.converter.GlueToTrinoConverter.getTableParameters; -import static io.trino.plugin.hive.metastore.glue.converter.GlueToTrinoConverter.getTableTypeNullable; -import static io.trino.plugin.hive.util.HiveUtil.DELTA_LAKE_PROVIDER; -import static io.trino.plugin.hive.util.HiveUtil.ICEBERG_TABLE_TYPE_NAME; -import static io.trino.plugin.hive.util.HiveUtil.ICEBERG_TABLE_TYPE_VALUE; -import static io.trino.plugin.hive.util.HiveUtil.SPARK_TABLE_PROVIDER_KEY; -import static org.apache.hadoop.hive.metastore.TableType.EXTERNAL_TABLE; -import static org.testng.Assert.assertEquals; -import static org.testng.Assert.assertNotNull; -import static org.testng.Assert.assertNotSame; -import static org.testng.Assert.assertNull; -import static org.testng.Assert.assertSame; -import static org.testng.Assert.assertTrue; - -@Test(singleThreaded = true) -public class TestGlueToTrinoConverter -{ - private static final String PUBLIC_OWNER = "PUBLIC"; - - private Database testDatabase; - private Table testTable; - private Partition testPartition; - - @BeforeMethod - public void setup() - { - testDatabase = getGlueTestDatabase(); - testTable = getGlueTestTable(testDatabase.getName()); - testPartition = getGlueTestPartition(testDatabase.getName(), testTable.getName(), ImmutableList.of("val1")); - } - - private static GluePartitionConverter createPartitionConverter(Table table) - { - return new GluePartitionConverter(GlueToTrinoConverter.convertTable(table, table.getDatabaseName())); - } - - @Test - public void testConvertDatabase() - { - io.trino.plugin.hive.metastore.Database trinoDatabase = GlueToTrinoConverter.convertDatabase(testDatabase); - assertEquals(trinoDatabase.getDatabaseName(), testDatabase.getName()); - assertEquals(trinoDatabase.getLocation().get(), testDatabase.getLocationUri()); - assertEquals(trinoDatabase.getComment().get(), testDatabase.getDescription()); - assertEquals(trinoDatabase.getParameters(), testDatabase.getParameters()); - assertEquals(trinoDatabase.getOwnerName(), Optional.of(PUBLIC_OWNER)); - assertEquals(trinoDatabase.getOwnerType(), Optional.of(PrincipalType.ROLE)); - } - - @Test - public void testConvertTable() - { - io.trino.plugin.hive.metastore.Table trinoTable = GlueToTrinoConverter.convertTable(testTable, testDatabase.getName()); - assertEquals(trinoTable.getTableName(), testTable.getName()); - assertEquals(trinoTable.getDatabaseName(), testDatabase.getName()); - assertEquals(trinoTable.getTableType(), getTableTypeNullable(testTable)); - assertEquals(trinoTable.getOwner().orElse(null), testTable.getOwner()); - assertEquals(trinoTable.getParameters(), getTableParameters(testTable)); - assertColumnList(trinoTable.getDataColumns(), testTable.getStorageDescriptor().getColumns()); - assertColumnList(trinoTable.getPartitionColumns(), testTable.getPartitionKeys()); - assertStorage(trinoTable.getStorage(), testTable.getStorageDescriptor()); - assertEquals(trinoTable.getViewOriginalText().get(), testTable.getViewOriginalText()); - assertEquals(trinoTable.getViewExpandedText().get(), testTable.getViewExpandedText()); - } - - @Test - public void testConvertTableWithOpenCSVSerDe() - { - Table glueTable = getGlueTestTable(testDatabase.getName()); - glueTable.setStorageDescriptor(getGlueTestStorageDescriptor( - ImmutableList.of(getGlueTestColumn("int")), - "org.apache.hadoop.hive.serde2.OpenCSVSerde")); - io.trino.plugin.hive.metastore.Table trinoTable = GlueToTrinoConverter.convertTable(glueTable, testDatabase.getName()); - - assertEquals(trinoTable.getTableName(), glueTable.getName()); - assertEquals(trinoTable.getDatabaseName(), testDatabase.getName()); - assertEquals(trinoTable.getTableType(), getTableTypeNullable(glueTable)); - assertEquals(trinoTable.getOwner().orElse(null), glueTable.getOwner()); - assertEquals(trinoTable.getParameters(), getTableParameters(glueTable)); - assertEquals(trinoTable.getDataColumns().size(), 1); - assertEquals(trinoTable.getDataColumns().get(0).getType(), HIVE_STRING); - - assertColumnList(trinoTable.getPartitionColumns(), glueTable.getPartitionKeys()); - assertStorage(trinoTable.getStorage(), glueTable.getStorageDescriptor()); - assertEquals(trinoTable.getViewOriginalText().get(), glueTable.getViewOriginalText()); - assertEquals(trinoTable.getViewExpandedText().get(), glueTable.getViewExpandedText()); - } - - @Test - public void testConvertTableWithoutTableType() - { - Table table = getGlueTestTable(testDatabase.getName()); - table.setTableType(null); - io.trino.plugin.hive.metastore.Table trinoTable = GlueToTrinoConverter.convertTable(table, testDatabase.getName()); - assertEquals(trinoTable.getTableType(), EXTERNAL_TABLE.name()); - } - - @Test - public void testConvertTableNullPartitions() - { - testTable.setPartitionKeys(null); - io.trino.plugin.hive.metastore.Table trinoTable = GlueToTrinoConverter.convertTable(testTable, testDatabase.getName()); - assertTrue(trinoTable.getPartitionColumns().isEmpty()); - } - - @Test - public void testConvertTableUppercaseColumnType() - { - com.amazonaws.services.glue.model.Column uppercaseColumn = getGlueTestColumn().withType("String"); - testTable.getStorageDescriptor().setColumns(ImmutableList.of(uppercaseColumn)); - GlueToTrinoConverter.convertTable(testTable, testDatabase.getName()); - } - - @Test - public void testConvertPartition() - { - GluePartitionConverter converter = createPartitionConverter(testTable); - io.trino.plugin.hive.metastore.Partition trinoPartition = converter.apply(testPartition); - assertEquals(trinoPartition.getDatabaseName(), testPartition.getDatabaseName()); - assertEquals(trinoPartition.getTableName(), testPartition.getTableName()); - assertColumnList(trinoPartition.getColumns(), testPartition.getStorageDescriptor().getColumns()); - assertEquals(trinoPartition.getValues(), testPartition.getValues()); - assertStorage(trinoPartition.getStorage(), testPartition.getStorageDescriptor()); - assertEquals(trinoPartition.getParameters(), getPartitionParameters(testPartition)); - } - - @Test - public void testPartitionConversionMemoization() - { - String fakeS3Location = "s3://some-fake-location"; - testPartition.getStorageDescriptor().setLocation(fakeS3Location); - // Second partition to convert with equal (but not aliased) values - Partition partitionTwo = getGlueTestPartition("" + testDatabase.getName(), "" + testTable.getName(), new ArrayList<>(testPartition.getValues())); - // Ensure storage fields match as well - partitionTwo.getStorageDescriptor().setColumns(new ArrayList<>(testPartition.getStorageDescriptor().getColumns())); - partitionTwo.getStorageDescriptor().setBucketColumns(new ArrayList<>(testPartition.getStorageDescriptor().getBucketColumns())); - partitionTwo.getStorageDescriptor().setLocation("" + fakeS3Location); - partitionTwo.getStorageDescriptor().setInputFormat("" + testPartition.getStorageDescriptor().getInputFormat()); - partitionTwo.getStorageDescriptor().setOutputFormat("" + testPartition.getStorageDescriptor().getOutputFormat()); - partitionTwo.getStorageDescriptor().setParameters(new HashMap<>(testPartition.getStorageDescriptor().getParameters())); - - GluePartitionConverter converter = createPartitionConverter(testTable); - io.trino.plugin.hive.metastore.Partition trinoPartition = converter.apply(testPartition); - io.trino.plugin.hive.metastore.Partition trinoPartition2 = converter.apply(partitionTwo); - - assertNotSame(trinoPartition, trinoPartition2); - assertSame(trinoPartition2.getDatabaseName(), trinoPartition.getDatabaseName()); - assertSame(trinoPartition2.getTableName(), trinoPartition.getTableName()); - assertSame(trinoPartition2.getColumns(), trinoPartition.getColumns()); - assertSame(trinoPartition2.getParameters(), trinoPartition.getParameters()); - assertNotSame(trinoPartition2.getValues(), trinoPartition.getValues()); - - Storage storage = trinoPartition.getStorage(); - Storage storage2 = trinoPartition2.getStorage(); - - assertSame(storage2.getStorageFormat(), storage.getStorageFormat()); - assertSame(storage2.getBucketProperty(), storage.getBucketProperty()); - assertSame(storage2.getSerdeParameters(), storage.getSerdeParameters()); - assertNotSame(storage2.getLocation(), storage.getLocation()); - } - - @Test - public void testDatabaseNullParameters() - { - testDatabase.setParameters(null); - assertNotNull(GlueToTrinoConverter.convertDatabase(testDatabase).getParameters()); - } - - @Test - public void testTableNullParameters() - { - testTable.setParameters(null); - testTable.getStorageDescriptor().getSerdeInfo().setParameters(null); - io.trino.plugin.hive.metastore.Table trinoTable = GlueToTrinoConverter.convertTable(testTable, testDatabase.getName()); - assertNotNull(trinoTable.getParameters()); - assertNotNull(trinoTable.getStorage().getSerdeParameters()); - } - - @Test - public void testIcebergTableNullStorageDescriptor() - { - testTable.setParameters(ImmutableMap.of(ICEBERG_TABLE_TYPE_NAME, ICEBERG_TABLE_TYPE_VALUE)); - testTable.setStorageDescriptor(null); - io.trino.plugin.hive.metastore.Table trinoTable = GlueToTrinoConverter.convertTable(testTable, testDatabase.getName()); - assertEquals(trinoTable.getDataColumns().size(), 1); - } - - @Test - public void testIcebergTableNonNullStorageDescriptor() - { - testTable.setParameters(ImmutableMap.of(ICEBERG_TABLE_TYPE_NAME, ICEBERG_TABLE_TYPE_VALUE)); - assertNotNull(testTable.getStorageDescriptor()); - io.trino.plugin.hive.metastore.Table trinoTable = GlueToTrinoConverter.convertTable(testTable, testDatabase.getName()); - assertEquals(trinoTable.getDataColumns().size(), 1); - } - - @Test - public void testDeltaTableNullStorageDescriptor() - { - testTable.setParameters(ImmutableMap.of(SPARK_TABLE_PROVIDER_KEY, DELTA_LAKE_PROVIDER)); - testTable.setStorageDescriptor(null); - io.trino.plugin.hive.metastore.Table trinoTable = GlueToTrinoConverter.convertTable(testTable, testDatabase.getName()); - assertEquals(trinoTable.getDataColumns().size(), 1); - } - - @Test - public void testDeltaTableNonNullStorageDescriptor() - { - testTable.setParameters(ImmutableMap.of(SPARK_TABLE_PROVIDER_KEY, DELTA_LAKE_PROVIDER)); - assertNotNull(testTable.getStorageDescriptor()); - io.trino.plugin.hive.metastore.Table trinoTable = GlueToTrinoConverter.convertTable(testTable, testDatabase.getName()); - assertEquals( - trinoTable.getDataColumns().stream() - .map(Column::getName) - .collect(toImmutableSet()), - testTable.getStorageDescriptor().getColumns().stream() - .map(com.amazonaws.services.glue.model.Column::getName) - .collect(toImmutableSet())); - } - - @Test - public void testIcebergMaterializedViewNullStorageDescriptor() - { - Table testMaterializedView = getGlueTestTrinoMaterializedView(testDatabase.getName()); - assertNull(testMaterializedView.getStorageDescriptor()); - io.trino.plugin.hive.metastore.Table trinoTable = GlueToTrinoConverter.convertTable(testMaterializedView, testDatabase.getName()); - assertEquals(trinoTable.getDataColumns().size(), 1); - } - - @Test - public void testPartitionNullParameters() - { - testPartition.setParameters(null); - assertNotNull(createPartitionConverter(testTable).apply(testPartition).getParameters()); - } - - private static void assertColumnList(List actual, List expected) - { - if (expected == null) { - assertNull(actual); - } - assertEquals(actual.size(), expected.size()); - - for (int i = 0; i < expected.size(); i++) { - assertColumn(actual.get(i), expected.get(i)); - } - } - - private static void assertColumn(Column actual, com.amazonaws.services.glue.model.Column expected) - { - assertEquals(actual.getName(), expected.getName()); - assertEquals(actual.getType().getHiveTypeName().toString(), expected.getType()); - assertEquals(actual.getComment().get(), expected.getComment()); - } - - private static void assertStorage(Storage actual, StorageDescriptor expected) - { - assertEquals(actual.getLocation(), expected.getLocation()); - assertEquals(actual.getStorageFormat().getSerde(), expected.getSerdeInfo().getSerializationLibrary()); - assertEquals(actual.getStorageFormat().getInputFormat(), expected.getInputFormat()); - assertEquals(actual.getStorageFormat().getOutputFormat(), expected.getOutputFormat()); - if (!isNullOrEmpty(expected.getBucketColumns())) { - HiveBucketProperty bucketProperty = actual.getBucketProperty().get(); - assertEquals(bucketProperty.getBucketedBy(), expected.getBucketColumns()); - assertEquals(bucketProperty.getBucketCount(), expected.getNumberOfBuckets().intValue()); - } - } -} diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/glue/TestHiveGlueMetastore.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/glue/TestHiveGlueMetastore.java deleted file mode 100644 index cda718aa14b5..000000000000 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/glue/TestHiveGlueMetastore.java +++ /dev/null @@ -1,1558 +0,0 @@ -/* - * 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.hive.metastore.glue; - -import com.amazonaws.auth.DefaultAWSCredentialsProviderChain; -import com.amazonaws.services.glue.AWSGlueAsync; -import com.amazonaws.services.glue.AWSGlueAsyncClientBuilder; -import com.amazonaws.services.glue.model.CreateTableRequest; -import com.amazonaws.services.glue.model.Database; -import com.amazonaws.services.glue.model.DeleteDatabaseRequest; -import com.amazonaws.services.glue.model.DeleteTableRequest; -import com.amazonaws.services.glue.model.EntityNotFoundException; -import com.amazonaws.services.glue.model.GetDatabasesRequest; -import com.amazonaws.services.glue.model.GetDatabasesResult; -import com.amazonaws.services.glue.model.TableInput; -import com.amazonaws.services.glue.model.UpdateTableRequest; -import com.google.common.collect.ImmutableList; -import com.google.common.collect.ImmutableMap; -import com.google.common.collect.ImmutableSet; -import io.airlift.concurrent.BoundedExecutor; -import io.airlift.log.Logger; -import io.airlift.slice.Slice; -import io.trino.plugin.hive.AbstractTestHiveLocal; -import io.trino.plugin.hive.HiveBasicStatistics; -import io.trino.plugin.hive.HiveMetastoreClosure; -import io.trino.plugin.hive.HiveType; -import io.trino.plugin.hive.PartitionStatistics; -import io.trino.plugin.hive.metastore.HiveColumnStatistics; -import io.trino.plugin.hive.metastore.HiveMetastore; -import io.trino.plugin.hive.metastore.PartitionWithStatistics; -import io.trino.plugin.hive.metastore.Table; -import io.trino.plugin.hive.metastore.glue.converter.GlueInputConverter; -import io.trino.spi.TrinoException; -import io.trino.spi.block.Block; -import io.trino.spi.block.BlockBuilder; -import io.trino.spi.connector.ColumnMetadata; -import io.trino.spi.connector.ConnectorMetadata; -import io.trino.spi.connector.ConnectorOutputTableHandle; -import io.trino.spi.connector.ConnectorPageSink; -import io.trino.spi.connector.ConnectorSession; -import io.trino.spi.connector.ConnectorTableMetadata; -import io.trino.spi.connector.SchemaTableName; -import io.trino.spi.connector.TableNotFoundException; -import io.trino.spi.predicate.Domain; -import io.trino.spi.predicate.Range; -import io.trino.spi.predicate.TupleDomain; -import io.trino.spi.predicate.ValueSet; -import io.trino.spi.statistics.ComputedStatistics; -import io.trino.spi.statistics.TableStatisticType; -import io.trino.spi.type.BigintType; -import io.trino.spi.type.DateType; -import io.trino.spi.type.IntegerType; -import io.trino.spi.type.SmallintType; -import io.trino.spi.type.TimestampType; -import io.trino.spi.type.TinyintType; -import io.trino.spi.type.Type; -import io.trino.spi.type.VarcharType; -import io.trino.testing.MaterializedResult; -import org.testng.annotations.BeforeClass; -import org.testng.annotations.DataProvider; -import org.testng.annotations.Test; - -import java.io.File; -import java.util.ArrayList; -import java.util.Arrays; -import java.util.Collection; -import java.util.Collections; -import java.util.List; -import java.util.Map; -import java.util.Optional; -import java.util.OptionalLong; -import java.util.Set; -import java.util.concurrent.Executor; -import java.util.function.Supplier; - -import static com.google.common.collect.ImmutableList.toImmutableList; -import static io.airlift.concurrent.MoreFutures.getFutureValue; -import static io.airlift.slice.Slices.utf8Slice; -import static io.trino.plugin.hive.HiveBasicStatistics.createEmptyStatistics; -import static io.trino.plugin.hive.HiveColumnStatisticType.MAX_VALUE; -import static io.trino.plugin.hive.HiveColumnStatisticType.MIN_VALUE; -import static io.trino.plugin.hive.HiveColumnStatisticType.NUMBER_OF_DISTINCT_VALUES; -import static io.trino.plugin.hive.HiveColumnStatisticType.NUMBER_OF_NON_NULL_VALUES; -import static io.trino.plugin.hive.HiveMetadata.TABLE_COMMENT; -import static io.trino.plugin.hive.HiveStorageFormat.ORC; -import static io.trino.plugin.hive.HiveStorageFormat.TEXTFILE; -import static io.trino.plugin.hive.HiveTestUtils.HDFS_ENVIRONMENT; -import static io.trino.plugin.hive.ViewReaderUtil.ICEBERG_MATERIALIZED_VIEW_COMMENT; -import static io.trino.plugin.hive.ViewReaderUtil.PRESTO_VIEW_FLAG; -import static io.trino.plugin.hive.ViewReaderUtil.isTrinoMaterializedView; -import static io.trino.plugin.hive.acid.AcidTransaction.NO_ACID_TRANSACTION; -import static io.trino.plugin.hive.metastore.HiveColumnStatistics.createIntegerColumnStatistics; -import static io.trino.plugin.hive.metastore.glue.AwsSdkUtil.getPaginatedResults; -import static io.trino.plugin.hive.metastore.glue.GlueClientUtil.createAsyncGlueClient; -import static io.trino.plugin.hive.metastore.glue.PartitionFilterBuilder.DECIMAL_TYPE; -import static io.trino.plugin.hive.metastore.glue.PartitionFilterBuilder.decimalOf; -import static io.trino.plugin.hive.util.HiveUtil.DELTA_LAKE_PROVIDER; -import static io.trino.plugin.hive.util.HiveUtil.ICEBERG_TABLE_TYPE_NAME; -import static io.trino.plugin.hive.util.HiveUtil.ICEBERG_TABLE_TYPE_VALUE; -import static io.trino.plugin.hive.util.HiveUtil.SPARK_TABLE_PROVIDER_KEY; -import static io.trino.plugin.hive.util.HiveUtil.isDeltaLakeTable; -import static io.trino.plugin.hive.util.HiveUtil.isIcebergTable; -import static io.trino.spi.connector.RetryMode.NO_RETRIES; -import static io.trino.spi.type.BigintType.BIGINT; -import static io.trino.spi.type.VarcharType.VARCHAR; -import static io.trino.spi.type.VarcharType.createUnboundedVarcharType; -import static io.trino.testing.TestingPageSinkId.TESTING_PAGE_SINK_ID; -import static java.lang.String.format; -import static java.lang.System.currentTimeMillis; -import static java.util.Collections.unmodifiableList; -import static java.util.Locale.ENGLISH; -import static java.util.Objects.requireNonNull; -import static java.util.UUID.randomUUID; -import static java.util.concurrent.TimeUnit.DAYS; -import static org.apache.hadoop.hive.common.FileUtils.makePartName; -import static org.apache.hadoop.hive.metastore.TableType.EXTERNAL_TABLE; -import static org.apache.hadoop.hive.metastore.TableType.VIRTUAL_VIEW; -import static org.assertj.core.api.Assertions.assertThat; -import static org.assertj.core.api.Assertions.assertThatThrownBy; -import static org.testng.Assert.assertEquals; -import static org.testng.Assert.assertTrue; - -/* - * GlueHiveMetastore currently uses AWS Default Credential Provider Chain, - * See https://docs.aws.amazon.com/sdk-for-java/v1/developer-guide/credentials.html#credentials-default - * on ways to set your AWS credentials which will be needed to run this test. - */ -@Test(singleThreaded = true) -public class TestHiveGlueMetastore - extends AbstractTestHiveLocal -{ - private static final Logger log = Logger.get(TestHiveGlueMetastore.class); - - private static final String PARTITION_KEY = "part_key_1"; - private static final String PARTITION_KEY2 = "part_key_2"; - private static final String TEST_DATABASE_NAME_PREFIX = "test_glue"; - - private static final List CREATE_TABLE_COLUMNS = ImmutableList.of(new ColumnMetadata("id", BIGINT)); - private static final List CREATE_TABLE_COLUMNS_PARTITIONED_VARCHAR = ImmutableList.builder() - .addAll(CREATE_TABLE_COLUMNS) - .add(new ColumnMetadata(PARTITION_KEY, VarcharType.VARCHAR)) - .build(); - private static final List CREATE_TABLE_COLUMNS_PARTITIONED_TWO_KEYS = ImmutableList.builder() - .addAll(CREATE_TABLE_COLUMNS) - .add(new ColumnMetadata(PARTITION_KEY, VarcharType.VARCHAR)) - .add(new ColumnMetadata(PARTITION_KEY2, BigintType.BIGINT)) - .build(); - private static final List CREATE_TABLE_COLUMNS_PARTITIONED_TINYINT = ImmutableList.builder() - .addAll(CREATE_TABLE_COLUMNS) - .add(new ColumnMetadata(PARTITION_KEY, TinyintType.TINYINT)) - .build(); - private static final List CREATE_TABLE_COLUMNS_PARTITIONED_SMALLINT = ImmutableList.builder() - .addAll(CREATE_TABLE_COLUMNS) - .add(new ColumnMetadata(PARTITION_KEY, SmallintType.SMALLINT)) - .build(); - private static final List CREATE_TABLE_COLUMNS_PARTITIONED_INTEGER = ImmutableList.builder() - .addAll(CREATE_TABLE_COLUMNS) - .add(new ColumnMetadata(PARTITION_KEY, IntegerType.INTEGER)) - .build(); - private static final List CREATE_TABLE_COLUMNS_PARTITIONED_BIGINT = ImmutableList.builder() - .addAll(CREATE_TABLE_COLUMNS) - .add(new ColumnMetadata(PARTITION_KEY, BigintType.BIGINT)) - .build(); - private static final List CREATE_TABLE_COLUMNS_PARTITIONED_DECIMAL = ImmutableList.builder() - .addAll(CREATE_TABLE_COLUMNS) - .add(new ColumnMetadata(PARTITION_KEY, DECIMAL_TYPE)) - .build(); - private static final List CREATE_TABLE_COLUMNS_PARTITIONED_DATE = ImmutableList.builder() - .addAll(CREATE_TABLE_COLUMNS) - .add(new ColumnMetadata(PARTITION_KEY, DateType.DATE)) - .build(); - private static final List CREATE_TABLE_COLUMNS_PARTITIONED_TIMESTAMP = ImmutableList.builder() - .addAll(CREATE_TABLE_COLUMNS) - .add(new ColumnMetadata(PARTITION_KEY, TimestampType.TIMESTAMP_MILLIS)) - .build(); - private static final List VARCHAR_PARTITION_VALUES = ImmutableList.of("2020-01-01", "2020-02-01", "2020-03-01", "2020-04-01"); - - protected static final HiveBasicStatistics HIVE_BASIC_STATISTICS = new HiveBasicStatistics(1000, 5000, 3000, 4000); - protected static final HiveColumnStatistics INTEGER_COLUMN_STATISTICS = createIntegerColumnStatistics( - OptionalLong.of(-1000), - OptionalLong.of(1000), - OptionalLong.of(1), - OptionalLong.of(2)); - - private HiveMetastoreClosure metastore; - private AWSGlueAsync glueClient; - - public TestHiveGlueMetastore() - { - super(TEST_DATABASE_NAME_PREFIX + randomUUID().toString().toLowerCase(ENGLISH).replace("-", "")); - } - - protected AWSGlueAsync getGlueClient() - { - return glueClient; - } - - @BeforeClass(alwaysRun = true) - @Override - public void initialize() - throws Exception - { - super.initialize(); - // uncomment to get extra AWS debug information -// Logging logging = Logging.initialize(); -// logging.setLevel("com.amazonaws.request", Level.DEBUG); - } - - @BeforeClass - public void setup() - { - metastore = new HiveMetastoreClosure(metastoreClient); - glueClient = AWSGlueAsyncClientBuilder.defaultClient(); - } - - @Override - protected HiveMetastore createMetastore(File tempDir) - { - GlueHiveMetastoreConfig glueConfig = new GlueHiveMetastoreConfig(); - glueConfig.setDefaultWarehouseDir(tempDir.toURI().toString()); - glueConfig.setAssumeCanonicalPartitionKeys(true); - - Executor executor = new BoundedExecutor(this.executor, 10); - GlueMetastoreStats stats = new GlueMetastoreStats(); - return new GlueHiveMetastore( - HDFS_ENVIRONMENT, - glueConfig, - executor, - new DefaultGlueColumnStatisticsProviderFactory(executor, executor), - createAsyncGlueClient(glueConfig, DefaultAWSCredentialsProviderChain.getInstance(), ImmutableSet.of(), stats.newRequestMetricsCollector()), - stats, - new DefaultGlueMetastoreTableFilterProvider(true).get()); - } - - @Test - public void cleanupOrphanedDatabases() - { - long creationTimeMillisThreshold = currentTimeMillis() - DAYS.toMillis(1); - GlueHiveMetastore metastore = (GlueHiveMetastore) getMetastoreClient(); - GlueMetastoreStats stats = metastore.getStats(); - List orphanedDatabases = getPaginatedResults( - glueClient::getDatabases, - new GetDatabasesRequest(), - GetDatabasesRequest::setNextToken, - GetDatabasesResult::getNextToken, - stats.getGetDatabases()) - .map(GetDatabasesResult::getDatabaseList) - .flatMap(List::stream) - .filter(database -> database.getName().startsWith(TEST_DATABASE_NAME_PREFIX) && - database.getCreateTime().getTime() <= creationTimeMillisThreshold) - .map(Database::getName) - .collect(toImmutableList()); - - log.info("Found %s %s* databases that look orphaned, removing", orphanedDatabases.size(), TEST_DATABASE_NAME_PREFIX); - orphanedDatabases.forEach(database -> { - try { - glueClient.deleteDatabase(new DeleteDatabaseRequest() - .withName(database)); - } - catch (EntityNotFoundException e) { - log.info("Database [%s] not found, could be removed by other cleanup process", database); - } - catch (RuntimeException e) { - log.warn(e, "Failed to remove database [%s]", database); - } - }); - } - - @Override - public void testRenameTable() - { - // rename table is not yet supported by Glue - } - - @Override - public void testUpdateTableColumnStatisticsEmptyOptionalFields() - { - // this test expects consistency between written and read stats but this is not provided by glue at the moment - // when writing empty min/max statistics glue will return 0 to the readers - // in order to avoid incorrect data we skip writes for statistics with min/max = null - } - - @Override - public void testUpdatePartitionColumnStatisticsEmptyOptionalFields() - { - // this test expects consistency between written and read stats but this is not provided by glue at the moment - // when writing empty min/max statistics glue will return 0 to the readers - // in order to avoid incorrect data we skip writes for statistics with min/max = null - } - - @Override - public void testUpdateBasicPartitionStatistics() - throws Exception - { - SchemaTableName tableName = temporaryTable("update_basic_partition_statistics"); - try { - createDummyPartitionedTable(tableName, STATISTICS_PARTITIONED_TABLE_COLUMNS); - testUpdatePartitionStatistics( - tableName, - PartitionStatistics.empty(), - ImmutableList.of(BASIC_STATISTICS_1, BASIC_STATISTICS_2), - ImmutableList.of(BASIC_STATISTICS_2, BASIC_STATISTICS_1)); - } - finally { - dropTable(tableName); - } - } - - @Override - public void testUpdatePartitionColumnStatistics() - throws Exception - { - SchemaTableName tableName = temporaryTable("update_partition_column_statistics"); - try { - createDummyPartitionedTable(tableName, STATISTICS_PARTITIONED_TABLE_COLUMNS); - // When the table has partitions, but row count statistics are set to zero, we treat this case as empty - // statistics to avoid underestimation in the CBO. This scenario may be caused when other engines are - // used to ingest data into partitioned hive tables. - testUpdatePartitionStatistics( - tableName, - PartitionStatistics.empty(), - ImmutableList.of(STATISTICS_1_1, STATISTICS_1_2, STATISTICS_2), - ImmutableList.of(STATISTICS_1_2, STATISTICS_1_1, STATISTICS_2)); - } - finally { - dropTable(tableName); - } - } - - @Override - public void testStorePartitionWithStatistics() - throws Exception - { - // When the table has partitions, but row count statistics are set to zero, we treat this case as empty - // statistics to avoid underestimation in the CBO. This scenario may be caused when other engines are - // used to ingest data into partitioned hive tables. - testStorePartitionWithStatistics(STATISTICS_PARTITIONED_TABLE_COLUMNS, BASIC_STATISTICS_1, BASIC_STATISTICS_2, BASIC_STATISTICS_1, PartitionStatistics.empty()); - } - - @Override - public void testGetPartitions() - throws Exception - { - try { - SchemaTableName tableName = temporaryTable("get_partitions"); - createDummyPartitionedTable(tableName, CREATE_TABLE_COLUMNS_PARTITIONED); - HiveMetastore metastoreClient = getMetastoreClient(); - Optional> partitionNames = metastoreClient.getPartitionNamesByFilter( - tableName.getSchemaName(), - tableName.getTableName(), - ImmutableList.of("ds"), TupleDomain.all()); - assertTrue(partitionNames.isPresent()); - assertEquals(partitionNames.get(), ImmutableList.of("ds=2016-01-01", "ds=2016-01-02")); - } - finally { - dropTable(tablePartitionFormat); - } - } - - @Test - public void testGetPartitionsWithFilterUsingReservedKeywordsAsColumnName() - throws Exception - { - SchemaTableName tableName = temporaryTable("get_partitions_with_filter_using_reserved_keyword_column_name"); - try { - String reservedKeywordPartitionColumnName = "key"; - String regularColumnPartitionName = "int_partition"; - List columns = ImmutableList.builder() - .add(new ColumnMetadata("t_string", createUnboundedVarcharType())) - .add(new ColumnMetadata(reservedKeywordPartitionColumnName, createUnboundedVarcharType())) - .add(new ColumnMetadata(regularColumnPartitionName, BIGINT)) - .build(); - List partitionedBy = ImmutableList.of(reservedKeywordPartitionColumnName, regularColumnPartitionName); - - doCreateEmptyTable(tableName, ORC, columns, partitionedBy); - - HiveMetastoreClosure metastoreClient = new HiveMetastoreClosure(getMetastoreClient()); - Table table = metastoreClient.getTable(tableName.getSchemaName(), tableName.getTableName()) - .orElseThrow(() -> new TableNotFoundException(tableName)); - - String partitionName1 = makePartName(ImmutableList.of(reservedKeywordPartitionColumnName, regularColumnPartitionName), ImmutableList.of("value1", "1")); - String partitionName2 = makePartName(ImmutableList.of(reservedKeywordPartitionColumnName, regularColumnPartitionName), ImmutableList.of("value2", "2")); - - List partitions = ImmutableList.of(partitionName1, partitionName2) - .stream() - .map(partitionName -> new PartitionWithStatistics(createDummyPartition(table, partitionName), partitionName, PartitionStatistics.empty())) - .collect(toImmutableList()); - metastoreClient.addPartitions(tableName.getSchemaName(), tableName.getTableName(), partitions); - metastoreClient.updatePartitionStatistics(tableName.getSchemaName(), tableName.getTableName(), partitionName1, currentStatistics -> ZERO_TABLE_STATISTICS); - metastoreClient.updatePartitionStatistics(tableName.getSchemaName(), tableName.getTableName(), partitionName2, currentStatistics -> ZERO_TABLE_STATISTICS); - - Optional> partitionNames = metastoreClient.getPartitionNamesByFilter( - tableName.getSchemaName(), - tableName.getTableName(), - ImmutableList.of(reservedKeywordPartitionColumnName, regularColumnPartitionName), - TupleDomain.withColumnDomains(ImmutableMap.of(regularColumnPartitionName, Domain.singleValue(BIGINT, 2L)))); - assertTrue(partitionNames.isPresent()); - assertEquals(partitionNames.get(), ImmutableList.of("key=value2/int_partition=2")); - - // KEY is a reserved keyword in the grammar of the SQL parser used internally by Glue API - // and therefore should not be used in the partition filter - partitionNames = metastoreClient.getPartitionNamesByFilter( - tableName.getSchemaName(), - tableName.getTableName(), - ImmutableList.of(reservedKeywordPartitionColumnName, regularColumnPartitionName), - TupleDomain.withColumnDomains(ImmutableMap.of(reservedKeywordPartitionColumnName, Domain.singleValue(VARCHAR, utf8Slice("value1"))))); - assertTrue(partitionNames.isPresent()); - assertEquals(partitionNames.get(), ImmutableList.of("key=value1/int_partition=1", "key=value2/int_partition=2")); - } - finally { - dropTable(tableName); - } - } - - @Test - public void testGetDatabasesLogsStats() - { - GlueHiveMetastore metastore = (GlueHiveMetastore) getMetastoreClient(); - GlueMetastoreStats stats = metastore.getStats(); - double initialCallCount = stats.getGetDatabases().getTime().getAllTime().getCount(); - long initialFailureCount = stats.getGetDatabases().getTotalFailures().getTotalCount(); - getMetastoreClient().getAllDatabases(); - assertThat(stats.getGetDatabases().getTime().getAllTime().getCount()).isGreaterThan(initialCallCount); - assertThat(stats.getGetDatabases().getTime().getAllTime().getAvg()).isGreaterThan(0.0); - assertEquals(stats.getGetDatabases().getTotalFailures().getTotalCount(), initialFailureCount); - } - - @Test - public void testGetDatabaseFailureLogsStats() - { - GlueHiveMetastore metastore = (GlueHiveMetastore) getMetastoreClient(); - GlueMetastoreStats stats = metastore.getStats(); - long initialFailureCount = stats.getGetDatabase().getTotalFailures().getTotalCount(); - assertThatThrownBy(() -> getMetastoreClient().getDatabase(null)) - .isInstanceOf(TrinoException.class) - .hasMessageStartingWith("Database name cannot be equal to null or empty"); - assertEquals(stats.getGetDatabase().getTotalFailures().getTotalCount(), initialFailureCount + 1); - } - - @Test - public void testGetPartitionsFilterVarChar() - throws Exception - { - TupleDomain singleEquals = new PartitionFilterBuilder() - .addStringValues(PARTITION_KEY, "2020-01-01") - .build(); - TupleDomain greaterThan = new PartitionFilterBuilder() - .addRanges(PARTITION_KEY, Range.greaterThan(VarcharType.VARCHAR, utf8Slice("2020-02-01"))) - .build(); - TupleDomain betweenInclusive = new PartitionFilterBuilder() - .addRanges(PARTITION_KEY, Range.range(VarcharType.VARCHAR, utf8Slice("2020-02-01"), true, utf8Slice("2020-03-01"), true)) - .build(); - TupleDomain greaterThanOrEquals = new PartitionFilterBuilder() - .addRanges(PARTITION_KEY, Range.greaterThanOrEqual(VarcharType.VARCHAR, utf8Slice("2020-03-01"))) - .build(); - TupleDomain inClause = new PartitionFilterBuilder() - .addStringValues(PARTITION_KEY, "2020-01-01", "2020-02-01") - .build(); - TupleDomain lessThan = new PartitionFilterBuilder() - .addRanges(PARTITION_KEY, Range.lessThan(VarcharType.VARCHAR, utf8Slice("2020-03-01"))) - .build(); - doGetPartitionsFilterTest( - CREATE_TABLE_COLUMNS_PARTITIONED_VARCHAR, - PARTITION_KEY, - VARCHAR_PARTITION_VALUES, - ImmutableList.of(singleEquals, greaterThan, betweenInclusive, greaterThanOrEquals, inClause, lessThan, TupleDomain.all()), - ImmutableList.of( - ImmutableList.of("2020-01-01"), - ImmutableList.of("2020-03-01", "2020-04-01"), - ImmutableList.of("2020-02-01", "2020-03-01"), - ImmutableList.of("2020-03-01", "2020-04-01"), - ImmutableList.of("2020-01-01", "2020-02-01"), - ImmutableList.of("2020-01-01", "2020-02-01"), - ImmutableList.of("2020-01-01", "2020-02-01", "2020-03-01", "2020-04-01"))); - } - - @Test - public void testGetPartitionsFilterBigInt() - throws Exception - { - TupleDomain singleEquals = new PartitionFilterBuilder() - .addBigintValues(PARTITION_KEY, 1000L) - .build(); - TupleDomain greaterThan = new PartitionFilterBuilder() - .addRanges(PARTITION_KEY, Range.greaterThan(BigintType.BIGINT, 100L)) - .build(); - TupleDomain betweenInclusive = new PartitionFilterBuilder() - .addRanges(PARTITION_KEY, Range.range(BigintType.BIGINT, 100L, true, 1000L, true)) - .build(); - TupleDomain greaterThanOrEquals = new PartitionFilterBuilder() - .addRanges(PARTITION_KEY, Range.greaterThanOrEqual(BigintType.BIGINT, 100L)) - .build(); - TupleDomain inClause = new PartitionFilterBuilder() - .addBigintValues(PARTITION_KEY, 1L, 1000000L) - .build(); - TupleDomain lessThan = new PartitionFilterBuilder() - .addRanges(PARTITION_KEY, Range.lessThan(BigintType.BIGINT, 1000L)) - .build(); - doGetPartitionsFilterTest( - CREATE_TABLE_COLUMNS_PARTITIONED_BIGINT, - PARTITION_KEY, - ImmutableList.of("1", "100", "1000", "1000000"), - ImmutableList.of(singleEquals, greaterThan, betweenInclusive, greaterThanOrEquals, inClause, lessThan, TupleDomain.all()), - ImmutableList.of( - ImmutableList.of("1000"), - ImmutableList.of("1000", "1000000"), - ImmutableList.of("100", "1000"), - ImmutableList.of("100", "1000", "1000000"), - ImmutableList.of("1", "1000000"), - ImmutableList.of("1", "100"), - ImmutableList.of("1", "100", "1000", "1000000"))); - } - - @Test - public void testGetPartitionsFilterInteger() - throws Exception - { - TupleDomain singleEquals = new PartitionFilterBuilder() - .addIntegerValues(PARTITION_KEY, 1000L) - .build(); - TupleDomain greaterThan = new PartitionFilterBuilder() - .addRanges(PARTITION_KEY, Range.greaterThan(IntegerType.INTEGER, 100L)) - .build(); - TupleDomain betweenInclusive = new PartitionFilterBuilder() - .addRanges(PARTITION_KEY, Range.range(IntegerType.INTEGER, 100L, true, 1000L, true)) - .build(); - TupleDomain greaterThanOrEquals = new PartitionFilterBuilder() - .addRanges(PARTITION_KEY, Range.greaterThanOrEqual(IntegerType.INTEGER, 100L)) - .build(); - TupleDomain inClause = new PartitionFilterBuilder() - .addIntegerValues(PARTITION_KEY, 1L, 1000000L) - .build(); - TupleDomain lessThan = new PartitionFilterBuilder() - .addRanges(PARTITION_KEY, Range.lessThan(IntegerType.INTEGER, 1000L)) - .build(); - doGetPartitionsFilterTest( - CREATE_TABLE_COLUMNS_PARTITIONED_INTEGER, - PARTITION_KEY, - ImmutableList.of("1", "100", "1000", "1000000"), - ImmutableList.of(singleEquals, greaterThan, betweenInclusive, greaterThanOrEquals, inClause, lessThan, TupleDomain.all()), - ImmutableList.of( - ImmutableList.of("1000"), - ImmutableList.of("1000", "1000000"), - ImmutableList.of("100", "1000"), - ImmutableList.of("100", "1000", "1000000"), - ImmutableList.of("1", "1000000"), - ImmutableList.of("1", "100"), - ImmutableList.of("1", "100", "1000", "1000000"))); - } - - @Test - public void testGetPartitionsFilterSmallInt() - throws Exception - { - TupleDomain singleEquals = new PartitionFilterBuilder() - .addSmallintValues(PARTITION_KEY, 1000L) - .build(); - TupleDomain greaterThan = new PartitionFilterBuilder() - .addRanges(PARTITION_KEY, Range.greaterThan(SmallintType.SMALLINT, 100L)) - .build(); - TupleDomain betweenInclusive = new PartitionFilterBuilder() - .addRanges(PARTITION_KEY, Range.range(SmallintType.SMALLINT, 100L, true, 1000L, true)) - .build(); - TupleDomain greaterThanOrEquals = new PartitionFilterBuilder() - .addRanges(PARTITION_KEY, Range.greaterThanOrEqual(SmallintType.SMALLINT, 100L)) - .build(); - TupleDomain inClause = new PartitionFilterBuilder() - .addSmallintValues(PARTITION_KEY, 1L, 10000L) - .build(); - TupleDomain lessThan = new PartitionFilterBuilder() - .addRanges(PARTITION_KEY, Range.lessThan(SmallintType.SMALLINT, 1000L)) - .build(); - doGetPartitionsFilterTest( - CREATE_TABLE_COLUMNS_PARTITIONED_SMALLINT, - PARTITION_KEY, - ImmutableList.of("1", "100", "1000", "10000"), - ImmutableList.of(singleEquals, greaterThan, betweenInclusive, greaterThanOrEquals, inClause, lessThan, TupleDomain.all()), - ImmutableList.of( - ImmutableList.of("1000"), - ImmutableList.of("1000", "10000"), - ImmutableList.of("100", "1000"), - ImmutableList.of("100", "1000", "10000"), - ImmutableList.of("1", "10000"), - ImmutableList.of("1", "100"), - ImmutableList.of("1", "100", "1000", "10000"))); - } - - @Test - public void testGetPartitionsFilterTinyInt() - throws Exception - { - TupleDomain singleEquals = new PartitionFilterBuilder() - .addTinyintValues(PARTITION_KEY, 127L) - .build(); - TupleDomain greaterThan = new PartitionFilterBuilder() - .addRanges(PARTITION_KEY, Range.greaterThan(TinyintType.TINYINT, 10L)) - .build(); - TupleDomain betweenInclusive = new PartitionFilterBuilder() - .addRanges(PARTITION_KEY, Range.range(TinyintType.TINYINT, 10L, true, 100L, true)) - .build(); - TupleDomain greaterThanOrEquals = new PartitionFilterBuilder() - .addRanges(PARTITION_KEY, Range.greaterThanOrEqual(TinyintType.TINYINT, 10L)) - .build(); - TupleDomain inClause = new PartitionFilterBuilder() - .addTinyintValues(PARTITION_KEY, 1L, 127L) - .build(); - TupleDomain lessThan = new PartitionFilterBuilder() - .addRanges(PARTITION_KEY, Range.lessThan(TinyintType.TINYINT, 100L)) - .build(); - doGetPartitionsFilterTest( - CREATE_TABLE_COLUMNS_PARTITIONED_TINYINT, - PARTITION_KEY, - ImmutableList.of("1", "10", "100", "127"), - ImmutableList.of(singleEquals, greaterThan, betweenInclusive, greaterThanOrEquals, inClause, lessThan, TupleDomain.all()), - ImmutableList.of( - ImmutableList.of("127"), - ImmutableList.of("100", "127"), - ImmutableList.of("10", "100"), - ImmutableList.of("10", "100", "127"), - ImmutableList.of("1", "127"), - ImmutableList.of("1", "10"), - ImmutableList.of("1", "10", "100", "127"))); - } - - @Test - public void testGetPartitionsFilterTinyIntNegatives() - throws Exception - { - TupleDomain singleEquals = new PartitionFilterBuilder() - .addTinyintValues(PARTITION_KEY, -128L) - .build(); - TupleDomain greaterThan = new PartitionFilterBuilder() - .addRanges(PARTITION_KEY, Range.greaterThan(TinyintType.TINYINT, 0L)) - .build(); - TupleDomain betweenInclusive = new PartitionFilterBuilder() - .addRanges(PARTITION_KEY, Range.range(TinyintType.TINYINT, 0L, true, 50L, true)) - .build(); - TupleDomain greaterThanOrEquals = new PartitionFilterBuilder() - .addRanges(PARTITION_KEY, Range.greaterThanOrEqual(TinyintType.TINYINT, 0L)) - .build(); - TupleDomain inClause = new PartitionFilterBuilder() - .addTinyintValues(PARTITION_KEY, 0L, -128L) - .build(); - TupleDomain lessThan = new PartitionFilterBuilder() - .addRanges(PARTITION_KEY, Range.lessThan(TinyintType.TINYINT, 0L)) - .build(); - doGetPartitionsFilterTest( - CREATE_TABLE_COLUMNS_PARTITIONED_TINYINT, - PARTITION_KEY, - ImmutableList.of("-128", "0", "50", "100"), - ImmutableList.of(singleEquals, greaterThan, betweenInclusive, greaterThanOrEquals, inClause, lessThan, TupleDomain.all()), - ImmutableList.of( - ImmutableList.of("-128"), - ImmutableList.of("100", "50"), - ImmutableList.of("0", "50"), - ImmutableList.of("0", "100", "50"), - ImmutableList.of("-128", "0"), - ImmutableList.of("-128"), - ImmutableList.of("-128", "0", "100", "50"))); - } - - @Test - public void testGetPartitionsFilterDecimal() - throws Exception - { - String value1 = "1.000"; - String value2 = "10.134"; - String value3 = "25.111"; - String value4 = "30.333"; - - TupleDomain singleEquals = new PartitionFilterBuilder() - .addDecimalValues(PARTITION_KEY, value1) - .build(); - TupleDomain greaterThan = new PartitionFilterBuilder() - .addRanges(PARTITION_KEY, Range.greaterThan(DECIMAL_TYPE, decimalOf(value2))) - .build(); - TupleDomain betweenInclusive = new PartitionFilterBuilder() - .addRanges(PARTITION_KEY, Range.range(DECIMAL_TYPE, decimalOf(value2), true, decimalOf(value3), true)) - .build(); - TupleDomain greaterThanOrEquals = new PartitionFilterBuilder() - .addRanges(PARTITION_KEY, Range.greaterThanOrEqual(DECIMAL_TYPE, decimalOf(value3))) - .build(); - TupleDomain inClause = new PartitionFilterBuilder() - .addDecimalValues(PARTITION_KEY, value1, value4) - .build(); - TupleDomain lessThan = new PartitionFilterBuilder() - .addRanges(PARTITION_KEY, Range.lessThan(DECIMAL_TYPE, decimalOf("25.5"))) - .build(); - doGetPartitionsFilterTest( - CREATE_TABLE_COLUMNS_PARTITIONED_DECIMAL, - PARTITION_KEY, - ImmutableList.of(value1, value2, value3, value4), - ImmutableList.of(singleEquals, greaterThan, betweenInclusive, greaterThanOrEquals, inClause, lessThan, TupleDomain.all()), - ImmutableList.of( - ImmutableList.of(value1), - ImmutableList.of(value3, value4), - ImmutableList.of(value2, value3), - ImmutableList.of(value3, value4), - ImmutableList.of(value1, value4), - ImmutableList.of(value1, value2, value3), - ImmutableList.of(value1, value2, value3, value4))); - } - - // we don't presently know how to properly convert a Date type into a string that is compatible with Glue. - @Test - public void testGetPartitionsFilterDate() - throws Exception - { - TupleDomain singleEquals = new PartitionFilterBuilder() - .addDateValues(PARTITION_KEY, 18000L) - .build(); - TupleDomain greaterThan = new PartitionFilterBuilder() - .addRanges(PARTITION_KEY, Range.greaterThan(DateType.DATE, 19000L)) - .build(); - TupleDomain betweenInclusive = new PartitionFilterBuilder() - .addRanges(PARTITION_KEY, Range.range(DateType.DATE, 19000L, true, 20000L, true)) - .build(); - TupleDomain greaterThanOrEquals = new PartitionFilterBuilder() - .addRanges(PARTITION_KEY, Range.greaterThanOrEqual(DateType.DATE, 19000L)) - .build(); - TupleDomain inClause = new PartitionFilterBuilder() - .addDateValues(PARTITION_KEY, 18000L, 21000L) - .build(); - TupleDomain lessThan = new PartitionFilterBuilder() - .addRanges(PARTITION_KEY, Range.lessThan(DateType.DATE, 20000L)) - .build(); - // we are unable to convert Date to a string format that Glue will accept, so it should translate to the wildcard in all cases. Commented out results are - // what we expect if we are able to do a proper conversion - doGetPartitionsFilterTest( - CREATE_TABLE_COLUMNS_PARTITIONED_DATE, - PARTITION_KEY, - ImmutableList.of("18000", "19000", "20000", "21000"), - ImmutableList.of( - singleEquals, greaterThan, betweenInclusive, greaterThanOrEquals, inClause, lessThan, TupleDomain.all()), - ImmutableList.of( -// ImmutableList.of("18000"), -// ImmutableList.of("20000", "21000"), -// ImmutableList.of("19000", "20000"), -// ImmutableList.of("19000", "20000", "21000"), -// ImmutableList.of("18000", "21000"), -// ImmutableList.of("18000", "19000"), - ImmutableList.of("18000", "19000", "20000", "21000"), - ImmutableList.of("18000", "19000", "20000", "21000"), - ImmutableList.of("18000", "19000", "20000", "21000"), - ImmutableList.of("18000", "19000", "20000", "21000"), - ImmutableList.of("18000", "19000", "20000", "21000"), - ImmutableList.of("18000", "19000", "20000", "21000"), - ImmutableList.of("18000", "19000", "20000", "21000"))); - } - - @Test - public void testGetPartitionsFilterTwoPartitionKeys() - throws Exception - { - TupleDomain equalsFilter = new PartitionFilterBuilder() - .addStringValues(PARTITION_KEY, "2020-03-01") - .addBigintValues(PARTITION_KEY2, 300L) - .build(); - TupleDomain rangeFilter = new PartitionFilterBuilder() - .addRanges(PARTITION_KEY, Range.greaterThanOrEqual(VarcharType.VARCHAR, utf8Slice("2020-02-01"))) - .addRanges(PARTITION_KEY2, Range.greaterThan(BigintType.BIGINT, 200L)) - .build(); - - doGetPartitionsFilterTest( - CREATE_TABLE_COLUMNS_PARTITIONED_TWO_KEYS, - ImmutableList.of(PARTITION_KEY, PARTITION_KEY2), - ImmutableList.of( - PartitionValues.make("2020-01-01", "100"), - PartitionValues.make("2020-02-01", "200"), - PartitionValues.make("2020-03-01", "300"), - PartitionValues.make("2020-04-01", "400")), - ImmutableList.of(equalsFilter, rangeFilter, TupleDomain.all()), - ImmutableList.of( - ImmutableList.of(PartitionValues.make("2020-03-01", "300")), - ImmutableList.of( - PartitionValues.make("2020-03-01", "300"), - PartitionValues.make("2020-04-01", "400")), - ImmutableList.of( - PartitionValues.make("2020-01-01", "100"), - PartitionValues.make("2020-02-01", "200"), - PartitionValues.make("2020-03-01", "300"), - PartitionValues.make("2020-04-01", "400")))); - } - - @Test - public void testGetPartitionsFilterMaxLengthWildcard() - throws Exception - { - // this filter string will exceed the 2048 char limit set by glue, and we expect the filter to revert to the wildcard - TupleDomain filter = new PartitionFilterBuilder() - .addStringValues(PARTITION_KEY, "x".repeat(2048)) - .build(); - - doGetPartitionsFilterTest( - CREATE_TABLE_COLUMNS_PARTITIONED_VARCHAR, - PARTITION_KEY, - VARCHAR_PARTITION_VALUES, - ImmutableList.of(filter), - ImmutableList.of( - ImmutableList.of("2020-01-01", "2020-02-01", "2020-03-01", "2020-04-01"))); - } - - @Test - public void testGetPartitionsFilterTwoPartitionKeysPartialQuery() - throws Exception - { - // we expect the second constraint to still be present and provide filtering - TupleDomain equalsFilter = new PartitionFilterBuilder() - .addStringValues(PARTITION_KEY, "x".repeat(2048)) - .addBigintValues(PARTITION_KEY2, 300L) - .build(); - - doGetPartitionsFilterTest( - CREATE_TABLE_COLUMNS_PARTITIONED_TWO_KEYS, - ImmutableList.of(PARTITION_KEY, PARTITION_KEY2), - ImmutableList.of( - PartitionValues.make("2020-01-01", "100"), - PartitionValues.make("2020-02-01", "200"), - PartitionValues.make("2020-03-01", "300"), - PartitionValues.make("2020-04-01", "400")), - ImmutableList.of(equalsFilter), - ImmutableList.of(ImmutableList.of(PartitionValues.make("2020-03-01", "300")))); - } - - @Test - public void testGetPartitionsFilterNone() - throws Exception - { - // test both a global none and that with a single column none, and a valid domain with none() - TupleDomain noneFilter = new PartitionFilterBuilder() - .addDomain(PARTITION_KEY, Domain.none(VarcharType.VARCHAR)) - .build(); - doGetPartitionsFilterTest( - CREATE_TABLE_COLUMNS_PARTITIONED_VARCHAR, - PARTITION_KEY, - VARCHAR_PARTITION_VALUES, - ImmutableList.of(TupleDomain.none(), noneFilter), - ImmutableList.of(ImmutableList.of(), ImmutableList.of())); - } - - @Test - public void testGetPartitionsFilterNotNull() - throws Exception - { - TupleDomain notNullFilter = new PartitionFilterBuilder() - .addDomain(PARTITION_KEY, Domain.notNull(VarcharType.VARCHAR)) - .build(); - doGetPartitionsFilterTest( - CREATE_TABLE_COLUMNS_PARTITIONED_VARCHAR, - PARTITION_KEY, - VARCHAR_PARTITION_VALUES, - ImmutableList.of(notNullFilter), - ImmutableList.of(ImmutableList.of("2020-01-01", "2020-02-01", "2020-03-01", "2020-04-01"))); - } - - @Test - public void testGetPartitionsFilterIsNull() - throws Exception - { - TupleDomain isNullFilter = new PartitionFilterBuilder() - .addDomain(PARTITION_KEY, Domain.onlyNull(VarcharType.VARCHAR)) - .build(); - doGetPartitionsFilterTest( - CREATE_TABLE_COLUMNS_PARTITIONED_VARCHAR, - PARTITION_KEY, - VARCHAR_PARTITION_VALUES, - ImmutableList.of(isNullFilter), - ImmutableList.of(ImmutableList.of())); - } - - @Test - public void testGetPartitionsFilterIsNullWithValue() - throws Exception - { - List partitionList = new ArrayList<>(); - partitionList.add("100"); - partitionList.add(null); - - doGetPartitionsFilterTest( - CREATE_TABLE_COLUMNS_PARTITIONED_VARCHAR, - PARTITION_KEY, - partitionList, - ImmutableList.of(new PartitionFilterBuilder() - // IS NULL - .addDomain(PARTITION_KEY, Domain.onlyNull(VarcharType.VARCHAR)) - .build()), - ImmutableList.of(ImmutableList.of(GlueExpressionUtil.NULL_STRING))); - - doGetPartitionsFilterTest( - CREATE_TABLE_COLUMNS_PARTITIONED_VARCHAR, - PARTITION_KEY, - partitionList, - ImmutableList.of(new PartitionFilterBuilder() - // IS NULL or is a specific value - .addDomain(PARTITION_KEY, Domain.create(ValueSet.of(VARCHAR, utf8Slice("100")), true)) - .build()), - ImmutableList.of(ImmutableList.of("100", GlueExpressionUtil.NULL_STRING))); - } - - @Test - public void testGetPartitionsFilterEqualsOrIsNullWithValue() - throws Exception - { - TupleDomain equalsOrIsNullFilter = new PartitionFilterBuilder() - .addStringValues(PARTITION_KEY, "2020-03-01") - .addDomain(PARTITION_KEY, Domain.onlyNull(VarcharType.VARCHAR)) - .build(); - List partitionList = new ArrayList<>(); - partitionList.add("2020-01-01"); - partitionList.add("2020-02-01"); - partitionList.add("2020-03-01"); - partitionList.add(null); - - doGetPartitionsFilterTest( - CREATE_TABLE_COLUMNS_PARTITIONED_VARCHAR, - PARTITION_KEY, - partitionList, - ImmutableList.of(equalsOrIsNullFilter), - ImmutableList.of(ImmutableList.of("2020-03-01", GlueExpressionUtil.NULL_STRING))); - } - - @Test - public void testGetPartitionsFilterIsNotNull() - throws Exception - { - TupleDomain isNotNullFilter = new PartitionFilterBuilder() - .addDomain(PARTITION_KEY, Domain.notNull(VarcharType.VARCHAR)) - .build(); - List partitionList = new ArrayList<>(); - partitionList.add("100"); - partitionList.add(null); - - doGetPartitionsFilterTest( - CREATE_TABLE_COLUMNS_PARTITIONED_VARCHAR, - PARTITION_KEY, - partitionList, - ImmutableList.of(isNotNullFilter), - ImmutableList.of(ImmutableList.of("100"))); - } - - @Test(dataProvider = "unsupportedNullPushdownTypes") - public void testGetPartitionsFilterUnsupportedIsNull(List columnMetadata, Type type, String partitionValue) - throws Exception - { - TupleDomain isNullFilter = new PartitionFilterBuilder() - .addDomain(PARTITION_KEY, Domain.onlyNull(type)) - .build(); - List partitionList = new ArrayList<>(); - partitionList.add(partitionValue); - partitionList.add(null); - - doGetPartitionsFilterTest( - columnMetadata, - PARTITION_KEY, - partitionList, - ImmutableList.of(isNullFilter), - // Currently, we get NULL partition from Glue and filter it in our side because - // (column = '__HIVE_DEFAULT_PARTITION__') on numeric types causes exception on Glue. e.g. 'input string: "__HIVE_D" is not an integer' - ImmutableList.of(ImmutableList.of(partitionValue, GlueExpressionUtil.NULL_STRING))); - } - - @Test(dataProvider = "unsupportedNullPushdownTypes") - public void testGetPartitionsFilterUnsupportedIsNotNull(List columnMetadata, Type type, String partitionValue) - throws Exception - { - TupleDomain isNotNullFilter = new PartitionFilterBuilder() - .addDomain(PARTITION_KEY, Domain.notNull(type)) - .build(); - List partitionList = new ArrayList<>(); - partitionList.add(partitionValue); - partitionList.add(null); - - doGetPartitionsFilterTest( - columnMetadata, - PARTITION_KEY, - partitionList, - ImmutableList.of(isNotNullFilter), - // Currently, we get NULL partition from Glue and filter it in our side because - // (column <> '__HIVE_DEFAULT_PARTITION__') on numeric types causes exception on Glue. e.g. 'input string: "__HIVE_D" is not an integer' - ImmutableList.of(ImmutableList.of(partitionValue, GlueExpressionUtil.NULL_STRING))); - } - - @DataProvider - public Object[][] unsupportedNullPushdownTypes() - { - return new Object[][] { - // Numeric types are unsupported for IS (NOT) NULL predicate pushdown - {CREATE_TABLE_COLUMNS_PARTITIONED_TINYINT, TinyintType.TINYINT, "127"}, - {CREATE_TABLE_COLUMNS_PARTITIONED_SMALLINT, SmallintType.SMALLINT, "32767"}, - {CREATE_TABLE_COLUMNS_PARTITIONED_INTEGER, IntegerType.INTEGER, "2147483647"}, - {CREATE_TABLE_COLUMNS_PARTITIONED_BIGINT, BigintType.BIGINT, "9223372036854775807"}, - {CREATE_TABLE_COLUMNS_PARTITIONED_DECIMAL, DECIMAL_TYPE, "12345.12345"}, - // Date and timestamp aren't numeric types, but the pushdown is unsupported because of GlueExpressionUtil.canConvertSqlTypeToStringForGlue - {CREATE_TABLE_COLUMNS_PARTITIONED_DATE, DateType.DATE, "2022-07-11"}, - {CREATE_TABLE_COLUMNS_PARTITIONED_TIMESTAMP, TimestampType.TIMESTAMP_MILLIS, "2022-07-11 01:02:03.123"}, - }; - } - - @Test - public void testGetPartitionsFilterEqualsAndIsNotNull() - throws Exception - { - TupleDomain equalsAndIsNotNullFilter = new PartitionFilterBuilder() - .addDomain(PARTITION_KEY, Domain.notNull(VarcharType.VARCHAR)) - .addBigintValues(PARTITION_KEY2, 300L) - .build(); - - doGetPartitionsFilterTest( - CREATE_TABLE_COLUMNS_PARTITIONED_TWO_KEYS, - ImmutableList.of(PARTITION_KEY, PARTITION_KEY2), - ImmutableList.of( - PartitionValues.make("2020-01-01", "100"), - PartitionValues.make("2020-02-01", "200"), - PartitionValues.make("2020-03-01", "300"), - PartitionValues.make(null, "300")), - ImmutableList.of(equalsAndIsNotNullFilter), - ImmutableList.of(ImmutableList.of(PartitionValues.make("2020-03-01", "300")))); - } - - @Test - public void testUpdateStatisticsOnCreate() - { - SchemaTableName tableName = temporaryTable("update_statistics_create"); - try (Transaction transaction = newTransaction()) { - ConnectorSession session = newSession(); - ConnectorMetadata metadata = transaction.getMetadata(); - - List columns = ImmutableList.of(new ColumnMetadata("a_column", BigintType.BIGINT)); - ConnectorTableMetadata tableMetadata = new ConnectorTableMetadata(tableName, columns, createTableProperties(TEXTFILE)); - ConnectorOutputTableHandle createTableHandle = metadata.beginCreateTable(session, tableMetadata, Optional.empty(), NO_RETRIES); - - // write data - ConnectorPageSink sink = pageSinkProvider.createPageSink(transaction.getTransactionHandle(), session, createTableHandle, TESTING_PAGE_SINK_ID); - MaterializedResult data = MaterializedResult.resultBuilder(session, BigintType.BIGINT) - .row(1L) - .row(2L) - .row(3L) - .row(4L) - .row(5L) - .build(); - sink.appendPage(data.toPage()); - Collection fragments = getFutureValue(sink.finish()); - - // prepare statistics - ComputedStatistics statistics = ComputedStatistics.builder(ImmutableList.of(), ImmutableList.of()) - .addTableStatistic(TableStatisticType.ROW_COUNT, singleValueBlock(5)) - .addColumnStatistic(MIN_VALUE.createColumnStatisticMetadata("a_column"), singleValueBlock(1)) - .addColumnStatistic(MAX_VALUE.createColumnStatisticMetadata("a_column"), singleValueBlock(1)) - .addColumnStatistic(NUMBER_OF_DISTINCT_VALUES.createColumnStatisticMetadata("a_column"), singleValueBlock(1)) - .addColumnStatistic(NUMBER_OF_NON_NULL_VALUES.createColumnStatisticMetadata("a_column"), singleValueBlock(1)) - .build(); - - // finish CTAS - metadata.finishCreateTable(session, createTableHandle, fragments, ImmutableList.of(statistics)); - transaction.commit(); - } - finally { - dropTable(tableName); - } - } - - @Test - public void testUpdatePartitionedStatisticsOnCreate() - { - SchemaTableName tableName = temporaryTable("update_partitioned_statistics_create"); - try (Transaction transaction = newTransaction()) { - ConnectorSession session = newSession(); - ConnectorMetadata metadata = transaction.getMetadata(); - - List columns = ImmutableList.of( - new ColumnMetadata("a_column", BigintType.BIGINT), - new ColumnMetadata("part_column", BigintType.BIGINT)); - - ConnectorTableMetadata tableMetadata = new ConnectorTableMetadata(tableName, columns, createTableProperties(TEXTFILE, ImmutableList.of("part_column"))); - ConnectorOutputTableHandle createTableHandle = metadata.beginCreateTable(session, tableMetadata, Optional.empty(), NO_RETRIES); - - // write data - ConnectorPageSink sink = pageSinkProvider.createPageSink(transaction.getTransactionHandle(), session, createTableHandle, TESTING_PAGE_SINK_ID); - MaterializedResult data = MaterializedResult.resultBuilder(session, BigintType.BIGINT, BigintType.BIGINT) - .row(1L, 1L) - .row(2L, 1L) - .row(3L, 1L) - .row(4L, 2L) - .row(5L, 2L) - .build(); - sink.appendPage(data.toPage()); - Collection fragments = getFutureValue(sink.finish()); - - // prepare statistics - ComputedStatistics statistics1 = ComputedStatistics.builder(ImmutableList.of("part_column"), ImmutableList.of(singleValueBlock(1))) - .addTableStatistic(TableStatisticType.ROW_COUNT, singleValueBlock(3)) - .addColumnStatistic(MIN_VALUE.createColumnStatisticMetadata("a_column"), singleValueBlock(1)) - .addColumnStatistic(MAX_VALUE.createColumnStatisticMetadata("a_column"), singleValueBlock(1)) - .addColumnStatistic(NUMBER_OF_DISTINCT_VALUES.createColumnStatisticMetadata("a_column"), singleValueBlock(1)) - .addColumnStatistic(NUMBER_OF_NON_NULL_VALUES.createColumnStatisticMetadata("a_column"), singleValueBlock(1)) - .build(); - ComputedStatistics statistics2 = ComputedStatistics.builder(ImmutableList.of("part_column"), ImmutableList.of(singleValueBlock(2))) - .addTableStatistic(TableStatisticType.ROW_COUNT, singleValueBlock(2)) - .addColumnStatistic(MIN_VALUE.createColumnStatisticMetadata("a_column"), singleValueBlock(4)) - .addColumnStatistic(MAX_VALUE.createColumnStatisticMetadata("a_column"), singleValueBlock(4)) - .addColumnStatistic(NUMBER_OF_DISTINCT_VALUES.createColumnStatisticMetadata("a_column"), singleValueBlock(4)) - .addColumnStatistic(NUMBER_OF_NON_NULL_VALUES.createColumnStatisticMetadata("a_column"), singleValueBlock(4)) - .build(); - - // finish CTAS - metadata.finishCreateTable(session, createTableHandle, fragments, ImmutableList.of(statistics1, statistics2)); - transaction.commit(); - } - finally { - dropTable(tableName); - } - } - - @Test - public void testStatisticsLargeNumberOfColumns() - throws Exception - { - SchemaTableName tableName = temporaryTable("test_statistics_large_number_of_columns"); - try { - ImmutableList.Builder columns = ImmutableList.builder(); - ImmutableMap.Builder columnStatistics = ImmutableMap.builder(); - for (int i = 1; i < 1500; ++i) { - String columnName = "t_bigint " + i + "_" + String.join("", Collections.nCopies(240, "x")); - columns.add(new ColumnMetadata(columnName, BIGINT)); - columnStatistics.put( - columnName, - createIntegerColumnStatistics( - OptionalLong.of(-1000 - i), - OptionalLong.of(1000 + i), - OptionalLong.of(i), - OptionalLong.of(2L * i))); - } - - PartitionStatistics partitionStatistics = PartitionStatistics.builder() - .setBasicStatistics(HIVE_BASIC_STATISTICS) - .setColumnStatistics(columnStatistics.buildOrThrow()).build(); - - doCreateEmptyTable(tableName, ORC, columns.build()); - testUpdateTableStatistics(tableName, ZERO_TABLE_STATISTICS, partitionStatistics); - } - finally { - dropTable(tableName); - } - } - - @Test - public void testStatisticsLongColumnNames() - throws Exception - { - SchemaTableName tableName = temporaryTable("test_statistics_long_column_name"); - try { - String columnName1 = String.join("", Collections.nCopies(255, "x")); - String columnName2 = String.join("", Collections.nCopies(255, "ӆ")); - String columnName3 = String.join("", Collections.nCopies(255, "ö")); - - List columns = List.of( - new ColumnMetadata(columnName1, BIGINT), - new ColumnMetadata(columnName2, BIGINT), - new ColumnMetadata(columnName3, BIGINT)); - - Map columnStatistics = Map.of( - columnName1, INTEGER_COLUMN_STATISTICS, - columnName2, INTEGER_COLUMN_STATISTICS, - columnName3, INTEGER_COLUMN_STATISTICS); - PartitionStatistics partitionStatistics = PartitionStatistics.builder() - .setBasicStatistics(HIVE_BASIC_STATISTICS) - .setColumnStatistics(columnStatistics).build(); - - doCreateEmptyTable(tableName, ORC, columns); - - assertThat(metastore.getTableStatistics(tableName.getSchemaName(), tableName.getTableName(), Optional.empty())) - .isEqualTo(ZERO_TABLE_STATISTICS); - testUpdateTableStatistics(tableName, ZERO_TABLE_STATISTICS, partitionStatistics); - } - finally { - dropTable(tableName); - } - } - - @Test - public void testStatisticsColumnModification() - throws Exception - { - SchemaTableName tableName = temporaryTable("test_statistics_column_modification"); - try { - List columns = List.of( - new ColumnMetadata("column1", BIGINT), - new ColumnMetadata("column2", BIGINT), - new ColumnMetadata("column3", BIGINT)); - - doCreateEmptyTable(tableName, ORC, columns); - - Map columnStatistics = Map.of( - "column1", INTEGER_COLUMN_STATISTICS, - "column2", INTEGER_COLUMN_STATISTICS); - PartitionStatistics partitionStatistics = PartitionStatistics.builder() - .setBasicStatistics(HIVE_BASIC_STATISTICS) - .setColumnStatistics(columnStatistics).build(); - - // set table statistics for column1 - metastore.updateTableStatistics( - tableName.getSchemaName(), - tableName.getTableName(), - NO_ACID_TRANSACTION, - actualStatistics -> { - assertThat(actualStatistics).isEqualTo(ZERO_TABLE_STATISTICS); - return partitionStatistics; - }); - - assertThat(metastore.getTableStatistics(tableName.getSchemaName(), tableName.getTableName(), Optional.empty())) - .isEqualTo(partitionStatistics); - - metastore.renameColumn(tableName.getSchemaName(), tableName.getTableName(), "column1", "column4"); - assertThat(metastore.getTableStatistics(tableName.getSchemaName(), tableName.getTableName(), Optional.empty())) - .isEqualTo(new PartitionStatistics( - HIVE_BASIC_STATISTICS, - Map.of("column2", INTEGER_COLUMN_STATISTICS))); - - metastore.dropColumn(tableName.getSchemaName(), tableName.getTableName(), "column2"); - assertThat(metastore.getTableStatistics(tableName.getSchemaName(), tableName.getTableName(), Optional.empty())) - .isEqualTo(new PartitionStatistics(HIVE_BASIC_STATISTICS, Map.of())); - - metastore.addColumn(tableName.getSchemaName(), tableName.getTableName(), "column5", HiveType.HIVE_INT, "comment"); - assertThat(metastore.getTableStatistics(tableName.getSchemaName(), tableName.getTableName(), Optional.empty())) - .isEqualTo(new PartitionStatistics(HIVE_BASIC_STATISTICS, Map.of())); - - // TODO: column1 stats should be removed on column delete. However this is tricky since stats can be stored in multiple partitions. - metastore.renameColumn(tableName.getSchemaName(), tableName.getTableName(), "column4", "column1"); - assertThat(metastore.getTableStatistics(tableName.getSchemaName(), tableName.getTableName(), Optional.empty())) - .isEqualTo(new PartitionStatistics( - HIVE_BASIC_STATISTICS, - Map.of("column1", INTEGER_COLUMN_STATISTICS))); - } - finally { - dropTable(tableName); - } - } - - @Test - public void testStatisticsPartitionedTableColumnModification() - throws Exception - { - SchemaTableName tableName = temporaryTable("test_partitioned_table_statistics_column_modification"); - try { - List columns = List.of( - new ColumnMetadata("column1", BIGINT), - new ColumnMetadata("column2", BIGINT), - new ColumnMetadata("ds", VARCHAR)); - - Map columnStatistics = Map.of( - "column1", INTEGER_COLUMN_STATISTICS, - "column2", INTEGER_COLUMN_STATISTICS); - PartitionStatistics partitionStatistics = PartitionStatistics.builder() - .setBasicStatistics(HIVE_BASIC_STATISTICS) - .setColumnStatistics(columnStatistics).build(); - - createDummyPartitionedTable(tableName, columns); - GlueHiveMetastore metastoreClient = (GlueHiveMetastore) getMetastoreClient(); - double countBefore = metastoreClient.getStats().getBatchUpdatePartition().getTime().getAllTime().getCount(); - - metastore.updatePartitionStatistics(tableName.getSchemaName(), tableName.getTableName(), "ds=2016-01-01", actualStatistics -> partitionStatistics); - - assertThat(metastoreClient.getStats().getBatchUpdatePartition().getTime().getAllTime().getCount()).isEqualTo(countBefore + 1); - PartitionStatistics tableStatistics = new PartitionStatistics(createEmptyStatistics(), Map.of()); - assertThat(metastore.getTableStatistics(tableName.getSchemaName(), tableName.getTableName(), Optional.empty())) - .isEqualTo(tableStatistics); - assertThat(metastore.getPartitionStatistics(tableName.getSchemaName(), tableName.getTableName(), Set.of("ds=2016-01-01"))) - .isEqualTo(Map.of("ds=2016-01-01", partitionStatistics)); - - // renaming table column does not rename partition columns - metastore.renameColumn(tableName.getSchemaName(), tableName.getTableName(), "column1", "column4"); - assertThat(metastore.getTableStatistics(tableName.getSchemaName(), tableName.getTableName(), Optional.empty())) - .isEqualTo(tableStatistics); - assertThat(metastore.getPartitionStatistics(tableName.getSchemaName(), tableName.getTableName(), Set.of("ds=2016-01-01"))) - .isEqualTo(Map.of("ds=2016-01-01", partitionStatistics)); - - // dropping table column does not drop partition columns - metastore.dropColumn(tableName.getSchemaName(), tableName.getTableName(), "column2"); - assertThat(metastore.getTableStatistics(tableName.getSchemaName(), tableName.getTableName(), Optional.empty())) - .isEqualTo(tableStatistics); - assertThat(metastore.getPartitionStatistics(tableName.getSchemaName(), tableName.getTableName(), Set.of("ds=2016-01-01"))) - .isEqualTo(Map.of("ds=2016-01-01", partitionStatistics)); - } - finally { - dropTable(tableName); - } - } - - @Test - public void testInvalidColumnStatisticsMetadata() - throws Exception - { - SchemaTableName tableName = temporaryTable("test_statistics_invalid_column_metadata"); - try { - List columns = List.of( - new ColumnMetadata("column1", BIGINT)); - - Map columnStatistics = Map.of( - "column1", INTEGER_COLUMN_STATISTICS); - PartitionStatistics partitionStatistics = PartitionStatistics.builder() - .setBasicStatistics(HIVE_BASIC_STATISTICS) - .setColumnStatistics(columnStatistics).build(); - - doCreateEmptyTable(tableName, ORC, columns); - - // set table statistics for column1 - metastore.updateTableStatistics( - tableName.getSchemaName(), - tableName.getTableName(), - NO_ACID_TRANSACTION, - actualStatistics -> { - assertThat(actualStatistics).isEqualTo(ZERO_TABLE_STATISTICS); - return partitionStatistics; - }); - - Table table = metastore.getTable(tableName.getSchemaName(), tableName.getTableName()).get(); - TableInput tableInput = GlueInputConverter.convertTable(table); - tableInput.setParameters(ImmutableMap.builder() - .putAll(tableInput.getParameters()) - .put("column_stats_bad_data", "bad data") - .buildOrThrow()); - getGlueClient().updateTable(new UpdateTableRequest() - .withDatabaseName(tableName.getSchemaName()) - .withTableInput(tableInput)); - - assertThat(metastore.getTableStatistics(tableName.getSchemaName(), tableName.getTableName(), Optional.empty())) - .isEqualTo(partitionStatistics); - } - finally { - dropTable(tableName); - } - } - - @Override - public void testPartitionColumnProperties() - { - // Glue currently does not support parameters on the partitioning columns - assertThatThrownBy(super::testPartitionColumnProperties) - .isInstanceOf(TrinoException.class) - .hasMessageStartingWith("Parameters not supported for partition columns (Service: AWSGlue; Status Code: 400; Error Code: InvalidInputException;"); - } - - @Test - public void testGlueObjectsWithoutStorageDescriptor() - { - // StorageDescriptor is an Optional field for Glue tables. - SchemaTableName table = temporaryTable("test_missing_storage_descriptor"); - DeleteTableRequest deleteTableRequest = new DeleteTableRequest() - .withDatabaseName(table.getSchemaName()) - .withName(table.getTableName()); - - try { - Supplier resetTableInput = () -> new TableInput() - .withStorageDescriptor(null) - .withName(table.getTableName()) - .withTableType(EXTERNAL_TABLE.name()); - - TableInput tableInput = resetTableInput.get(); - glueClient.createTable(new CreateTableRequest() - .withDatabaseName(database) - .withTableInput(tableInput)); - - assertThatThrownBy(() -> metastore.getTable(table.getSchemaName(), table.getTableName())) - .hasMessageStartingWith("Table StorageDescriptor is null for table"); - glueClient.deleteTable(deleteTableRequest); - - // Iceberg table - tableInput = resetTableInput.get().withParameters(ImmutableMap.of(ICEBERG_TABLE_TYPE_NAME, ICEBERG_TABLE_TYPE_VALUE)); - glueClient.createTable(new CreateTableRequest() - .withDatabaseName(database) - .withTableInput(tableInput)); - assertTrue(isIcebergTable(metastore.getTable(table.getSchemaName(), table.getTableName()).orElseThrow())); - glueClient.deleteTable(deleteTableRequest); - - // Delta Lake table - tableInput = resetTableInput.get().withParameters(ImmutableMap.of(SPARK_TABLE_PROVIDER_KEY, DELTA_LAKE_PROVIDER)); - glueClient.createTable(new CreateTableRequest() - .withDatabaseName(database) - .withTableInput(tableInput)); - assertTrue(isDeltaLakeTable(metastore.getTable(table.getSchemaName(), table.getTableName()).orElseThrow())); - glueClient.deleteTable(deleteTableRequest); - - // Iceberg materialized view - tableInput = resetTableInput.get().withTableType(VIRTUAL_VIEW.name()) - .withViewOriginalText("/* Presto Materialized View: eyJvcmlnaW5hbFNxbCI6IlNFTEVDVCAxIiwiY29sdW1ucyI6W3sibmFtZSI6ImEiLCJ0eXBlIjoiaW50ZWdlciJ9XX0= */") - .withViewExpandedText(ICEBERG_MATERIALIZED_VIEW_COMMENT) - .withParameters(ImmutableMap.of( - PRESTO_VIEW_FLAG, "true", - TABLE_COMMENT, ICEBERG_MATERIALIZED_VIEW_COMMENT)); - glueClient.createTable(new CreateTableRequest() - .withDatabaseName(database) - .withTableInput(tableInput)); - assertTrue(isTrinoMaterializedView(metastore.getTable(table.getSchemaName(), table.getTableName()).orElseThrow())); - materializedViews.add(table); - try (Transaction transaction = newTransaction()) { - ConnectorSession session = newSession(); - ConnectorMetadata metadata = transaction.getMetadata(); - // Not a view - assertThat(metadata.listViews(session, Optional.empty())) - .doesNotContain(table); - assertThat(metadata.listViews(session, Optional.of(table.getSchemaName()))) - .doesNotContain(table); - assertThat(metadata.getView(session, table)).isEmpty(); - } - finally { - materializedViews.remove(table); - } - } - finally { - // Table cannot be dropped through HiveMetastore since a TableHandle cannot be created - glueClient.deleteTable(new DeleteTableRequest() - .withDatabaseName(table.getSchemaName()) - .withName(table.getTableName())); - } - } - - private Block singleValueBlock(long value) - { - BlockBuilder blockBuilder = BIGINT.createBlockBuilder(null, 1); - BIGINT.writeLong(blockBuilder, value); - return blockBuilder.build(); - } - - private void doGetPartitionsFilterTest( - List columnMetadata, - String partitionColumnName, - List partitionStringValues, - List> filterList, - List> expectedSingleValueList) - throws Exception - { - List partitionValuesList = partitionStringValues.stream() - .map(PartitionValues::make) - .collect(toImmutableList()); - List> expectedPartitionValuesList = expectedSingleValueList.stream() - .map(expectedValue -> expectedValue.stream() - .map(PartitionValues::make) - .collect(toImmutableList())) - .collect(toImmutableList()); - doGetPartitionsFilterTest(columnMetadata, ImmutableList.of(partitionColumnName), partitionValuesList, filterList, expectedPartitionValuesList); - } - - /** - * @param filterList should be same sized list as expectedValuesList - */ - private void doGetPartitionsFilterTest( - List columnMetadata, - List partitionColumnNames, - List partitionValues, - List> filterList, - List> expectedValuesList) - throws Exception - { - try (CloseableSchamaTableName closeableTableName = new CloseableSchamaTableName(temporaryTable(("get_partitions")))) { - SchemaTableName tableName = closeableTableName.getSchemaTableName(); - createDummyPartitionedTable(tableName, columnMetadata, partitionColumnNames, partitionValues); - HiveMetastore metastoreClient = getMetastoreClient(); - - for (int i = 0; i < filterList.size(); i++) { - TupleDomain filter = filterList.get(i); - List expectedValues = expectedValuesList.get(i); - List expectedResults = expectedValues.stream() - .map(expectedPartitionValues -> makePartName(partitionColumnNames, expectedPartitionValues.getValues())) - .collect(toImmutableList()); - - Optional> partitionNames = metastoreClient.getPartitionNamesByFilter( - tableName.getSchemaName(), - tableName.getTableName(), - partitionColumnNames, - filter); - assertTrue(partitionNames.isPresent()); - assertEquals( - partitionNames.get(), - expectedResults, - format("lists \nactual: %s\nexpected: %s\nmismatch for filter %s (input index %d)\n", partitionNames.get(), expectedResults, filter, i)); - } - } - } - - private void createDummyPartitionedTable(SchemaTableName tableName, List columns, List partitionColumnNames, List partitionValues) - throws Exception - { - doCreateEmptyTable(tableName, ORC, columns, partitionColumnNames); - - HiveMetastoreClosure metastoreClient = new HiveMetastoreClosure(getMetastoreClient()); - Table table = metastoreClient.getTable(tableName.getSchemaName(), tableName.getTableName()) - .orElseThrow(() -> new TableNotFoundException(tableName)); - List partitions = new ArrayList<>(); - List partitionNames = new ArrayList<>(); - partitionValues.stream() - .map(partitionValue -> makePartName(partitionColumnNames, partitionValue.values)) - .forEach( - partitionName -> { - partitions.add(new PartitionWithStatistics(createDummyPartition(table, partitionName), partitionName, PartitionStatistics.empty())); - partitionNames.add(partitionName); - }); - metastoreClient.addPartitions(tableName.getSchemaName(), tableName.getTableName(), partitions); - partitionNames.forEach( - partitionName -> metastoreClient.updatePartitionStatistics( - tableName.getSchemaName(), tableName.getTableName(), partitionName, currentStatistics -> ZERO_TABLE_STATISTICS)); - } - - private class CloseableSchamaTableName - implements AutoCloseable - { - private final SchemaTableName schemaTableName; - - private CloseableSchamaTableName(SchemaTableName schemaTableName) - { - this.schemaTableName = schemaTableName; - } - - public SchemaTableName getSchemaTableName() - { - return schemaTableName; - } - - @Override - public void close() - { - dropTable(schemaTableName); - } - } - - // container class for readability. Each value is one for a partitionKey, in order they appear in the schema - private static class PartitionValues - { - private final List values; - - private static PartitionValues make(String... values) - { - return new PartitionValues(Arrays.asList(values)); - } - - private PartitionValues(List values) - { - // Elements are nullable - //noinspection Java9CollectionFactory - this.values = unmodifiableList(new ArrayList<>(requireNonNull(values, "values is null"))); - } - - public List getValues() - { - return values; - } - } -} diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/glue/TestingGlueHiveMetastore.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/glue/TestingGlueHiveMetastore.java new file mode 100644 index 000000000000..6c490989f11b --- /dev/null +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/glue/TestingGlueHiveMetastore.java @@ -0,0 +1,65 @@ +/* + * 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.hive.metastore.glue; + +import com.google.common.collect.ImmutableSet; +import io.trino.plugin.hive.metastore.glue.GlueHiveMetastore.TableKind; +import software.amazon.awssdk.services.glue.GlueClient; + +import java.io.IOException; +import java.net.URI; +import java.nio.file.Path; +import java.util.EnumSet; +import java.util.function.Consumer; + +import static com.google.common.base.Verify.verify; +import static io.trino.plugin.hive.HiveTestUtils.HDFS_FILE_SYSTEM_FACTORY; +import static io.trino.plugin.hive.metastore.glue.GlueMetastoreModule.createGlueClient; +import static java.nio.file.Files.createDirectories; +import static java.nio.file.Files.exists; +import static java.nio.file.Files.isDirectory; + +public final class TestingGlueHiveMetastore +{ + private TestingGlueHiveMetastore() {} + + public static GlueHiveMetastore createTestingGlueHiveMetastore(Path defaultWarehouseDir, Consumer registerResource) + { + if (!exists(defaultWarehouseDir)) { + try { + createDirectories(defaultWarehouseDir); + } + catch (IOException e) { + throw new RuntimeException("Could not create directory: %s".formatted(defaultWarehouseDir), e); + } + } + verify(isDirectory(defaultWarehouseDir), "%s is not a directory", defaultWarehouseDir); + return createTestingGlueHiveMetastore(defaultWarehouseDir.toUri(), registerResource); + } + + public static GlueHiveMetastore createTestingGlueHiveMetastore(URI warehouseUri, Consumer registerResource) + { + GlueHiveMetastoreConfig glueConfig = new GlueHiveMetastoreConfig() + .setDefaultWarehouseDir(warehouseUri.toString()); + GlueClient glueClient = createGlueClient(glueConfig, ImmutableSet.of()); + registerResource.accept(glueClient); + return new GlueHiveMetastore( + glueClient, + new GlueContext(glueConfig), + GlueCache.NOOP, + HDFS_FILE_SYSTEM_FACTORY, + glueConfig, + EnumSet.allOf(TableKind.class)); + } +} diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/glue/TestingMetastoreObjects.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/glue/TestingMetastoreObjects.java index 0594cbf8814b..d3a617f2197d 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/glue/TestingMetastoreObjects.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/glue/TestingMetastoreObjects.java @@ -32,9 +32,9 @@ import java.util.concurrent.ThreadLocalRandom; import java.util.function.Consumer; -import static io.trino.plugin.hive.HiveMetadata.TABLE_COMMENT; -import static io.trino.plugin.hive.ViewReaderUtil.ICEBERG_MATERIALIZED_VIEW_COMMENT; import static io.trino.plugin.hive.ViewReaderUtil.PRESTO_VIEW_FLAG; +import static io.trino.plugin.hive.metastore.Table.TABLE_COMMENT; +import static io.trino.plugin.hive.metastore.TableInfo.ICEBERG_MATERIALIZED_VIEW_COMMENT; import static java.lang.String.format; public final class TestingMetastoreObjects diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/recording/TestRecordingHiveMetastore.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/recording/TestRecordingHiveMetastore.java index 798087d841f4..417b23a6a55e 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/recording/TestRecordingHiveMetastore.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/recording/TestRecordingHiveMetastore.java @@ -41,10 +41,12 @@ import io.trino.plugin.hive.metastore.Storage; import io.trino.plugin.hive.metastore.StorageFormat; import io.trino.plugin.hive.metastore.Table; +import io.trino.plugin.hive.metastore.TableInfo; import io.trino.plugin.hive.metastore.UnimplementedHiveMetastore; import io.trino.plugin.hive.util.HiveBlockEncodingSerde; import io.trino.spi.block.Block; import io.trino.spi.block.TestingBlockJsonSerde; +import io.trino.spi.connector.SchemaTableName; import io.trino.spi.predicate.Domain; import io.trino.spi.predicate.TupleDomain; import io.trino.spi.security.RoleGrant; @@ -58,12 +60,12 @@ import java.util.List; import java.util.Map; import java.util.Optional; +import java.util.OptionalDouble; import java.util.OptionalLong; import java.util.Set; import java.util.concurrent.TimeUnit; import static io.trino.plugin.hive.HiveBasicStatistics.createEmptyStatistics; -import static io.trino.plugin.hive.util.HiveBucketing.BucketingVersion.BUCKETING_V1; import static io.trino.spi.security.PrincipalType.USER; import static io.trino.spi.type.VarcharType.createUnboundedVarcharType; import static org.testng.Assert.assertEquals; @@ -84,7 +86,7 @@ public class TestRecordingHiveMetastore private static final Storage TABLE_STORAGE = new Storage( StorageFormat.create("serde", "input", "output"), Optional.of("location"), - Optional.of(new HiveBucketProperty(ImmutableList.of("column"), BUCKETING_V1, 10, ImmutableList.of(new SortingColumn("column", Order.ASCENDING)))), + Optional.of(new HiveBucketProperty(ImmutableList.of("column"), 10, ImmutableList.of(new SortingColumn("column", Order.ASCENDING)))), true, ImmutableMap.of("param", "value2")); private static final Table TABLE = new Table( @@ -124,9 +126,9 @@ public class TestRecordingHiveMetastore Optional.empty(), Optional.empty(), OptionalLong.of(1234), + OptionalDouble.of(1235), OptionalLong.of(1235), - OptionalLong.of(1), - OptionalLong.of(8)))); + OptionalLong.of(1)))); private static final HivePrivilegeInfo PRIVILEGE_INFO = new HivePrivilegeInfo(HivePrivilege.SELECT, true, new HivePrincipal(USER, "grantor"), new HivePrincipal(USER, "grantee")); private static final RoleGrant ROLE_GRANT = new RoleGrant(new TrinoPrincipal(USER, "grantee"), "role", true); private static final List PARTITION_COLUMN_NAMES = ImmutableList.of(TABLE_COLUMN.getName()); @@ -175,13 +177,7 @@ private void validateMetadata(HiveMetastore hiveMetastore) assertEquals(hiveMetastore.getDatabase("database"), Optional.of(DATABASE)); assertEquals(hiveMetastore.getAllDatabases(), ImmutableList.of("database")); assertEquals(hiveMetastore.getTable("database", "table"), Optional.of(TABLE)); - assertEquals(hiveMetastore.getTableStatistics(TABLE), PARTITION_STATISTICS); - assertEquals(hiveMetastore.getPartitionStatistics(TABLE, ImmutableList.of(PARTITION, OTHER_PARTITION)), ImmutableMap.of( - "column=value", PARTITION_STATISTICS, - "column=other_value", PARTITION_STATISTICS)); - assertEquals(hiveMetastore.getAllTables("database"), ImmutableList.of("table")); - assertEquals(hiveMetastore.getTablesWithParameter("database", "param", "value3"), ImmutableList.of("table")); - assertEquals(hiveMetastore.getAllViews("database"), ImmutableList.of()); + assertEquals(hiveMetastore.getTables("database"), ImmutableList.of(new TableInfo(new SchemaTableName("database", "table"), TableInfo.ExtendedRelationType.TABLE))); assertEquals(hiveMetastore.getPartition(TABLE, ImmutableList.of("value")), Optional.of(PARTITION)); assertEquals(hiveMetastore.getPartitionNamesByFilter("database", "table", PARTITION_COLUMN_NAMES, TupleDomain.all()), Optional.of(ImmutableList.of("value"))); assertEquals(hiveMetastore.getPartitionNamesByFilter("database", "table", PARTITION_COLUMN_NAMES, TUPLE_DOMAIN), Optional.of(ImmutableList.of("value"))); @@ -191,13 +187,10 @@ private void validateMetadata(HiveMetastore hiveMetastore) assertEquals(hiveMetastore.listTablePrivileges("database", "table", Optional.of("owner"), Optional.of(new HivePrincipal(USER, "user"))), ImmutableSet.of(PRIVILEGE_INFO)); assertEquals(hiveMetastore.listRoles(), ImmutableSet.of("role")); assertEquals(hiveMetastore.listRoleGrants(new HivePrincipal(USER, "user")), ImmutableSet.of(ROLE_GRANT)); - assertEquals(hiveMetastore.listGrantedPrincipals("role"), ImmutableSet.of(ROLE_GRANT)); } private void validatePartitionSubset(HiveMetastore hiveMetastore) { - assertEquals(hiveMetastore.getPartitionStatistics(TABLE, ImmutableList.of(PARTITION)), ImmutableMap.of("column=value", PARTITION_STATISTICS)); - assertEquals(hiveMetastore.getPartitionStatistics(TABLE, ImmutableList.of(OTHER_PARTITION)), ImmutableMap.of("column=other_value", PARTITION_STATISTICS)); assertEquals(hiveMetastore.getPartitionsByNames(TABLE, ImmutableList.of("column=value")), ImmutableMap.of("column=value", Optional.of(PARTITION))); assertEquals(hiveMetastore.getPartitionsByNames(TABLE, ImmutableList.of("column=other_value")), ImmutableMap.of("column=other_value", Optional.of(OTHER_PARTITION))); } @@ -232,6 +225,15 @@ public Optional
getTable(String databaseName, String tableName) } @Override + public List getTables(String databaseName) + { + if (databaseName.equals("database")) { + return List.of(new TableInfo(TABLE.getSchemaTableName(), TableInfo.ExtendedRelationType.TABLE)); + } + + return List.of(); + } + public PartitionStatistics getTableStatistics(Table table) { if (table.getDatabaseName().equals("database") && table.getTableName().equals("table")) { @@ -241,7 +243,6 @@ public PartitionStatistics getTableStatistics(Table table) return new PartitionStatistics(createEmptyStatistics(), ImmutableMap.of()); } - @Override public Map getPartitionStatistics(Table table, List partitions) { ImmutableMap.Builder result = ImmutableMap.builder(); @@ -256,7 +257,6 @@ public Map getPartitionStatistics(Table table, List return result.buildOrThrow(); } - @Override public List getAllTables(String databaseName) { if (databaseName.equals("database")) { @@ -266,7 +266,6 @@ public List getAllTables(String databaseName) return ImmutableList.of(); } - @Override public List getTablesWithParameter(String databaseName, String parameterKey, String parameterValue) { if (databaseName.equals("database") && parameterKey.equals("param") && parameterValue.equals("value3")) { @@ -275,7 +274,6 @@ public List getTablesWithParameter(String databaseName, String parameter return ImmutableList.of(); } - @Override public List getAllViews(String databaseName) { return ImmutableList.of(); @@ -346,7 +344,6 @@ public Set listRoles() return ImmutableSet.of("role"); } - @Override public Set listGrantedPrincipals(String role) { return ImmutableSet.of(ROLE_GRANT); diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/thrift/InMemoryThriftMetastore.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/thrift/InMemoryThriftMetastore.java deleted file mode 100644 index cfb678985521..000000000000 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/thrift/InMemoryThriftMetastore.java +++ /dev/null @@ -1,718 +0,0 @@ -/* - * 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.hive.metastore.thrift; - -import com.google.common.collect.ImmutableList; -import com.google.common.collect.ImmutableMap; -import com.google.common.collect.ImmutableSet; -import com.google.errorprone.annotations.concurrent.GuardedBy; -import io.trino.hive.thrift.metastore.Database; -import io.trino.hive.thrift.metastore.FieldSchema; -import io.trino.hive.thrift.metastore.Partition; -import io.trino.hive.thrift.metastore.PrincipalPrivilegeSet; -import io.trino.hive.thrift.metastore.PrincipalType; -import io.trino.hive.thrift.metastore.Table; -import io.trino.plugin.hive.HiveColumnStatisticType; -import io.trino.plugin.hive.PartitionStatistics; -import io.trino.plugin.hive.SchemaAlreadyExistsException; -import io.trino.plugin.hive.TableAlreadyExistsException; -import io.trino.plugin.hive.acid.AcidTransaction; -import io.trino.plugin.hive.metastore.HivePrincipal; -import io.trino.plugin.hive.metastore.HivePrivilegeInfo; -import io.trino.plugin.hive.metastore.HivePrivilegeInfo.HivePrivilege; -import io.trino.plugin.hive.metastore.PartitionWithStatistics; -import io.trino.spi.TrinoException; -import io.trino.spi.connector.SchemaNotFoundException; -import io.trino.spi.connector.SchemaTableName; -import io.trino.spi.connector.TableNotFoundException; -import io.trino.spi.predicate.TupleDomain; -import io.trino.spi.security.RoleGrant; -import io.trino.spi.type.Type; -import org.apache.hadoop.fs.Path; -import org.apache.hadoop.hive.metastore.TableType; - -import java.io.File; -import java.io.IOException; -import java.io.UncheckedIOException; -import java.net.URI; -import java.util.EnumSet; -import java.util.HashMap; -import java.util.List; -import java.util.Map; -import java.util.Objects; -import java.util.Optional; -import java.util.Set; -import java.util.function.Function; - -import static com.google.common.base.MoreObjects.toStringHelper; -import static com.google.common.base.Preconditions.checkArgument; -import static com.google.common.collect.ImmutableList.toImmutableList; -import static com.google.common.collect.ImmutableSet.toImmutableSet; -import static com.google.common.io.MoreFiles.deleteRecursively; -import static com.google.common.io.RecursiveDeleteOption.ALLOW_INSECURE; -import static io.trino.plugin.hive.HiveBasicStatistics.createEmptyStatistics; -import static io.trino.plugin.hive.metastore.MetastoreUtil.partitionKeyFilterToStringList; -import static io.trino.plugin.hive.metastore.thrift.ThriftMetastoreUtil.toMetastoreApiPartition; -import static io.trino.plugin.hive.util.HiveUtil.toPartitionValues; -import static io.trino.spi.StandardErrorCode.SCHEMA_NOT_EMPTY; -import static java.util.Locale.US; -import static java.util.Objects.requireNonNull; -import static org.apache.hadoop.hive.common.FileUtils.makePartName; -import static org.apache.hadoop.hive.metastore.TableType.EXTERNAL_TABLE; -import static org.apache.hadoop.hive.metastore.TableType.MANAGED_TABLE; -import static org.apache.hadoop.hive.metastore.TableType.VIRTUAL_VIEW; - -public class InMemoryThriftMetastore - implements ThriftMetastore -{ - @GuardedBy("this") - private final Map databases = new HashMap<>(); - @GuardedBy("this") - private final Map relations = new HashMap<>(); - @GuardedBy("this") - private final Map views = new HashMap<>(); - @GuardedBy("this") - private final Map partitions = new HashMap<>(); - @GuardedBy("this") - private final Map columnStatistics = new HashMap<>(); - @GuardedBy("this") - private final Map partitionColumnStatistics = new HashMap<>(); - @GuardedBy("this") - private final Map> tablePrivileges = new HashMap<>(); - - private final File baseDirectory; - private final boolean assumeCanonicalPartitionKeys; - - public InMemoryThriftMetastore(File baseDirectory, ThriftMetastoreConfig metastoreConfig) - { - this.baseDirectory = requireNonNull(baseDirectory, "baseDirectory is null"); - this.assumeCanonicalPartitionKeys = requireNonNull(metastoreConfig).isAssumeCanonicalPartitionKeys(); - checkArgument(!baseDirectory.exists(), "Base directory already exists"); - checkArgument(baseDirectory.mkdirs(), "Could not create base directory"); - } - - @Override - public synchronized void createDatabase(Database database) - { - requireNonNull(database, "database is null"); - - File directory; - if (database.getLocationUri() != null) { - directory = new File(URI.create(database.getLocationUri())); - } - else { - // use Hive default naming convention - directory = new File(baseDirectory, database.getName() + ".db"); - database = database.deepCopy(); - database.setLocationUri(directory.toURI().toString()); - } - - checkArgument(!directory.exists(), "Database directory already exists"); - checkArgument(isParentDir(directory, baseDirectory), "Database directory must be inside of the metastore base directory"); - checkArgument(directory.mkdirs(), "Could not create database directory"); - - if (databases.putIfAbsent(database.getName(), database) != null) { - throw new SchemaAlreadyExistsException(database.getName()); - } - } - - // TODO: respect deleteData - @Override - public synchronized void dropDatabase(String databaseName, boolean deleteData) - { - if (!databases.containsKey(databaseName)) { - throw new SchemaNotFoundException(databaseName); - } - if (!getAllTables(databaseName).isEmpty()) { - throw new TrinoException(SCHEMA_NOT_EMPTY, "Schema not empty: " + databaseName); - } - databases.remove(databaseName); - } - - @Override - public synchronized void alterDatabase(String databaseName, Database newDatabase) - { - String newDatabaseName = newDatabase.getName(); - - if (databaseName.equals(newDatabaseName)) { - if (databases.replace(databaseName, newDatabase) == null) { - throw new SchemaNotFoundException(databaseName); - } - return; - } - - Database database = databases.get(databaseName); - if (database == null) { - throw new SchemaNotFoundException(databaseName); - } - if (databases.putIfAbsent(newDatabaseName, database) != null) { - throw new SchemaAlreadyExistsException(newDatabaseName); - } - databases.remove(databaseName); - - rewriteKeys(relations, name -> new SchemaTableName(newDatabaseName, name.getTableName())); - rewriteKeys(views, name -> new SchemaTableName(newDatabaseName, name.getTableName())); - rewriteKeys(partitions, name -> name.withSchemaName(newDatabaseName)); - rewriteKeys(tablePrivileges, name -> name.withDatabase(newDatabaseName)); - } - - @Override - public synchronized List getAllDatabases() - { - return ImmutableList.copyOf(databases.keySet()); - } - - @Override - public synchronized void createTable(Table table) - { - TableType tableType = TableType.valueOf(table.getTableType()); - checkArgument(EnumSet.of(MANAGED_TABLE, EXTERNAL_TABLE, VIRTUAL_VIEW).contains(tableType), "Invalid table type: %s", tableType); - - if (tableType == VIRTUAL_VIEW) { - checkArgument(table.getSd().getLocation() == null, "Storage location for view must be null"); - } - else { - File directory = new File(new Path(table.getSd().getLocation()).toUri()); - checkArgument(directory.exists(), "Table directory does not exist"); - if (tableType == MANAGED_TABLE) { - checkArgument(isParentDir(directory, baseDirectory), "Table directory must be inside of the metastore base directory"); - } - } - - SchemaTableName schemaTableName = new SchemaTableName(table.getDbName(), table.getTableName()); - Table tableCopy = table.deepCopy(); - - if (relations.putIfAbsent(schemaTableName, tableCopy) != null) { - throw new TableAlreadyExistsException(schemaTableName); - } - - if (tableType == VIRTUAL_VIEW) { - views.put(schemaTableName, tableCopy); - } - - PrincipalPrivilegeSet privileges = table.getPrivileges(); - if (privileges != null) { - throw new UnsupportedOperationException(); - } - } - - @Override - public synchronized void dropTable(String databaseName, String tableName, boolean deleteData) - { - List locations = listAllDataPaths(this, databaseName, tableName); - - SchemaTableName schemaTableName = new SchemaTableName(databaseName, tableName); - Table table = relations.remove(schemaTableName); - if (table == null) { - throw new TableNotFoundException(schemaTableName); - } - views.remove(schemaTableName); - partitions.keySet().removeIf(partitionName -> partitionName.matches(databaseName, tableName)); - - // remove data - if (deleteData && table.getTableType().equals(MANAGED_TABLE.name())) { - for (String location : locations) { - if (location != null) { - File directory = new File(new Path(location).toUri()); - checkArgument(isParentDir(directory, baseDirectory), "Table directory must be inside of the metastore base directory"); - deleteDirectory(directory); - } - } - } - } - - private static List listAllDataPaths(ThriftMetastore metastore, String schemaName, String tableName) - { - ImmutableList.Builder locations = ImmutableList.builder(); - Table table = metastore.getTable(schemaName, tableName).get(); - if (table.getSd().getLocation() != null) { - // For unpartitioned table, there should be nothing directly under this directory. - // But including this location in the set makes the directory content assert more - // extensive, which is desirable. - locations.add(table.getSd().getLocation()); - } - List partitionColumnNames = table.getPartitionKeys().stream() - .map(FieldSchema::getName) - .collect(toImmutableList()); - Optional> partitionNames = metastore.getPartitionNamesByFilter(schemaName, tableName, partitionColumnNames, TupleDomain.all()); - if (partitionNames.isPresent()) { - metastore.getPartitionsByNames(schemaName, tableName, partitionNames.get()).stream() - .map(partition -> partition.getSd().getLocation()) - .filter(location -> !location.startsWith(table.getSd().getLocation())) - .forEach(locations::add); - } - - return locations.build(); - } - - @Override - public synchronized void alterTable(String databaseName, String tableName, Table newTable) - { - SchemaTableName oldName = new SchemaTableName(databaseName, tableName); - SchemaTableName newName = new SchemaTableName(newTable.getDbName(), newTable.getTableName()); - - // if the name did not change, this is a simple schema change - if (oldName.equals(newName)) { - if (relations.replace(oldName, newTable) == null) { - throw new TableNotFoundException(oldName); - } - return; - } - - // remove old table definition and add the new one - Table table = relations.get(oldName); - if (table == null) { - throw new TableNotFoundException(oldName); - } - - if (relations.putIfAbsent(newName, newTable) != null) { - throw new TableAlreadyExistsException(newName); - } - relations.remove(oldName); - } - - @Override - public void alterTransactionalTable(Table table, long transactionId, long writeId) - { - alterTable(table.getDbName(), table.getTableName(), table); - } - - @Override - public synchronized List getAllTables(String databaseName) - { - ImmutableList.Builder tables = ImmutableList.builder(); - for (SchemaTableName schemaTableName : this.relations.keySet()) { - if (schemaTableName.getSchemaName().equals(databaseName)) { - tables.add(schemaTableName.getTableName()); - } - } - return tables.build(); - } - - @Override - public synchronized List getTablesWithParameter(String databaseName, String parameterKey, String parameterValue) - { - requireNonNull(parameterKey, "parameterKey is null"); - requireNonNull(parameterValue, "parameterValue is null"); - - return relations.entrySet().stream() - .filter(entry -> entry.getKey().getSchemaName().equals(databaseName) - && parameterValue.equals(entry.getValue().getParameters().get(parameterKey))) - .map(entry -> entry.getKey().getTableName()) - .collect(toImmutableList()); - } - - @Override - public synchronized List getAllViews(String databaseName) - { - ImmutableList.Builder tables = ImmutableList.builder(); - for (SchemaTableName schemaTableName : this.views.keySet()) { - if (schemaTableName.getSchemaName().equals(databaseName)) { - tables.add(schemaTableName.getTableName()); - } - } - return tables.build(); - } - - @Override - public synchronized Optional> getAllTables() - { - return Optional.of(ImmutableList.copyOf(relations.keySet())); - } - - @Override - public synchronized Optional> getAllViews() - { - return Optional.of(ImmutableList.copyOf(views.keySet())); - } - - @Override - public synchronized Optional getDatabase(String databaseName) - { - return Optional.ofNullable(databases.get(databaseName)); - } - - @Override - public synchronized void addPartitions(String databaseName, String tableName, List partitionsWithStatistics) - { - for (PartitionWithStatistics partitionWithStatistics : partitionsWithStatistics) { - Partition partition = toMetastoreApiPartition(partitionWithStatistics.getPartition()); - if (partition.getParameters() == null) { - partition.setParameters(ImmutableMap.of()); - } - PartitionName partitionKey = PartitionName.partition(databaseName, tableName, partitionWithStatistics.getPartitionName()); - partitions.put(partitionKey, partition); - partitionColumnStatistics.put(partitionKey, partitionWithStatistics.getStatistics()); - } - } - - @Override - public synchronized void dropPartition(String databaseName, String tableName, List parts, boolean deleteData) - { - partitions.entrySet().removeIf(entry -> - entry.getKey().matches(databaseName, tableName) && entry.getValue().getValues().equals(parts)); - } - - @Override - public synchronized void alterPartition(String databaseName, String tableName, PartitionWithStatistics partitionWithStatistics) - { - Partition partition = toMetastoreApiPartition(partitionWithStatistics.getPartition()); - if (partition.getParameters() == null) { - partition.setParameters(ImmutableMap.of()); - } - PartitionName partitionKey = PartitionName.partition(databaseName, tableName, partitionWithStatistics.getPartitionName()); - partitions.put(partitionKey, partition); - partitionColumnStatistics.put(partitionKey, partitionWithStatistics.getStatistics()); - } - - @Override - public synchronized Optional getPartition(String databaseName, String tableName, List partitionValues) - { - PartitionName name = PartitionName.partition(databaseName, tableName, partitionValues); - Partition partition = partitions.get(name); - if (partition == null) { - return Optional.empty(); - } - return Optional.of(partition.deepCopy()); - } - - @Override - public synchronized Optional> getPartitionNamesByFilter(String databaseName, String tableName, List columnNames, TupleDomain partitionKeysFilter) - { - Optional> parts = partitionKeyFilterToStringList(columnNames, partitionKeysFilter, assumeCanonicalPartitionKeys); - - if (parts.isEmpty()) { - return Optional.of(ImmutableList.of()); - } - return Optional.of(partitions.entrySet().stream() - .filter(entry -> partitionMatches(entry.getValue(), databaseName, tableName, parts.get())) - .map(entry -> entry.getKey().getPartitionName()) - .collect(toImmutableList())); - } - - private static boolean partitionMatches(Partition partition, String databaseName, String tableName, List parts) - { - if (!partition.getDbName().equals(databaseName) || - !partition.getTableName().equals(tableName)) { - return false; - } - List values = partition.getValues(); - if (values.size() != parts.size()) { - return false; - } - for (int i = 0; i < values.size(); i++) { - String part = parts.get(i); - if (!part.isEmpty() && !values.get(i).equals(part)) { - return false; - } - } - return true; - } - - @Override - public synchronized List getPartitionsByNames(String databaseName, String tableName, List partitionNames) - { - ImmutableList.Builder builder = ImmutableList.builder(); - for (String name : partitionNames) { - PartitionName partitionName = PartitionName.partition(databaseName, tableName, name); - Partition partition = partitions.get(partitionName); - if (partition == null) { - return ImmutableList.of(); - } - builder.add(partition.deepCopy()); - } - return builder.build(); - } - - @Override - public synchronized Optional
getTable(String databaseName, String tableName) - { - SchemaTableName schemaTableName = new SchemaTableName(databaseName, tableName); - return Optional.ofNullable(relations.get(schemaTableName)); - } - - @Override - public Set getSupportedColumnStatistics(Type type) - { - return ThriftMetastoreUtil.getSupportedColumnStatistics(type); - } - - @Override - public synchronized PartitionStatistics getTableStatistics(Table table) - { - return getTableStatistics(table.getDbName(), table.getTableName()); - } - - private synchronized PartitionStatistics getTableStatistics(String databaseName, String tableName) - { - SchemaTableName schemaTableName = new SchemaTableName(databaseName, tableName); - PartitionStatistics statistics = columnStatistics.get(schemaTableName); - if (statistics == null) { - statistics = new PartitionStatistics(createEmptyStatistics(), ImmutableMap.of()); - } - return statistics; - } - - @Override - public synchronized Map getPartitionStatistics(Table table, List partitions) - { - List partitionColumns = table.getPartitionKeys().stream() - .map(FieldSchema::getName) - .collect(toImmutableList()); - Set partitionNames = partitions.stream() - .map(partition -> makePartName(partitionColumns, partition.getValues())) - .collect(toImmutableSet()); - return getPartitionStatistics(table.getDbName(), table.getTableName(), partitionNames); - } - - private synchronized Map getPartitionStatistics(String databaseName, String tableName, Set partitionNames) - { - ImmutableMap.Builder result = ImmutableMap.builder(); - for (String partitionName : partitionNames) { - PartitionName partitionKey = PartitionName.partition(databaseName, tableName, partitionName); - PartitionStatistics statistics = partitionColumnStatistics.get(partitionKey); - if (statistics == null) { - statistics = new PartitionStatistics(createEmptyStatistics(), ImmutableMap.of()); - } - result.put(partitionName, statistics); - } - return result.buildOrThrow(); - } - - @Override - public synchronized void updateTableStatistics(String databaseName, String tableName, AcidTransaction transaction, Function update) - { - columnStatistics.put(new SchemaTableName(databaseName, tableName), update.apply(getTableStatistics(databaseName, tableName))); - } - - @Override - public synchronized void updatePartitionStatistics(Table table, String partitionName, Function update) - { - PartitionName partitionKey = PartitionName.partition(table.getDbName(), table.getTableName(), partitionName); - partitionColumnStatistics.put(partitionKey, update.apply(getPartitionStatistics(table.getDbName(), table.getTableName(), ImmutableSet.of(partitionName)).get(partitionName))); - } - - @Override - public void createRole(String role, String grantor) - { - throw new UnsupportedOperationException(); - } - - @Override - public void dropRole(String role) - { - throw new UnsupportedOperationException(); - } - - @Override - public Set listRoles() - { - throw new UnsupportedOperationException(); - } - - @Override - public void grantRoles(Set roles, Set grantees, boolean adminOption, HivePrincipal grantor) - { - throw new UnsupportedOperationException(); - } - - @Override - public void revokeRoles(Set roles, Set grantees, boolean adminOption, HivePrincipal grantor) - { - throw new UnsupportedOperationException(); - } - - @Override - public Set listGrantedPrincipals(String role) - { - throw new UnsupportedOperationException(); - } - - @Override - public Set listRoleGrants(HivePrincipal principal) - { - throw new UnsupportedOperationException(); - } - - @Override - public Set listTablePrivileges(String databaseName, String tableName, Optional tableOwner, Optional principal) - { - return ImmutableSet.of(); - } - - @Override - public void grantTablePrivileges(String databaseName, String tableName, String tableOwner, HivePrincipal grantee, HivePrincipal grantor, Set privileges, boolean grantOption) - { - throw new UnsupportedOperationException(); - } - - @Override - public void revokeTablePrivileges(String databaseName, String tableName, String tableOwner, HivePrincipal grantee, HivePrincipal grantor, Set privileges, boolean grantOption) - { - throw new UnsupportedOperationException(); - } - - private static boolean isParentDir(File directory, File baseDirectory) - { - for (File parent = directory.getParentFile(); parent != null; parent = parent.getParentFile()) { - if (parent.equals(baseDirectory)) { - return true; - } - } - return false; - } - - private static void deleteDirectory(File dir) - { - try { - deleteRecursively(dir.toPath(), ALLOW_INSECURE); - } - catch (IOException e) { - throw new UncheckedIOException(e); - } - } - - private static class PartitionName - { - private final String schemaName; - private final String tableName; - private final List partitionValues; - private final String partitionName; // does not participate in equals and hashValue - - private PartitionName(String schemaName, String tableName, List partitionValues, String partitionName) - { - this.schemaName = schemaName.toLowerCase(US); - this.tableName = tableName.toLowerCase(US); - this.partitionValues = requireNonNull(partitionValues, "partitionValues is null"); - this.partitionName = partitionName; - } - - public static PartitionName partition(String schemaName, String tableName, String partitionName) - { - return new PartitionName(schemaName.toLowerCase(US), tableName.toLowerCase(US), toPartitionValues(partitionName), partitionName); - } - - public static PartitionName partition(String schemaName, String tableName, List partitionValues) - { - return new PartitionName(schemaName.toLowerCase(US), tableName.toLowerCase(US), partitionValues, null); - } - - public String getPartitionName() - { - return requireNonNull(partitionName, "partitionName is null"); - } - - public boolean matches(String schemaName, String tableName) - { - return this.schemaName.equals(schemaName) && - this.tableName.equals(tableName); - } - - public PartitionName withSchemaName(String schemaName) - { - return new PartitionName(schemaName, tableName, partitionValues, partitionName); - } - - @Override - public int hashCode() - { - return Objects.hash(schemaName, tableName, partitionValues); - } - - @Override - public boolean equals(Object obj) - { - if (this == obj) { - return true; - } - if (obj == null || getClass() != obj.getClass()) { - return false; - } - PartitionName other = (PartitionName) obj; - return Objects.equals(this.schemaName, other.schemaName) - && Objects.equals(this.tableName, other.tableName) - && Objects.equals(this.partitionValues, other.partitionValues); - } - - @Override - public String toString() - { - return schemaName + "/" + tableName + "/" + partitionName; - } - } - - private static class PrincipalTableKey - { - private final String principalName; - private final PrincipalType principalType; - private final String database; - private final String table; - - public PrincipalTableKey(String principalName, PrincipalType principalType, String table, String database) - { - this.principalName = requireNonNull(principalName, "principalName is null"); - this.principalType = requireNonNull(principalType, "principalType is null"); - this.table = requireNonNull(table, "table is null"); - this.database = requireNonNull(database, "database is null"); - } - - public PrincipalTableKey withDatabase(String database) - { - return new PrincipalTableKey(principalName, principalType, table, database); - } - - @Override - public boolean equals(Object o) - { - if (this == o) { - return true; - } - if (o == null || getClass() != o.getClass()) { - return false; - } - PrincipalTableKey that = (PrincipalTableKey) o; - return Objects.equals(principalName, that.principalName) && - principalType == that.principalType && - Objects.equals(table, that.table) && - Objects.equals(database, that.database); - } - - @Override - public int hashCode() - { - return Objects.hash(principalName, principalType, table, database); - } - - @Override - public String toString() - { - return toStringHelper(this) - .add("principalName", principalName) - .add("principalType", principalType) - .add("table", table) - .add("database", database) - .toString(); - } - } - - private static void rewriteKeys(Map map, Function keyRewriter) - { - for (K key : ImmutableSet.copyOf(map.keySet())) { - K newKey = keyRewriter.apply(key); - if (!newKey.equals(key)) { - map.put(newKey, map.remove(key)); - } - } - } -} diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/thrift/MetastoreClientAdapterProvider.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/thrift/MetastoreClientAdapterProvider.java new file mode 100644 index 000000000000..479448cc9a25 --- /dev/null +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/thrift/MetastoreClientAdapterProvider.java @@ -0,0 +1,19 @@ +/* + * 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.hive.metastore.thrift; + +public interface MetastoreClientAdapterProvider +{ + ThriftMetastoreClient createThriftMetastoreClientAdapter(ThriftMetastoreClient delegate); +} diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/thrift/MockThriftMetastoreClient.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/thrift/MockThriftMetastoreClient.java index 50c4a4f3140c..9822cf389f9a 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/thrift/MockThriftMetastoreClient.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/thrift/MockThriftMetastoreClient.java @@ -19,6 +19,7 @@ import com.google.common.collect.Maps; import io.trino.hive.thrift.metastore.ColumnStatisticsData; import io.trino.hive.thrift.metastore.ColumnStatisticsObj; +import io.trino.hive.thrift.metastore.DataOperationType; import io.trino.hive.thrift.metastore.Database; import io.trino.hive.thrift.metastore.EnvironmentContext; import io.trino.hive.thrift.metastore.FieldSchema; @@ -36,7 +37,7 @@ import io.trino.hive.thrift.metastore.SerDeInfo; import io.trino.hive.thrift.metastore.StorageDescriptor; import io.trino.hive.thrift.metastore.Table; -import io.trino.plugin.hive.acid.AcidOperation; +import io.trino.hive.thrift.metastore.TableMeta; import io.trino.spi.connector.SchemaTableName; import io.trino.testng.services.ManageTestResources; import org.apache.hadoop.hive.metastore.TableType; @@ -48,11 +49,13 @@ import java.util.List; import java.util.Map; import java.util.Optional; +import java.util.Set; import java.util.concurrent.atomic.AtomicInteger; import static com.google.common.collect.ImmutableList.toImmutableList; import static io.trino.hive.thrift.metastore.PrincipalType.ROLE; import static io.trino.hive.thrift.metastore.PrincipalType.USER; +import static io.trino.plugin.hive.TableType.MANAGED_TABLE; @ManageTestResources.Suppress(because = "close() is no-op and instance's resources are negligible") public class MockThriftMetastoreClient @@ -157,7 +160,6 @@ public List getAllDatabases() return ImmutableList.of(TEST_DATABASE); } - @Override public List getAllTables(String dbName) { accessCount.incrementAndGet(); @@ -170,7 +172,6 @@ public List getAllTables(String dbName) return ImmutableList.of(TEST_TABLE); } - @Override public Optional> getAllTables() throws TException { @@ -181,19 +182,16 @@ public Optional> getAllTables() return Optional.of(ImmutableList.of(new SchemaTableName(TEST_DATABASE, TEST_TABLE))); } - @Override public List getAllViews(String databaseName) { throw new UnsupportedOperationException(); } - @Override public Optional> getAllViews() { throw new UnsupportedOperationException(); } - @Override public List getTablesWithParameter(String databaseName, String parameterKey, String parameterValue) { throw new UnsupportedOperationException(); @@ -213,6 +211,25 @@ public Database getDatabase(String name) return new Database(TEST_DATABASE, null, null, null); } + @Override + public List getTableMeta(String databaseName) + { + accessCount.incrementAndGet(); + if (throwException) { + throw new RuntimeException(); + } + if (!databaseName.equals(TEST_DATABASE)) { + return ImmutableList.of(); // As specified by Hive specification + } + return ImmutableList.of(new TableMeta(TEST_DATABASE, TEST_TABLE, MANAGED_TABLE.name())); + } + + @Override + public List getTableNamesWithParameters(String databaseName, String parameterKey, Set parameterValues) + { + throw new UnsupportedOperationException(); + } + @Override public Table getTable(String dbName, String tableName) throws TException @@ -502,7 +519,6 @@ public void revokeRole(String role, String granteeName, PrincipalType granteeTyp // No-op } - @Override public List listGrantedPrincipals(String role) { throw new UnsupportedOperationException(); @@ -585,7 +601,8 @@ public void alterPartitions(String dbName, String tableName, List par } @Override - public void addDynamicPartitions(String dbName, String tableName, List partitionNames, long transactionId, long writeId, AcidOperation operation) + public void addDynamicPartitions(String dbName, String tableName, List partitionNames, long transactionId, long writeId, DataOperationType operation) + throws TException { throw new UnsupportedOperationException(); } diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/thrift/MockThriftMetastoreClientFactory.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/thrift/MockThriftMetastoreClientFactory.java index 17abe7b24c39..c7154072a24d 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/thrift/MockThriftMetastoreClientFactory.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/thrift/MockThriftMetastoreClientFactory.java @@ -26,16 +26,16 @@ public class MockThriftMetastoreClientFactory implements ThriftMetastoreClientFactory { - private final Map> clients; + private final Map> clients; public MockThriftMetastoreClientFactory(Map> clients) { this.clients = clients.entrySet().stream() - .collect(toImmutableMap(entry -> createHostAndPort(entry.getKey()), Map.Entry::getValue)); + .collect(toImmutableMap(entry -> URI.create(entry.getKey()), Map.Entry::getValue)); } @Override - public ThriftMetastoreClient create(HostAndPort address, Optional delegationToken) + public ThriftMetastoreClient create(URI address, Optional delegationToken) throws TTransportException { checkArgument(delegationToken.isEmpty(), "delegation token is not supported"); diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/thrift/TestHiveMetastoreAccessOperations.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/thrift/TestHiveMetastoreAccessOperations.java index e61c06faf9a4..0866fb485be2 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/thrift/TestHiveMetastoreAccessOperations.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/thrift/TestHiveMetastoreAccessOperations.java @@ -28,15 +28,15 @@ import java.io.File; -import static io.trino.plugin.hive.metastore.CountingAccessHiveMetastore.Method.CREATE_TABLE; -import static io.trino.plugin.hive.metastore.CountingAccessHiveMetastore.Method.GET_DATABASE; -import static io.trino.plugin.hive.metastore.CountingAccessHiveMetastore.Method.GET_PARTITIONS_BY_NAMES; -import static io.trino.plugin.hive.metastore.CountingAccessHiveMetastore.Method.GET_PARTITION_NAMES_BY_FILTER; -import static io.trino.plugin.hive.metastore.CountingAccessHiveMetastore.Method.GET_PARTITION_STATISTICS; -import static io.trino.plugin.hive.metastore.CountingAccessHiveMetastore.Method.GET_TABLE; -import static io.trino.plugin.hive.metastore.CountingAccessHiveMetastore.Method.GET_TABLE_STATISTICS; -import static io.trino.plugin.hive.metastore.CountingAccessHiveMetastore.Method.UPDATE_PARTITION_STATISTICS; -import static io.trino.plugin.hive.metastore.CountingAccessHiveMetastore.Method.UPDATE_TABLE_STATISTICS; +import static io.trino.plugin.hive.metastore.MetastoreMethod.CREATE_TABLE; +import static io.trino.plugin.hive.metastore.MetastoreMethod.GET_DATABASE; +import static io.trino.plugin.hive.metastore.MetastoreMethod.GET_PARTITIONS_BY_NAMES; +import static io.trino.plugin.hive.metastore.MetastoreMethod.GET_PARTITION_NAMES_BY_FILTER; +import static io.trino.plugin.hive.metastore.MetastoreMethod.GET_PARTITION_STATISTICS; +import static io.trino.plugin.hive.metastore.MetastoreMethod.GET_TABLE; +import static io.trino.plugin.hive.metastore.MetastoreMethod.GET_TABLE_STATISTICS; +import static io.trino.plugin.hive.metastore.MetastoreMethod.UPDATE_PARTITION_STATISTICS; +import static io.trino.plugin.hive.metastore.MetastoreMethod.UPDATE_TABLE_STATISTICS; import static io.trino.plugin.hive.metastore.file.TestingFileHiveMetastore.createTestingFileHiveMetastore; import static io.trino.testing.TestingSession.testSessionBuilder; @@ -60,7 +60,7 @@ protected QueryRunner createQueryRunner() File baseDir = queryRunner.getCoordinator().getBaseDataDir().resolve("hive").toFile(); metastore = new CountingAccessHiveMetastore(createTestingFileHiveMetastore(baseDir)); - queryRunner.installPlugin(new TestingHivePlugin(metastore)); + queryRunner.installPlugin(new TestingHivePlugin(baseDir.toPath(), metastore)); queryRunner.createCatalog("hive", "hive", ImmutableMap.of()); queryRunner.execute("CREATE SCHEMA test_schema"); @@ -109,7 +109,7 @@ public void testSelectPartitionedTable() assertMetastoreInvocations("SELECT * FROM test_select_partition", ImmutableMultiset.builder() - .addCopies(GET_TABLE, 2) + .add(GET_TABLE) .add(GET_PARTITION_NAMES_BY_FILTER) .add(GET_PARTITIONS_BY_NAMES) .build()); @@ -117,7 +117,7 @@ public void testSelectPartitionedTable() assertUpdate("INSERT INTO test_select_partition SELECT 2 AS data, 20 AS part", 1); assertMetastoreInvocations("SELECT * FROM test_select_partition", ImmutableMultiset.builder() - .addCopies(GET_TABLE, 2) + .add(GET_TABLE) .add(GET_PARTITION_NAMES_BY_FILTER) .add(GET_PARTITIONS_BY_NAMES) .build()); @@ -125,7 +125,7 @@ public void testSelectPartitionedTable() // Specify a specific partition assertMetastoreInvocations("SELECT * FROM test_select_partition WHERE part = 10", ImmutableMultiset.builder() - .addCopies(GET_TABLE, 2) + .add(GET_TABLE) .add(GET_PARTITION_NAMES_BY_FILTER) .add(GET_PARTITIONS_BY_NAMES) .build()); @@ -187,7 +187,7 @@ public void testSelfJoin() assertMetastoreInvocations("SELECT child.age, parent.age FROM test_self_join_table child JOIN test_self_join_table parent ON child.parent = parent.id", ImmutableMultiset.builder() .add(GET_TABLE) - .add(GET_TABLE_STATISTICS) + .addCopies(GET_TABLE_STATISTICS, 2) .build()); } @@ -246,7 +246,7 @@ public void testAnalyzePartitionedTable() assertMetastoreInvocations("ANALYZE test_analyze_partition", ImmutableMultiset.builder() - .addCopies(GET_TABLE, 2) + .add(GET_TABLE) .add(GET_PARTITION_NAMES_BY_FILTER) .add(GET_PARTITIONS_BY_NAMES) .add(GET_PARTITION_STATISTICS) @@ -257,7 +257,7 @@ public void testAnalyzePartitionedTable() assertMetastoreInvocations("ANALYZE test_analyze_partition", ImmutableMultiset.builder() - .addCopies(GET_TABLE, 2) + .add(GET_TABLE) .add(GET_PARTITION_NAMES_BY_FILTER) .add(GET_PARTITIONS_BY_NAMES) .add(GET_PARTITION_STATISTICS) @@ -284,7 +284,7 @@ public void testDropStatsPartitionedTable() assertMetastoreInvocations("CALL system.drop_stats('test_schema', 'drop_stats_partition')", ImmutableMultiset.builder() - .addCopies(GET_TABLE, 2) + .add(GET_TABLE) .add(GET_PARTITION_NAMES_BY_FILTER) .add(UPDATE_PARTITION_STATISTICS) .build()); @@ -293,7 +293,7 @@ public void testDropStatsPartitionedTable() assertMetastoreInvocations("CALL system.drop_stats('test_schema', 'drop_stats_partition')", ImmutableMultiset.builder() - .addCopies(GET_TABLE, 2) + .add(GET_TABLE) .add(GET_PARTITION_NAMES_BY_FILTER) .addCopies(UPDATE_PARTITION_STATISTICS, 2) .build()); diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/thrift/TestHiveMetastoreMetadataQueriesAccessOperations.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/thrift/TestHiveMetastoreMetadataQueriesAccessOperations.java index 91f881655ddd..8e373d4bd33d 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/thrift/TestHiveMetastoreMetadataQueriesAccessOperations.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/thrift/TestHiveMetastoreMetadataQueriesAccessOperations.java @@ -15,39 +15,38 @@ import com.google.common.collect.HashMultiset; import com.google.common.collect.ImmutableList; -import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableMultiset; import com.google.common.collect.Multiset; +import io.trino.Session; +import io.trino.plugin.hive.HiveQueryRunner; import io.trino.plugin.hive.HiveType; -import io.trino.plugin.hive.TestingHivePlugin; import io.trino.plugin.hive.metastore.Column; import io.trino.plugin.hive.metastore.CountingAccessHiveMetastore; -import io.trino.plugin.hive.metastore.CountingAccessHiveMetastore.Method; import io.trino.plugin.hive.metastore.CountingAccessHiveMetastoreUtil; +import io.trino.plugin.hive.metastore.HivePrincipal; +import io.trino.plugin.hive.metastore.HivePrivilegeInfo; +import io.trino.plugin.hive.metastore.MetastoreMethod; import io.trino.plugin.hive.metastore.Table; +import io.trino.plugin.hive.metastore.TableInfo; import io.trino.plugin.hive.metastore.UnimplementedHiveMetastore; import io.trino.spi.connector.SchemaTableName; +import io.trino.spi.security.RoleGrant; import io.trino.testing.AbstractTestQueryFramework; -import io.trino.testing.DataProviders; -import io.trino.testing.DistributedQueryRunner; import io.trino.testing.QueryRunner; import org.intellij.lang.annotations.Language; -import org.testng.annotations.BeforeMethod; import org.testng.annotations.Test; import java.util.List; import java.util.Optional; +import java.util.Set; import java.util.stream.IntStream; import static com.google.common.collect.ImmutableList.toImmutableList; import static io.trino.plugin.hive.HiveStorageFormat.ORC; import static io.trino.plugin.hive.TableType.MANAGED_TABLE; -import static io.trino.plugin.hive.metastore.CountingAccessHiveMetastore.Method.GET_ALL_DATABASES; -import static io.trino.plugin.hive.metastore.CountingAccessHiveMetastore.Method.GET_ALL_TABLES; -import static io.trino.plugin.hive.metastore.CountingAccessHiveMetastore.Method.GET_ALL_TABLES_FROM_DATABASE; -import static io.trino.plugin.hive.metastore.CountingAccessHiveMetastore.Method.GET_ALL_VIEWS; -import static io.trino.plugin.hive.metastore.CountingAccessHiveMetastore.Method.GET_ALL_VIEWS_FROM_DATABASE; -import static io.trino.plugin.hive.metastore.CountingAccessHiveMetastore.Method.GET_TABLE; +import static io.trino.plugin.hive.metastore.MetastoreMethod.GET_ALL_DATABASES; +import static io.trino.plugin.hive.metastore.MetastoreMethod.GET_ALL_TABLES_FROM_DATABASE; +import static io.trino.plugin.hive.metastore.MetastoreMethod.GET_TABLE; import static io.trino.plugin.hive.metastore.StorageFormat.fromHiveStorageFormat; import static io.trino.testing.TestingSession.testSessionBuilder; import static org.assertj.core.api.Assertions.assertThat; @@ -64,33 +63,30 @@ public class TestHiveMetastoreMetadataQueriesAccessOperations private MockHiveMetastore mockMetastore; private CountingAccessHiveMetastore metastore; + private static final Session SESSION = testSessionBuilder() + .setCatalog("hive") + .setSchema(Optional.empty()) + .build(); + @Override protected QueryRunner createQueryRunner() throws Exception { - DistributedQueryRunner queryRunner = DistributedQueryRunner.builder( - testSessionBuilder() - .setCatalog("hive") - .setSchema(Optional.empty()) - .build()) + mockMetastore = new MockHiveMetastore(); + metastore = new CountingAccessHiveMetastore(mockMetastore); + + QueryRunner queryRunner = HiveQueryRunner.builder(SESSION) // metadata queries do not use workers .setNodeCount(1) .addCoordinatorProperty("optimizer.experimental-max-prefetched-information-schema-prefixes", Integer.toString(MAX_PREFIXES_COUNT)) + .addHiveProperty("hive.hive-views.enabled", "true") + .setCreateTpchSchemas(false) + .setMetastore(runner -> metastore) .build(); - mockMetastore = new MockHiveMetastore(); - metastore = new CountingAccessHiveMetastore(mockMetastore); - queryRunner.installPlugin(new TestingHivePlugin(metastore)); - queryRunner.createCatalog("hive", "hive", ImmutableMap.of()); return queryRunner; } - @BeforeMethod - public void resetMetastoreSetup() - { - mockMetastore.setAllTablesViewsImplemented(false); - } - @Test public void testSelectSchemasWithoutPredicate() { @@ -112,31 +108,18 @@ public void testSelectSchemasWithLikeOverSchemaName() assertMetastoreInvocations("SELECT * FROM system.jdbc.schemas WHERE table_schem LIKE 'test%'", ImmutableMultiset.of(GET_ALL_DATABASES)); } - @Test(dataProviderClass = DataProviders.class, dataProvider = "trueFalse") - public void testSelectTablesWithoutPredicate(boolean allTablesViewsImplemented) + @Test + public void testSelectTablesWithoutPredicate() { - mockMetastore.setAllTablesViewsImplemented(allTablesViewsImplemented); assertMetastoreInvocations("SELECT * FROM information_schema.tables", - allTablesViewsImplemented - ? ImmutableMultiset.builder() - .add(GET_ALL_TABLES) - .add(GET_ALL_VIEWS) - .build() - : ImmutableMultiset.builder() + ImmutableMultiset.builder() .add(GET_ALL_DATABASES) .addCopies(GET_ALL_TABLES_FROM_DATABASE, TEST_SCHEMAS_COUNT) - .addCopies(GET_ALL_VIEWS_FROM_DATABASE, TEST_SCHEMAS_COUNT) .build()); assertMetastoreInvocations("SELECT * FROM system.jdbc.tables", - allTablesViewsImplemented - ? ImmutableMultiset.builder() - .add(GET_ALL_TABLES) - .add(GET_ALL_VIEWS) - .build() - : ImmutableMultiset.builder() + ImmutableMultiset.builder() .add(GET_ALL_DATABASES) .addCopies(GET_ALL_TABLES_FROM_DATABASE, TEST_SCHEMAS_COUNT) - .addCopies(GET_ALL_VIEWS_FROM_DATABASE, TEST_SCHEMAS_COUNT) .build()); } @@ -153,145 +136,80 @@ public void testSelectTablesWithFilterBySchema() assertMetastoreInvocations("SELECT * FROM information_schema.tables WHERE table_schema = 'test_schema_0'", ImmutableMultiset.builder() .add(GET_ALL_TABLES_FROM_DATABASE) - .add(GET_ALL_VIEWS_FROM_DATABASE) .build()); assertMetastoreInvocations("SELECT * FROM system.jdbc.tables WHERE table_schem = 'test_schema_0'", ImmutableMultiset.builder() .add(GET_ALL_TABLES_FROM_DATABASE) - .add(GET_ALL_VIEWS_FROM_DATABASE) .build()); } - @Test(dataProviderClass = DataProviders.class, dataProvider = "trueFalse") - public void testSelectTablesWithLikeOverSchema(boolean allTablesViewsImplemented) + @Test + public void testSelectTablesWithLikeOverSchema() { - mockMetastore.setAllTablesViewsImplemented(allTablesViewsImplemented); assertMetastoreInvocations("SELECT * FROM information_schema.tables WHERE table_schema LIKE 'test%'", - allTablesViewsImplemented - ? ImmutableMultiset.builder() - .add(GET_ALL_DATABASES) - .add(GET_ALL_TABLES) - .add(GET_ALL_VIEWS) - .build() - : ImmutableMultiset.builder() + ImmutableMultiset.builder() .add(GET_ALL_DATABASES) .addCopies(GET_ALL_TABLES_FROM_DATABASE, TEST_SCHEMAS_COUNT) - .addCopies(GET_ALL_VIEWS_FROM_DATABASE, TEST_SCHEMAS_COUNT) .build()); assertMetastoreInvocations("SELECT * FROM system.jdbc.tables WHERE table_schem LIKE 'test%'", - allTablesViewsImplemented - ? ImmutableMultiset.builder() - .add(GET_ALL_TABLES) - .add(GET_ALL_VIEWS) - .build() - : ImmutableMultiset.builder() + ImmutableMultiset.builder() .add(GET_ALL_DATABASES) .addCopies(GET_ALL_TABLES_FROM_DATABASE, TEST_SCHEMAS_COUNT) - .addCopies(GET_ALL_VIEWS_FROM_DATABASE, TEST_SCHEMAS_COUNT) .build()); } - @Test(dataProviderClass = DataProviders.class, dataProvider = "trueFalse") - public void testSelectTablesWithFilterByTableName(boolean allTablesViewsImplemented) + @Test + public void testSelectTablesWithFilterByTableName() { - mockMetastore.setAllTablesViewsImplemented(allTablesViewsImplemented); assertMetastoreInvocations("SELECT * FROM information_schema.tables WHERE table_name = 'test_table_0'", - allTablesViewsImplemented - ? ImmutableMultiset.builder() - .add(GET_ALL_DATABASES) - .add(GET_ALL_TABLES) - .add(GET_ALL_VIEWS) - .build() - : ImmutableMultiset.builder() + ImmutableMultiset.builder() .add(GET_ALL_DATABASES) .addCopies(GET_ALL_TABLES_FROM_DATABASE, TEST_SCHEMAS_COUNT) - .addCopies(GET_ALL_VIEWS_FROM_DATABASE, TEST_SCHEMAS_COUNT) .build()); assertMetastoreInvocations("SELECT * FROM system.jdbc.tables WHERE table_name = 'test_table_0'", - allTablesViewsImplemented - ? ImmutableMultiset.builder() - .add(GET_ALL_TABLES) - .add(GET_ALL_VIEWS) - .build() - : ImmutableMultiset.builder() + ImmutableMultiset.builder() .add(GET_ALL_DATABASES) .addCopies(GET_ALL_TABLES_FROM_DATABASE, TEST_SCHEMAS_COUNT) - .addCopies(GET_ALL_VIEWS_FROM_DATABASE, TEST_SCHEMAS_COUNT) .build()); assertMetastoreInvocations("SELECT * FROM system.jdbc.tables WHERE table_name LIKE 'test\\_table\\_0' ESCAPE '\\'", - allTablesViewsImplemented - ? ImmutableMultiset.builder() - .add(GET_ALL_TABLES) - .add(GET_ALL_VIEWS) - .build() - : ImmutableMultiset.builder() + ImmutableMultiset.builder() .add(GET_ALL_DATABASES) .addCopies(GET_ALL_TABLES_FROM_DATABASE, TEST_SCHEMAS_COUNT) - .addCopies(GET_ALL_VIEWS_FROM_DATABASE, TEST_SCHEMAS_COUNT) .build()); assertMetastoreInvocations("SELECT * FROM system.jdbc.tables WHERE table_name LIKE 'test_table_0' ESCAPE '\\'", - allTablesViewsImplemented - ? ImmutableMultiset.builder() - .add(GET_ALL_TABLES) - .add(GET_ALL_VIEWS) - .build() - : ImmutableMultiset.builder() + ImmutableMultiset.builder() .add(GET_ALL_DATABASES) .addCopies(GET_ALL_TABLES_FROM_DATABASE, TEST_SCHEMAS_COUNT) - .addCopies(GET_ALL_VIEWS_FROM_DATABASE, TEST_SCHEMAS_COUNT) .build()); } - @Test(dataProviderClass = DataProviders.class, dataProvider = "trueFalse") - public void testSelectTablesWithLikeOverTableName(boolean allTablesViewsImplemented) + @Test + public void testSelectTablesWithLikeOverTableName() { - mockMetastore.setAllTablesViewsImplemented(allTablesViewsImplemented); assertMetastoreInvocations("SELECT * FROM information_schema.tables WHERE table_name LIKE 'test%'", - allTablesViewsImplemented - ? ImmutableMultiset.builder() - .add(GET_ALL_DATABASES) - .add(GET_ALL_TABLES) - .add(GET_ALL_VIEWS) - .build() - : ImmutableMultiset.builder() + ImmutableMultiset.builder() .add(GET_ALL_DATABASES) .addCopies(GET_ALL_TABLES_FROM_DATABASE, TEST_SCHEMAS_COUNT) - .addCopies(GET_ALL_VIEWS_FROM_DATABASE, TEST_SCHEMAS_COUNT) .build()); assertMetastoreInvocations("SELECT * FROM system.jdbc.tables WHERE table_name LIKE 'test%'", - allTablesViewsImplemented - ? ImmutableMultiset.builder() - .add(GET_ALL_TABLES) - .add(GET_ALL_VIEWS) - .build() - : ImmutableMultiset.builder() + ImmutableMultiset.builder() .add(GET_ALL_DATABASES) .addCopies(GET_ALL_TABLES_FROM_DATABASE, TEST_SCHEMAS_COUNT) - .addCopies(GET_ALL_VIEWS_FROM_DATABASE, TEST_SCHEMAS_COUNT) .build()); } - @Test(dataProviderClass = DataProviders.class, dataProvider = "trueFalse") - public void testSelectViewsWithoutPredicate(boolean allTablesViewsImplemented) + @Test + public void testSelectViewsWithoutPredicate() { - mockMetastore.setAllTablesViewsImplemented(allTablesViewsImplemented); assertMetastoreInvocations("SELECT * FROM information_schema.views", - allTablesViewsImplemented - ? ImmutableMultiset.of(GET_ALL_VIEWS) - : ImmutableMultiset.builder() + ImmutableMultiset.builder() .add(GET_ALL_DATABASES) - .addCopies(GET_ALL_VIEWS_FROM_DATABASE, TEST_SCHEMAS_COUNT) + .addCopies(GET_ALL_TABLES_FROM_DATABASE, TEST_SCHEMAS_COUNT) .build()); assertMetastoreInvocations("SELECT * FROM system.jdbc.tables WHERE table_type = 'VIEW'", - allTablesViewsImplemented - ? ImmutableMultiset.builder() - .add(GET_ALL_TABLES) - .add(GET_ALL_VIEWS) - .build() - : ImmutableMultiset.builder() + ImmutableMultiset.builder() .add(GET_ALL_DATABASES) .addCopies(GET_ALL_TABLES_FROM_DATABASE, TEST_SCHEMAS_COUNT) - .addCopies(GET_ALL_VIEWS_FROM_DATABASE, TEST_SCHEMAS_COUNT) .build()); } @@ -305,123 +223,72 @@ public void testSelectViewsWithFilterByInformationSchema() @Test public void testSelectViewsWithFilterBySchema() { - assertMetastoreInvocations("SELECT * FROM information_schema.views WHERE table_schema = 'test_schema_0'", ImmutableMultiset.of(GET_ALL_VIEWS_FROM_DATABASE)); + assertMetastoreInvocations("SELECT * FROM information_schema.views WHERE table_schema = 'test_schema_0'", ImmutableMultiset.of(GET_ALL_TABLES_FROM_DATABASE)); assertMetastoreInvocations("SELECT * FROM system.jdbc.tables WHERE table_type = 'VIEW' AND table_schem = 'test_schema_0'", ImmutableMultiset.builder() .add(GET_ALL_TABLES_FROM_DATABASE) - .add(GET_ALL_VIEWS_FROM_DATABASE) .build()); } - @Test(dataProviderClass = DataProviders.class, dataProvider = "trueFalse") - public void testSelectViewsWithLikeOverSchema(boolean allTablesViewsImplemented) + @Test + public void testSelectViewsWithLikeOverSchema() { - mockMetastore.setAllTablesViewsImplemented(allTablesViewsImplemented); assertMetastoreInvocations("SELECT * FROM information_schema.views WHERE table_schema LIKE 'test%'", - allTablesViewsImplemented - ? ImmutableMultiset.builder() - .add(GET_ALL_DATABASES) - .add(GET_ALL_VIEWS) - .build() - : ImmutableMultiset.builder() + ImmutableMultiset.builder() .add(GET_ALL_DATABASES) - .addCopies(GET_ALL_VIEWS_FROM_DATABASE, TEST_SCHEMAS_COUNT) + .addCopies(GET_ALL_TABLES_FROM_DATABASE, TEST_SCHEMAS_COUNT) .build()); assertMetastoreInvocations("SELECT * FROM system.jdbc.tables WHERE table_type = 'VIEW' AND table_schem LIKE 'test%'", - allTablesViewsImplemented - ? ImmutableMultiset.builder() - .add(GET_ALL_TABLES) - .add(GET_ALL_VIEWS) - .build() - : ImmutableMultiset.builder() + ImmutableMultiset.builder() .add(GET_ALL_DATABASES) .addCopies(GET_ALL_TABLES_FROM_DATABASE, TEST_SCHEMAS_COUNT) - .addCopies(GET_ALL_VIEWS_FROM_DATABASE, TEST_SCHEMAS_COUNT) .build()); } - @Test(dataProviderClass = DataProviders.class, dataProvider = "trueFalse") - public void testSelectViewsWithFilterByTableName(boolean allTablesViewsImplemented) + @Test + public void testSelectViewsWithFilterByTableName() { - mockMetastore.setAllTablesViewsImplemented(allTablesViewsImplemented); assertMetastoreInvocations("SELECT * FROM information_schema.views WHERE table_name = 'test_table_0'", - allTablesViewsImplemented - ? ImmutableMultiset.builder() - .add(GET_ALL_DATABASES) - .add(GET_ALL_VIEWS) - .build() - : ImmutableMultiset.builder() + ImmutableMultiset.builder() .add(GET_ALL_DATABASES) - .addCopies(GET_ALL_VIEWS_FROM_DATABASE, TEST_SCHEMAS_COUNT) + .addCopies(GET_ALL_TABLES_FROM_DATABASE, TEST_SCHEMAS_COUNT) .build()); assertMetastoreInvocations("SELECT * FROM system.jdbc.tables WHERE table_type = 'VIEW' AND table_name = 'test_table_0'", - allTablesViewsImplemented - ? ImmutableMultiset.builder() - .add(GET_ALL_TABLES) - .add(GET_ALL_VIEWS) - .build() - : ImmutableMultiset.builder() + ImmutableMultiset.builder() .add(GET_ALL_DATABASES) .addCopies(GET_ALL_TABLES_FROM_DATABASE, TEST_SCHEMAS_COUNT) - .addCopies(GET_ALL_VIEWS_FROM_DATABASE, TEST_SCHEMAS_COUNT) .build()); } - @Test(dataProviderClass = DataProviders.class, dataProvider = "trueFalse") - public void testSelectViewsWithLikeOverTableName(boolean allTablesViewsImplemented) + @Test + public void testSelectViewsWithLikeOverTableName() { - mockMetastore.setAllTablesViewsImplemented(allTablesViewsImplemented); + // mockMetastore.setAllTablesViewsImplemented(allTablesViewsImplemented); assertMetastoreInvocations("SELECT * FROM information_schema.views WHERE table_name LIKE 'test%'", - allTablesViewsImplemented - ? ImmutableMultiset.builder() - .add(GET_ALL_DATABASES) - .add(GET_ALL_VIEWS) - .build() - : ImmutableMultiset.builder() + ImmutableMultiset.builder() .add(GET_ALL_DATABASES) - .addCopies(GET_ALL_VIEWS_FROM_DATABASE, TEST_SCHEMAS_COUNT) + .addCopies(GET_ALL_TABLES_FROM_DATABASE, TEST_SCHEMAS_COUNT) .build()); assertMetastoreInvocations("SELECT * FROM system.jdbc.tables WHERE table_type = 'VIEW' AND table_name LIKE 'test%'", - allTablesViewsImplemented - ? ImmutableMultiset.builder() - .add(GET_ALL_TABLES) - .add(GET_ALL_VIEWS) - .build() - : ImmutableMultiset.builder() + ImmutableMultiset.builder() .add(GET_ALL_DATABASES) .addCopies(GET_ALL_TABLES_FROM_DATABASE, TEST_SCHEMAS_COUNT) - .addCopies(GET_ALL_VIEWS_FROM_DATABASE, TEST_SCHEMAS_COUNT) .build()); } - @Test(dataProviderClass = DataProviders.class, dataProvider = "trueFalse") - public void testSelectColumnsWithoutPredicate(boolean allTablesViewsImplemented) + @Test + public void testSelectColumnsWithoutPredicate() { - mockMetastore.setAllTablesViewsImplemented(allTablesViewsImplemented); assertMetastoreInvocations("SELECT * FROM information_schema.columns", - allTablesViewsImplemented - ? ImmutableMultiset.builder() - .add(GET_ALL_TABLES) - .add(GET_ALL_VIEWS) - .addCopies(GET_TABLE, TEST_ALL_TABLES_COUNT) - .build() - : ImmutableMultiset.builder() + ImmutableMultiset.builder() .add(GET_ALL_DATABASES) .addCopies(GET_ALL_TABLES_FROM_DATABASE, TEST_SCHEMAS_COUNT) - .addCopies(GET_ALL_VIEWS_FROM_DATABASE, TEST_SCHEMAS_COUNT) .addCopies(GET_TABLE, TEST_ALL_TABLES_COUNT) .build()); assertMetastoreInvocations("SELECT * FROM system.jdbc.columns", - allTablesViewsImplemented - ? ImmutableMultiset.builder() - .add(GET_ALL_TABLES) - .add(GET_ALL_VIEWS) - .addCopies(GET_TABLE, TEST_ALL_TABLES_COUNT) - .build() - : ImmutableMultiset.builder() + ImmutableMultiset.builder() .add(GET_ALL_DATABASES) .addCopies(GET_ALL_TABLES_FROM_DATABASE, TEST_SCHEMAS_COUNT) - .addCopies(GET_ALL_VIEWS_FROM_DATABASE, TEST_SCHEMAS_COUNT) .addCopies(GET_TABLE, TEST_ALL_TABLES_COUNT) .build()); } @@ -439,46 +306,33 @@ public void testSelectColumnsFilterBySchema() assertMetastoreInvocations("SELECT * FROM information_schema.columns WHERE table_schema = 'test_schema_0'", ImmutableMultiset.builder() .add(GET_ALL_TABLES_FROM_DATABASE) - .add(GET_ALL_VIEWS_FROM_DATABASE) .addCopies(GET_TABLE, TEST_TABLES_IN_SCHEMA_COUNT) .build()); assertMetastoreInvocations("SELECT * FROM system.jdbc.columns WHERE table_schem = 'test_schema_0'", ImmutableMultiset.builder() .add(GET_ALL_TABLES_FROM_DATABASE) - .add(GET_ALL_VIEWS_FROM_DATABASE) .addCopies(GET_TABLE, TEST_TABLES_IN_SCHEMA_COUNT) .build()); assertMetastoreInvocations("SELECT * FROM system.jdbc.columns WHERE table_schem LIKE 'test\\_schema\\_0' ESCAPE '\\'", ImmutableMultiset.builder() .add(GET_ALL_TABLES_FROM_DATABASE) - .add(GET_ALL_VIEWS_FROM_DATABASE) .addCopies(GET_TABLE, TEST_TABLES_IN_SCHEMA_COUNT) .build()); assertMetastoreInvocations("SELECT * FROM system.jdbc.columns WHERE table_schem LIKE 'test_schema_0' ESCAPE '\\'", ImmutableMultiset.builder() .add(GET_ALL_DATABASES) .add(GET_ALL_TABLES_FROM_DATABASE) - .add(GET_ALL_VIEWS_FROM_DATABASE) .addCopies(GET_TABLE, TEST_TABLES_IN_SCHEMA_COUNT) .build()); } - @Test(dataProviderClass = DataProviders.class, dataProvider = "trueFalse") - public void testSelectColumnsWithLikeOverSchema(boolean allTablesViewsImplemented) + @Test + public void testSelectColumnsWithLikeOverSchema() { - mockMetastore.setAllTablesViewsImplemented(allTablesViewsImplemented); assertMetastoreInvocations("SELECT * FROM information_schema.columns WHERE table_schema LIKE 'test%'", - allTablesViewsImplemented - ? ImmutableMultiset.builder() - .add(GET_ALL_DATABASES) - .add(GET_ALL_TABLES) - .add(GET_ALL_VIEWS) - .addCopies(GET_TABLE, TEST_ALL_TABLES_COUNT) - .build() - : ImmutableMultiset.builder() + ImmutableMultiset.builder() .add(GET_ALL_DATABASES) .addCopies(GET_ALL_TABLES_FROM_DATABASE, TEST_SCHEMAS_COUNT) - .addCopies(GET_ALL_VIEWS_FROM_DATABASE, TEST_SCHEMAS_COUNT) .addCopies(GET_TABLE, TEST_ALL_TABLES_COUNT) .build()); assertMetastoreInvocations("SELECT * FROM system.jdbc.columns WHERE table_schem LIKE 'test%'", @@ -489,13 +343,12 @@ public void testSelectColumnsWithLikeOverSchema(boolean allTablesViewsImplemente .build()); } - @Test(dataProviderClass = DataProviders.class, dataProvider = "trueFalse") - public void testSelectColumnsFilterByTableName(boolean allTablesViewsImplemented) + @Test + public void testSelectColumnsFilterByTableName() { - mockMetastore.setAllTablesViewsImplemented(allTablesViewsImplemented); metastore.resetCounters(); computeActual("SELECT * FROM information_schema.columns WHERE table_name = 'test_table_0'"); - Multiset invocations = metastore.getMethodInvocations(); + Multiset invocations = metastore.getMethodInvocations(); assertThat(invocations.count(GET_TABLE)).as("GET_TABLE invocations") // some lengthy explanatory comment why variable count @@ -506,16 +359,9 @@ public void testSelectColumnsFilterByTableName(boolean allTablesViewsImplemented invocations.elementSet().remove(GET_TABLE); assertThat(invocations).as("invocations except of GET_TABLE") - .isEqualTo(allTablesViewsImplemented - ? ImmutableMultiset.builder() - .add(GET_ALL_DATABASES) - .add(GET_ALL_TABLES) - .add(GET_ALL_VIEWS) - .build() - : ImmutableMultiset.builder() + .isEqualTo(ImmutableMultiset.builder() .add(GET_ALL_DATABASES) .addCopies(GET_ALL_TABLES_FROM_DATABASE, TEST_SCHEMAS_COUNT) - .addCopies(GET_ALL_VIEWS_FROM_DATABASE, TEST_SCHEMAS_COUNT) .build()); assertMetastoreInvocations("SELECT * FROM system.jdbc.columns WHERE table_name = 'test_table_0'", @@ -536,22 +382,13 @@ public void testSelectColumnsFilterByTableName(boolean allTablesViewsImplemented .build()); } - @Test(dataProviderClass = DataProviders.class, dataProvider = "trueFalse") - public void testSelectColumnsWithLikeOverTableName(boolean allTablesViewsImplemented) + @Test + public void testSelectColumnsWithLikeOverTableName() { - mockMetastore.setAllTablesViewsImplemented(allTablesViewsImplemented); assertMetastoreInvocations("SELECT * FROM information_schema.columns WHERE table_name LIKE 'test%'", - allTablesViewsImplemented - ? ImmutableMultiset.builder() - .add(GET_ALL_DATABASES) - .add(GET_ALL_TABLES) - .add(GET_ALL_VIEWS) - .addCopies(GET_TABLE, TEST_ALL_TABLES_COUNT) - .build() - : ImmutableMultiset.builder() + ImmutableMultiset.builder() .add(GET_ALL_DATABASES) .addCopies(GET_ALL_TABLES_FROM_DATABASE, TEST_SCHEMAS_COUNT) - .addCopies(GET_ALL_VIEWS_FROM_DATABASE, TEST_SCHEMAS_COUNT) .addCopies(GET_TABLE, TEST_ALL_TABLES_COUNT) .build()); assertMetastoreInvocations("SELECT * FROM system.jdbc.columns WHERE table_name LIKE 'test%'", @@ -562,67 +399,36 @@ public void testSelectColumnsWithLikeOverTableName(boolean allTablesViewsImpleme .build()); } - @Test(dataProviderClass = DataProviders.class, dataProvider = "trueFalse") - public void testSelectColumnsFilterByColumn(boolean allTablesViewsImplemented) + @Test + public void testSelectColumnsFilterByColumn() { - mockMetastore.setAllTablesViewsImplemented(allTablesViewsImplemented); assertMetastoreInvocations("SELECT * FROM information_schema.columns WHERE column_name = 'name'", - allTablesViewsImplemented - ? ImmutableMultiset.builder() - .add(GET_ALL_TABLES) - .add(GET_ALL_VIEWS) - .addCopies(GET_TABLE, TEST_ALL_TABLES_COUNT) - .build() - : ImmutableMultiset.builder() + ImmutableMultiset.builder() .add(GET_ALL_DATABASES) .addCopies(GET_ALL_TABLES_FROM_DATABASE, TEST_SCHEMAS_COUNT) - .addCopies(GET_ALL_VIEWS_FROM_DATABASE, TEST_SCHEMAS_COUNT) .addCopies(GET_TABLE, TEST_ALL_TABLES_COUNT) .build()); assertMetastoreInvocations("SELECT * FROM system.jdbc.columns WHERE column_name = 'name'", - allTablesViewsImplemented - ? ImmutableMultiset.builder() - .add(GET_ALL_TABLES) - .add(GET_ALL_VIEWS) - .addCopies(GET_TABLE, TEST_ALL_TABLES_COUNT) - .build() - : ImmutableMultiset.builder() + ImmutableMultiset.builder() .add(GET_ALL_DATABASES) .addCopies(GET_ALL_TABLES_FROM_DATABASE, TEST_SCHEMAS_COUNT) - .addCopies(GET_ALL_VIEWS_FROM_DATABASE, TEST_SCHEMAS_COUNT) .addCopies(GET_TABLE, TEST_ALL_TABLES_COUNT) .build()); } - @Test(dataProviderClass = DataProviders.class, dataProvider = "trueFalse") - public void testSelectColumnsWithLikeOverColumn(boolean allTablesViewsImplemented) + @Test + public void testSelectColumnsWithLikeOverColumn() { - mockMetastore.setAllTablesViewsImplemented(allTablesViewsImplemented); assertMetastoreInvocations("SELECT * FROM information_schema.columns WHERE column_name LIKE 'n%'", - allTablesViewsImplemented - ? ImmutableMultiset.builder() - .add(GET_ALL_DATABASES) - .add(GET_ALL_TABLES) - .add(GET_ALL_VIEWS) - .addCopies(GET_TABLE, TEST_ALL_TABLES_COUNT) - .build() - : ImmutableMultiset.builder() + ImmutableMultiset.builder() .add(GET_ALL_DATABASES) .addCopies(GET_ALL_TABLES_FROM_DATABASE, TEST_SCHEMAS_COUNT) - .addCopies(GET_ALL_VIEWS_FROM_DATABASE, TEST_SCHEMAS_COUNT) .addCopies(GET_TABLE, TEST_ALL_TABLES_COUNT) .build()); assertMetastoreInvocations("SELECT * FROM system.jdbc.columns WHERE column_name LIKE 'n%'", - allTablesViewsImplemented - ? ImmutableMultiset.builder() - .add(GET_ALL_TABLES) - .add(GET_ALL_VIEWS) - .addCopies(GET_TABLE, TEST_ALL_TABLES_COUNT) - .build() - : ImmutableMultiset.builder() + ImmutableMultiset.builder() .add(GET_ALL_DATABASES) .addCopies(GET_ALL_TABLES_FROM_DATABASE, TEST_SCHEMAS_COUNT) - .addCopies(GET_ALL_VIEWS_FROM_DATABASE, TEST_SCHEMAS_COUNT) .addCopies(GET_TABLE, TEST_ALL_TABLES_COUNT) .build()); } @@ -668,13 +474,11 @@ public List getAllDatabases() return SCHEMAS; } - @Override public List getAllTables(String databaseName) { return TABLES_PER_SCHEMA; } - @Override public Optional> getAllTables() { if (allTablesViewsImplemented) { @@ -683,13 +487,11 @@ public Optional> getAllTables() return Optional.empty(); } - @Override public List getAllViews(String databaseName) { return ImmutableList.of(); } - @Override public Optional> getAllViews() { if (allTablesViewsImplemented) { @@ -715,9 +517,24 @@ public Optional
getTable(String databaseName, String tableName) .build()); } - public void setAllTablesViewsImplemented(boolean allTablesViewsImplemented) + @Override + public List getTables(String databaseName) + { + return TABLES_PER_SCHEMA.stream() + .map(name -> new TableInfo(new SchemaTableName(databaseName, name), TableInfo.ExtendedRelationType.TABLE)) + .toList(); + } + + @Override + public Set listTablePrivileges(String databaseName, String tableName, Optional tableOwner, Optional prestoPrincipal) + { + return Set.of(); + } + + @Override + public Set listRoleGrants(HivePrincipal principal) { - this.allTablesViewsImplemented = allTablesViewsImplemented; + return Set.of(); } } } diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/thrift/TestThriftHiveMetastoreClient.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/thrift/TestThriftHiveMetastoreClient.java index 1c5a7951e2be..ad01bd1122cf 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/thrift/TestThriftHiveMetastoreClient.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/thrift/TestThriftHiveMetastoreClient.java @@ -19,6 +19,7 @@ import org.apache.thrift.transport.TTransport; import org.testng.annotations.Test; +import java.util.Optional; import java.util.concurrent.atomic.AtomicInteger; import static org.assertj.core.api.Assertions.assertThat; @@ -44,10 +45,9 @@ public void testAlternativeCall() return new TTransportMock(); }, "dummy", + Optional.empty(), new MetastoreSupportsDateStatistics(), - new AtomicInteger(), - new AtomicInteger(), - new AtomicInteger(), + true, new AtomicInteger(), new AtomicInteger(), new AtomicInteger(), diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/thrift/TestThriftMetastoreConfig.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/thrift/TestThriftMetastoreConfig.java index da3fa3277a96..f151f904c94c 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/thrift/TestThriftMetastoreConfig.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/thrift/TestThriftMetastoreConfig.java @@ -37,7 +37,8 @@ public class TestThriftMetastoreConfig public void testDefaults() { assertRecordedDefaults(recordDefaults(ThriftMetastoreConfig.class) - .setMetastoreTimeout(new Duration(10, SECONDS)) + .setConnectTimeout(new Duration(10, SECONDS)) + .setReadTimeout(new Duration(10, SECONDS)) .setSocksProxy(null) .setMaxRetries(9) .setBackoffScaleFactor(2.0) @@ -56,7 +57,8 @@ public void testDefaults() .setDeleteFilesOnDrop(false) .setMaxWaitForTransactionLock(new Duration(10, MINUTES)) .setAssumeCanonicalPartitionKeys(false) - .setWriteStatisticsThreads(20)); + .setWriteStatisticsThreads(20) + .setCatalogName(null)); } @Test @@ -67,7 +69,8 @@ public void testExplicitPropertyMappings() Path truststoreFile = Files.createTempFile(null, null); Map properties = ImmutableMap.builder() - .put("hive.metastore-timeout", "20s") + .put("hive.metastore.thrift.client.connect-timeout", "22s") + .put("hive.metastore.thrift.client.read-timeout", "44s") .put("hive.metastore.thrift.client.socks-proxy", "localhost:1234") .put("hive.metastore.thrift.client.max-retries", "15") .put("hive.metastore.thrift.client.backoff-scale-factor", "3.0") @@ -87,10 +90,12 @@ public void testExplicitPropertyMappings() .put("hive.metastore.thrift.write-statistics-threads", "10") .put("hive.metastore.thrift.assume-canonical-partition-keys", "true") .put("hive.metastore.thrift.use-spark-table-statistics-fallback", "false") + .put("hive.metastore.thrift.catalog-name", "custom_catalog_name") .buildOrThrow(); ThriftMetastoreConfig expected = new ThriftMetastoreConfig() - .setMetastoreTimeout(new Duration(20, SECONDS)) + .setConnectTimeout(new Duration(22, SECONDS)) + .setReadTimeout(new Duration(44, SECONDS)) .setSocksProxy(HostAndPort.fromParts("localhost", 1234)) .setMaxRetries(15) .setBackoffScaleFactor(3.0) @@ -109,7 +114,8 @@ public void testExplicitPropertyMappings() .setMaxWaitForTransactionLock(new Duration(5, MINUTES)) .setAssumeCanonicalPartitionKeys(true) .setWriteStatisticsThreads(10) - .setUseSparkTableStatisticsFallback(false); + .setUseSparkTableStatisticsFallback(false) + .setCatalogName("custom_catalog_name"); assertFullMapping(properties, expected); } diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/thrift/TestThriftMetastoreUtil.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/thrift/TestThriftMetastoreUtil.java index 4e13a7599394..7c24f0cec2a0 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/thrift/TestThriftMetastoreUtil.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/thrift/TestThriftMetastoreUtil.java @@ -80,14 +80,14 @@ public void testLongStatsToColumnStatistics() longColumnStatsData.setNumNulls(1); longColumnStatsData.setNumDVs(20); ColumnStatisticsObj columnStatisticsObj = new ColumnStatisticsObj("my_col", BIGINT_TYPE_NAME, longStats(longColumnStatsData)); - HiveColumnStatistics actual = fromMetastoreApiColumnStatistics(columnStatisticsObj, OptionalLong.of(1000)); + HiveColumnStatistics actual = fromMetastoreApiColumnStatistics(columnStatisticsObj); assertEquals( actual, HiveColumnStatistics.builder() .setIntegerStatistics(new IntegerStatistics(OptionalLong.of(0), OptionalLong.of(100))) .setNullsCount(1) - .setDistinctValuesCount(19) + .setDistinctValuesWithNullCount(20) .build()); } @@ -96,7 +96,7 @@ public void testEmptyLongStatsToColumnStatistics() { LongColumnStatsData emptyLongColumnStatsData = new LongColumnStatsData(); ColumnStatisticsObj columnStatisticsObj = new ColumnStatisticsObj("my_col", BIGINT_TYPE_NAME, longStats(emptyLongColumnStatsData)); - HiveColumnStatistics actual = fromMetastoreApiColumnStatistics(columnStatisticsObj, OptionalLong.empty()); + HiveColumnStatistics actual = fromMetastoreApiColumnStatistics(columnStatisticsObj); assertEquals( actual, @@ -114,14 +114,14 @@ public void testDoubleStatsToColumnStatistics() doubleColumnStatsData.setNumNulls(1); doubleColumnStatsData.setNumDVs(20); ColumnStatisticsObj columnStatisticsObj = new ColumnStatisticsObj("my_col", DOUBLE_TYPE_NAME, doubleStats(doubleColumnStatsData)); - HiveColumnStatistics actual = fromMetastoreApiColumnStatistics(columnStatisticsObj, OptionalLong.of(1000)); + HiveColumnStatistics actual = fromMetastoreApiColumnStatistics(columnStatisticsObj); assertEquals( actual, HiveColumnStatistics.builder() .setDoubleStatistics(new DoubleStatistics(OptionalDouble.of(0), OptionalDouble.of(100))) .setNullsCount(1) - .setDistinctValuesCount(19) + .setDistinctValuesWithNullCount(20) .build()); } @@ -130,7 +130,7 @@ public void testEmptyDoubleStatsToColumnStatistics() { DoubleColumnStatsData emptyDoubleColumnStatsData = new DoubleColumnStatsData(); ColumnStatisticsObj columnStatisticsObj = new ColumnStatisticsObj("my_col", DOUBLE_TYPE_NAME, doubleStats(emptyDoubleColumnStatsData)); - HiveColumnStatistics actual = fromMetastoreApiColumnStatistics(columnStatisticsObj, OptionalLong.empty()); + HiveColumnStatistics actual = fromMetastoreApiColumnStatistics(columnStatisticsObj); assertEquals( actual, @@ -150,14 +150,14 @@ public void testDecimalStatsToColumnStatistics() decimalColumnStatsData.setNumNulls(1); decimalColumnStatsData.setNumDVs(20); ColumnStatisticsObj columnStatisticsObj = new ColumnStatisticsObj("my_col", DECIMAL_TYPE_NAME, decimalStats(decimalColumnStatsData)); - HiveColumnStatistics actual = fromMetastoreApiColumnStatistics(columnStatisticsObj, OptionalLong.of(1000)); + HiveColumnStatistics actual = fromMetastoreApiColumnStatistics(columnStatisticsObj); assertEquals( actual, HiveColumnStatistics.builder() .setDecimalStatistics(new DecimalStatistics(Optional.of(low), Optional.of(high))) .setNullsCount(1) - .setDistinctValuesCount(19) + .setDistinctValuesWithNullCount(20) .build()); } @@ -166,7 +166,7 @@ public void testEmptyDecimalStatsToColumnStatistics() { DecimalColumnStatsData emptyDecimalColumnStatsData = new DecimalColumnStatsData(); ColumnStatisticsObj columnStatisticsObj = new ColumnStatisticsObj("my_col", DECIMAL_TYPE_NAME, decimalStats(emptyDecimalColumnStatsData)); - HiveColumnStatistics actual = fromMetastoreApiColumnStatistics(columnStatisticsObj, OptionalLong.empty()); + HiveColumnStatistics actual = fromMetastoreApiColumnStatistics(columnStatisticsObj); assertEquals( actual, @@ -183,7 +183,7 @@ public void testBooleanStatsToColumnStatistics() booleanColumnStatsData.setNumFalses(10); booleanColumnStatsData.setNumNulls(0); ColumnStatisticsObj columnStatisticsObj = new ColumnStatisticsObj("my_col", BOOLEAN_TYPE_NAME, booleanStats(booleanColumnStatsData)); - HiveColumnStatistics actual = fromMetastoreApiColumnStatistics(columnStatisticsObj, OptionalLong.empty()); + HiveColumnStatistics actual = fromMetastoreApiColumnStatistics(columnStatisticsObj); assertEquals( actual, @@ -198,7 +198,7 @@ public void testImpalaGeneratedBooleanStatistics() { BooleanColumnStatsData statsData = new BooleanColumnStatsData(1L, -1L, 2L); ColumnStatisticsObj columnStatisticsObj = new ColumnStatisticsObj("my_col", BOOLEAN_TYPE_NAME, booleanStats(statsData)); - HiveColumnStatistics actual = fromMetastoreApiColumnStatistics(columnStatisticsObj, OptionalLong.empty()); + HiveColumnStatistics actual = fromMetastoreApiColumnStatistics(columnStatisticsObj); assertEquals( actual, @@ -213,7 +213,7 @@ public void testEmptyBooleanStatsToColumnStatistics() { BooleanColumnStatsData emptyBooleanColumnStatsData = new BooleanColumnStatsData(); ColumnStatisticsObj columnStatisticsObj = new ColumnStatisticsObj("my_col", BOOLEAN_TYPE_NAME, booleanStats(emptyBooleanColumnStatsData)); - HiveColumnStatistics actual = fromMetastoreApiColumnStatistics(columnStatisticsObj, OptionalLong.empty()); + HiveColumnStatistics actual = fromMetastoreApiColumnStatistics(columnStatisticsObj); assertEquals( actual, @@ -231,14 +231,14 @@ public void testDateStatsToColumnStatistics() dateColumnStatsData.setNumNulls(1); dateColumnStatsData.setNumDVs(20); ColumnStatisticsObj columnStatisticsObj = new ColumnStatisticsObj("my_col", DATE_TYPE_NAME, dateStats(dateColumnStatsData)); - HiveColumnStatistics actual = fromMetastoreApiColumnStatistics(columnStatisticsObj, OptionalLong.of(1000)); + HiveColumnStatistics actual = fromMetastoreApiColumnStatistics(columnStatisticsObj); assertEquals( actual, HiveColumnStatistics.builder() .setDateStatistics(new DateStatistics(Optional.of(LocalDate.ofEpochDay(1000)), Optional.of(LocalDate.ofEpochDay(2000)))) .setNullsCount(1) - .setDistinctValuesCount(19) + .setDistinctValuesWithNullCount(20) .build()); } @@ -247,7 +247,7 @@ public void testEmptyDateStatsToColumnStatistics() { DateColumnStatsData emptyDateColumnStatsData = new DateColumnStatsData(); ColumnStatisticsObj columnStatisticsObj = new ColumnStatisticsObj("my_col", DATE_TYPE_NAME, dateStats(emptyDateColumnStatsData)); - HiveColumnStatistics actual = fromMetastoreApiColumnStatistics(columnStatisticsObj, OptionalLong.empty()); + HiveColumnStatistics actual = fromMetastoreApiColumnStatistics(columnStatisticsObj); assertEquals( actual, @@ -265,15 +265,15 @@ public void testStringStatsToColumnStatistics() stringColumnStatsData.setNumNulls(1); stringColumnStatsData.setNumDVs(20); ColumnStatisticsObj columnStatisticsObj = new ColumnStatisticsObj("my_col", STRING_TYPE_NAME, stringStats(stringColumnStatsData)); - HiveColumnStatistics actual = fromMetastoreApiColumnStatistics(columnStatisticsObj, OptionalLong.of(2)); + HiveColumnStatistics actual = fromMetastoreApiColumnStatistics(columnStatisticsObj); assertEquals( actual, HiveColumnStatistics.builder() .setMaxValueSizeInBytes(100) - .setTotalSizeInBytes(23) + .setAverageColumnLength(23.333) .setNullsCount(1) - .setDistinctValuesCount(1) + .setDistinctValuesWithNullCount(20) .build()); } @@ -282,7 +282,7 @@ public void testEmptyStringColumnStatsData() { StringColumnStatsData emptyStringColumnStatsData = new StringColumnStatsData(); ColumnStatisticsObj columnStatisticsObj = new ColumnStatisticsObj("my_col", STRING_TYPE_NAME, stringStats(emptyStringColumnStatsData)); - HiveColumnStatistics actual = fromMetastoreApiColumnStatistics(columnStatisticsObj, OptionalLong.empty()); + HiveColumnStatistics actual = fromMetastoreApiColumnStatistics(columnStatisticsObj); assertEquals(actual, HiveColumnStatistics.builder().build()); } @@ -295,13 +295,13 @@ public void testBinaryStatsToColumnStatistics() binaryColumnStatsData.setAvgColLen(22.2); binaryColumnStatsData.setNumNulls(2); ColumnStatisticsObj columnStatisticsObj = new ColumnStatisticsObj("my_col", BINARY_TYPE_NAME, binaryStats(binaryColumnStatsData)); - HiveColumnStatistics actual = fromMetastoreApiColumnStatistics(columnStatisticsObj, OptionalLong.of(4)); + HiveColumnStatistics actual = fromMetastoreApiColumnStatistics(columnStatisticsObj); assertEquals( actual, HiveColumnStatistics.builder() .setMaxValueSizeInBytes(100) - .setTotalSizeInBytes(44) + .setAverageColumnLength(22.2) .setNullsCount(2) .build()); } @@ -311,7 +311,7 @@ public void testEmptyBinaryStatsToColumnStatistics() { BinaryColumnStatsData emptyBinaryColumnStatsData = new BinaryColumnStatsData(); ColumnStatisticsObj columnStatisticsObj = new ColumnStatisticsObj("my_col", BINARY_TYPE_NAME, binaryStats(emptyBinaryColumnStatsData)); - HiveColumnStatistics actual = fromMetastoreApiColumnStatistics(columnStatisticsObj, OptionalLong.empty()); + HiveColumnStatistics actual = fromMetastoreApiColumnStatistics(columnStatisticsObj); assertEquals(actual, HiveColumnStatistics.builder().build()); } @@ -323,19 +323,19 @@ public void testSingleDistinctValue() doubleColumnStatsData.setNumNulls(10); doubleColumnStatsData.setNumDVs(1); ColumnStatisticsObj columnStatisticsObj = new ColumnStatisticsObj("my_col", DOUBLE_TYPE_NAME, doubleStats(doubleColumnStatsData)); - HiveColumnStatistics actual = fromMetastoreApiColumnStatistics(columnStatisticsObj, OptionalLong.of(10)); + HiveColumnStatistics actual = fromMetastoreApiColumnStatistics(columnStatisticsObj); assertEquals(actual.getNullsCount(), OptionalLong.of(10)); - assertEquals(actual.getDistinctValuesCount(), OptionalLong.of(0)); + assertEquals(actual.getDistinctValuesWithNullCount(), OptionalLong.of(1)); doubleColumnStatsData = new DoubleColumnStatsData(); doubleColumnStatsData.setNumNulls(10); doubleColumnStatsData.setNumDVs(1); columnStatisticsObj = new ColumnStatisticsObj("my_col", DOUBLE_TYPE_NAME, doubleStats(doubleColumnStatsData)); - actual = fromMetastoreApiColumnStatistics(columnStatisticsObj, OptionalLong.of(11)); + actual = fromMetastoreApiColumnStatistics(columnStatisticsObj); assertEquals(actual.getNullsCount(), OptionalLong.of(10)); - assertEquals(actual.getDistinctValuesCount(), OptionalLong.of(1)); + assertEquals(actual.getDistinctValuesWithNullCount(), OptionalLong.of(1)); } @Test diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/thrift/TestThriftSparkMetastoreUtil.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/thrift/TestThriftSparkMetastoreUtil.java index 171856d550dc..25e48e601fea 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/thrift/TestThriftSparkMetastoreUtil.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/thrift/TestThriftSparkMetastoreUtil.java @@ -64,7 +64,7 @@ public void testSparkLongStatsToColumnStatistics() HiveColumnStatistics.builder() .setIntegerStatistics(new IntegerStatistics(OptionalLong.of(1), OptionalLong.of(4))) .setNullsCount(0) - .setDistinctValuesCount(4) + .setDistinctValuesWithNullCount(4) .build()); } @@ -91,7 +91,7 @@ public void testSparkDoubleStatsToColumnStatistics() HiveColumnStatistics.builder() .setDoubleStatistics(new DoubleStatistics(OptionalDouble.of(0.3), OptionalDouble.of(3.3))) .setNullsCount(1) - .setDistinctValuesCount(9) + .setDistinctValuesWithNullCount(10) .build()); } @@ -118,7 +118,7 @@ public void testSparkDecimalStatsToColumnStatistics() HiveColumnStatistics.builder() .setDecimalStatistics(new DecimalStatistics(Optional.of(new BigDecimal("0.3")), Optional.of(new BigDecimal("3.3")))) .setNullsCount(1) - .setDistinctValuesCount(9) + .setDistinctValuesWithNullCount(10) .build()); } @@ -171,7 +171,7 @@ public void testSparkDateStatsToColumnStatistics() HiveColumnStatistics.builder() .setDateStatistics((new DateStatistics(Optional.of(LocalDate.of(2000, 1, 1)), Optional.of(LocalDate.of(2030, 12, 31))))) .setNullsCount(3) - .setDistinctValuesCount(7) + .setDistinctValuesWithNullCount(10) .build()); } @@ -192,8 +192,8 @@ public void testSparkStringStatsToColumnStatistics() actual, HiveColumnStatistics.builder() .setNullsCount(7) - .setDistinctValuesCount(3) - .setTotalSizeInBytes(30) + .setDistinctValuesWithNullCount(3) + .setAverageColumnLength(10) .build()); } @@ -214,7 +214,7 @@ public void testSparkBinaryStatsToColumnStatistics() actual, HiveColumnStatistics.builder() .setNullsCount(3) - .setTotalSizeInBytes(70) + .setAverageColumnLength(10) .setMaxValueSizeInBytes(10) .build()); } diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/thrift/TestingTokenAwareMetastoreClientFactory.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/thrift/TestingTokenAwareMetastoreClientFactory.java index 4f65cb429a6c..9d960e86d6ee 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/thrift/TestingTokenAwareMetastoreClientFactory.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/thrift/TestingTokenAwareMetastoreClientFactory.java @@ -17,6 +17,7 @@ import io.airlift.units.Duration; import org.apache.thrift.TException; +import java.net.URI; import java.util.Optional; import static java.util.Objects.requireNonNull; @@ -29,23 +30,31 @@ public class TestingTokenAwareMetastoreClientFactory public static final Duration TIMEOUT = new Duration(20, SECONDS); private final DefaultThriftMetastoreClientFactory factory; - private final HostAndPort address; + private final URI address; - public TestingTokenAwareMetastoreClientFactory(Optional socksProxy, HostAndPort address) + private final MetastoreClientAdapterProvider metastoreClientAdapterProvider; + + public TestingTokenAwareMetastoreClientFactory(Optional socksProxy, URI uri) + { + this(socksProxy, uri, TIMEOUT, delegate -> delegate); + } + + public TestingTokenAwareMetastoreClientFactory(Optional socksProxy, URI address, Duration timeout) { - this(socksProxy, address, TIMEOUT); + this(socksProxy, address, timeout, delegate -> delegate); } - public TestingTokenAwareMetastoreClientFactory(Optional socksProxy, HostAndPort address, Duration timeout) + public TestingTokenAwareMetastoreClientFactory(Optional socksProxy, URI uri, Duration timeout, MetastoreClientAdapterProvider metastoreClientAdapterProvider) { - this.factory = new DefaultThriftMetastoreClientFactory(Optional.empty(), socksProxy, timeout, AUTHENTICATION, "localhost"); - this.address = requireNonNull(address, "address is null"); + this.factory = new DefaultThriftMetastoreClientFactory(Optional.empty(), socksProxy, timeout, timeout, AUTHENTICATION, "localhost", Optional.empty()); + this.address = requireNonNull(uri, "uri is null"); + this.metastoreClientAdapterProvider = requireNonNull(metastoreClientAdapterProvider, "metastoreClientAdapterProvider is null"); } @Override public ThriftMetastoreClient createMetastoreClient(Optional delegationToken) throws TException { - return factory.create(address, delegationToken); + return metastoreClientAdapterProvider.createThriftMetastoreClientAdapter(factory.create(address, delegationToken)); } } diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/optimizer/TestConnectorPushdownRulesWithHive.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/optimizer/TestConnectorPushdownRulesWithHive.java index f58af7606135..8a59708af024 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/optimizer/TestConnectorPushdownRulesWithHive.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/optimizer/TestConnectorPushdownRulesWithHive.java @@ -113,7 +113,7 @@ protected Optional createLocalQueryRunner() metastore.createDatabase(database); LocalQueryRunner queryRunner = LocalQueryRunner.create(HIVE_SESSION); - queryRunner.createCatalog(TEST_CATALOG_NAME, new TestingHiveConnectorFactory(metastore), ImmutableMap.of()); + queryRunner.createCatalog(TEST_CATALOG_NAME, new TestingHiveConnectorFactory(baseDir.toPath(), Optional.of(metastore)), ImmutableMap.of()); catalogHandle = queryRunner.getCatalogHandle(TEST_CATALOG_NAME); return Optional.of(queryRunner); diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/optimizer/TestHivePlans.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/optimizer/TestHivePlans.java index 07e0c129f006..3c6c303bf16a 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/optimizer/TestHivePlans.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/optimizer/TestHivePlans.java @@ -89,7 +89,7 @@ protected LocalQueryRunner createLocalQueryRunner() protected LocalQueryRunner createQueryRunner(Session session, HiveMetastore metastore) { LocalQueryRunner queryRunner = LocalQueryRunner.create(session); - queryRunner.createCatalog(HIVE_CATALOG_NAME, new TestingHiveConnectorFactory(metastore), Map.of("hive.max-partitions-for-eager-load", "5")); + queryRunner.createCatalog(HIVE_CATALOG_NAME, new TestingHiveConnectorFactory(baseDir.toPath(), Optional.of(metastore)), Map.of("hive.max-partitions-for-eager-load", "5")); return queryRunner; } diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/optimizer/TestHiveProjectionPushdownIntoTableScan.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/optimizer/TestHiveProjectionPushdownIntoTableScan.java index 679bce14b825..1261bc7d05d5 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/optimizer/TestHiveProjectionPushdownIntoTableScan.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/optimizer/TestHiveProjectionPushdownIntoTableScan.java @@ -89,7 +89,7 @@ protected LocalQueryRunner createLocalQueryRunner() metastore.createDatabase(database); LocalQueryRunner queryRunner = LocalQueryRunner.create(HIVE_SESSION); - queryRunner.createCatalog(HIVE_CATALOG_NAME, new TestingHiveConnectorFactory(metastore), ImmutableMap.of()); + queryRunner.createCatalog(HIVE_CATALOG_NAME, new TestingHiveConnectorFactory(baseDir.toPath(), Optional.of(metastore)), ImmutableMap.of()); return queryRunner; } diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/orc/TestHiveOrcWithShortZoneId.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/orc/TestHiveOrcWithShortZoneId.java index 876ae2958490..4b2fcf879468 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/orc/TestHiveOrcWithShortZoneId.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/orc/TestHiveOrcWithShortZoneId.java @@ -14,28 +14,49 @@ package io.trino.plugin.hive.orc; import com.google.common.collect.ImmutableList; +import com.google.common.io.Resources; +import io.trino.filesystem.Location; +import io.trino.filesystem.TrinoFileSystem; +import io.trino.filesystem.TrinoFileSystemFactory; import io.trino.plugin.hive.HiveQueryRunner; +import io.trino.spi.security.ConnectorIdentity; import io.trino.testing.AbstractTestQueryFramework; import io.trino.testing.QueryRunner; import io.trino.testing.sql.TestTable; import org.testng.annotations.Test; -import static io.trino.testing.containers.TestContainers.getPathFromClassPathResource; +import java.io.OutputStream; +import java.net.URL; +import java.util.UUID; + +import static io.trino.plugin.hive.TestingHiveUtils.getConnectorService; public class TestHiveOrcWithShortZoneId extends AbstractTestQueryFramework { - private String resourceLocation; + private Location dataFile; @Override protected QueryRunner createQueryRunner() throws Exception { - // See README.md to know how resource is generated - resourceLocation = getPathFromClassPathResource("with_short_zone_id/data"); - return HiveQueryRunner.builder() + QueryRunner queryRunner = HiveQueryRunner.builder() .addHiveProperty("hive.orc.read-legacy-short-zone-id", "true") .build(); + + URL resourceLocation = Resources.getResource("with_short_zone_id/data/data.orc"); + + TrinoFileSystem fileSystem = getConnectorService(queryRunner, TrinoFileSystemFactory.class) + .create(ConnectorIdentity.ofUser("test")); + + Location tempDir = Location.of("local:///temp_" + UUID.randomUUID()); + fileSystem.createDirectory(tempDir); + dataFile = tempDir.appendPath("data.orc"); + try (OutputStream out = fileSystem.newOutputFile(dataFile).create()) { + Resources.copy(resourceLocation, out); + } + + return queryRunner; } @Test @@ -45,7 +66,7 @@ public void testSelectWithShortZoneId() try (TestTable testTable = new TestTable( getQueryRunner()::execute, "test_select_with_short_zone_id_", - "(id INT, firstName VARCHAR, lastName VARCHAR) WITH (external_location = '%s')".formatted(resourceLocation))) { + "(id INT, firstName VARCHAR, lastName VARCHAR) WITH (external_location = '%s')".formatted(dataFile.parentDirectory()))) { assertQuery("SELECT * FROM " + testTable.getName(), "VALUES (1, 'John', 'Doe')"); } } diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/orc/TestOrcPageSourceFactory.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/orc/TestOrcPageSourceFactory.java index bcdc09ae6b87..39661c274a9d 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/orc/TestOrcPageSourceFactory.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/orc/TestOrcPageSourceFactory.java @@ -15,33 +15,37 @@ import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; +import com.google.common.io.Resources; import io.trino.filesystem.Location; -import io.trino.filesystem.hdfs.HdfsFileSystemFactory; +import io.trino.filesystem.TrinoFileSystem; +import io.trino.filesystem.TrinoFileSystemFactory; +import io.trino.filesystem.memory.MemoryFileSystemFactory; import io.trino.plugin.hive.AcidInfo; import io.trino.plugin.hive.FileFormatDataSourceStats; import io.trino.plugin.hive.HiveColumnHandle; import io.trino.plugin.hive.HiveConfig; import io.trino.plugin.hive.HivePageSourceFactory; import io.trino.plugin.hive.ReaderPageSource; +import io.trino.plugin.hive.Schema; import io.trino.spi.Page; import io.trino.spi.connector.ConnectorPageSource; import io.trino.spi.predicate.Domain; import io.trino.spi.predicate.TupleDomain; +import io.trino.spi.security.ConnectorIdentity; import io.trino.spi.type.Type; import io.trino.tpch.Nation; import io.trino.tpch.NationColumn; import io.trino.tpch.NationGenerator; import org.testng.annotations.Test; -import java.io.File; -import java.net.URISyntaxException; +import java.io.IOException; +import java.io.OutputStream; import java.util.ArrayList; import java.util.List; import java.util.Map; import java.util.Optional; import java.util.OptionalInt; import java.util.OptionalLong; -import java.util.Properties; import java.util.Set; import java.util.function.LongPredicate; @@ -51,8 +55,6 @@ import static io.trino.plugin.hive.HiveColumnHandle.ColumnType.REGULAR; import static io.trino.plugin.hive.HiveColumnHandle.createBaseColumn; import static io.trino.plugin.hive.HiveStorageFormat.ORC; -import static io.trino.plugin.hive.HiveTestUtils.HDFS_ENVIRONMENT; -import static io.trino.plugin.hive.HiveTestUtils.HDFS_FILE_SYSTEM_STATS; import static io.trino.plugin.hive.HiveTestUtils.SESSION; import static io.trino.plugin.hive.HiveType.toHiveType; import static io.trino.plugin.hive.acid.AcidTransaction.NO_ACID_TRANSACTION; @@ -64,10 +66,7 @@ import static io.trino.tpch.NationColumn.NATION_KEY; import static io.trino.tpch.NationColumn.REGION_KEY; import static java.util.Collections.nCopies; -import static org.apache.hadoop.hive.metastore.api.hive_metastoreConstants.FILE_INPUT_FORMAT; -import static org.apache.hadoop.hive.metastore.api.hive_metastoreConstants.TABLE_IS_TRANSACTIONAL; -import static org.apache.hadoop.hive.ql.io.AcidUtils.deleteDeltaSubdir; -import static org.apache.hadoop.hive.serde.serdeConstants.SERIALIZATION_LIB; +import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatThrownBy; import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertFalse; @@ -75,20 +74,17 @@ public class TestOrcPageSourceFactory { private static final Map ALL_COLUMNS = ImmutableMap.of(NATION_KEY, 0, NAME, 1, REGION_KEY, 2, COMMENT, 3); - private static final HivePageSourceFactory PAGE_SOURCE_FACTORY = new OrcPageSourceFactory( - new OrcReaderConfig(), - new HdfsFileSystemFactory(HDFS_ENVIRONMENT, HDFS_FILE_SYSTEM_STATS), - new FileFormatDataSourceStats(), - new HiveConfig()); @Test public void testFullFileRead() + throws IOException { assertRead(ImmutableMap.of(NATION_KEY, 0, NAME, 1, REGION_KEY, 2, COMMENT, 3), OptionalLong.empty(), Optional.empty(), nationKey -> false); } @Test public void testSingleColumnRead() + throws IOException { assertRead(ImmutableMap.of(REGION_KEY, ALL_COLUMNS.get(REGION_KEY)), OptionalLong.empty(), Optional.empty(), nationKey -> false); } @@ -98,6 +94,7 @@ public void testSingleColumnRead() */ @Test public void testFullFileSkipped() + throws IOException { assertRead(ALL_COLUMNS, OptionalLong.of(100L), Optional.empty(), nationKey -> false); } @@ -107,49 +104,59 @@ public void testFullFileSkipped() */ @Test public void testSomeStripesAndRowGroupRead() + throws IOException { assertRead(ALL_COLUMNS, OptionalLong.of(5L), Optional.empty(), nationKey -> false); } @Test public void testDeletedRows() + throws IOException { - Location partitionLocation = Location.of(getResource("nation_delete_deltas").toString()); - Optional acidInfo = AcidInfo.builder(partitionLocation) - .addDeleteDelta(partitionLocation.appendPath(deleteDeltaSubdir(3L, 3L, 0))) - .addDeleteDelta(partitionLocation.appendPath(deleteDeltaSubdir(4L, 4L, 0))) + TrinoFileSystemFactory fileSystemFactory = new MemoryFileSystemFactory(); + Location fileLocation = copyResource(fileSystemFactory, "nationFile25kRowsSortedOnNationKey/bucket_00000"); + long fileLength = fileSystemFactory.create(ConnectorIdentity.ofUser("test")).newInputFile(fileLocation).length(); + + Location deleteFile3 = copyResource(fileSystemFactory, "nation_delete_deltas/delete_delta_0000003_0000003_0000/bucket_00000"); + Location deleteFile4 = copyResource(fileSystemFactory, "nation_delete_deltas/delete_delta_0000004_0000004_0000/bucket_00000"); + Optional acidInfo = AcidInfo.builder(deleteFile3.parentDirectory().parentDirectory()) + .addDeleteDelta(deleteFile3.parentDirectory()) + .addDeleteDelta(deleteFile4.parentDirectory()) .build(); - assertRead(ALL_COLUMNS, OptionalLong.empty(), acidInfo, nationKey -> nationKey == 5 || nationKey == 19); + List actual = readFile(fileSystemFactory, ALL_COLUMNS, OptionalLong.empty(), acidInfo, fileLocation, fileLength); + + List expected = expectedResult(OptionalLong.empty(), nationKey -> nationKey == 5 || nationKey == 19, 1000); + assertEqualsByColumns(ALL_COLUMNS.keySet(), actual, expected); } @Test public void testReadWithAcidVersionValidationHive3() throws Exception { - File tableFile = new File(getResource("acid_version_validation/acid_version_hive_3/00000_0").toURI()); - Location tablePath = Location.of(tableFile.getParentFile().toURI().toString()); + TrinoFileSystemFactory fileSystemFactory = new MemoryFileSystemFactory(); + Location fileLocation = copyResource(fileSystemFactory, "acid_version_validation/acid_version_hive_3/00000_0"); - Optional acidInfo = AcidInfo.builder(tablePath) + Optional acidInfo = AcidInfo.builder(fileLocation.parentDirectory()) .setOrcAcidVersionValidated(false) .build(); - List result = readFile(Map.of(), OptionalLong.empty(), acidInfo, tableFile.getPath(), 625); - assertEquals(result.size(), 1); + List result = readFile(fileSystemFactory, Map.of(), OptionalLong.empty(), acidInfo, fileLocation, 625); + assertThat(result).hasSize(1); } @Test public void testReadWithAcidVersionValidationNoVersionInMetadata() throws Exception { - File tableFile = new File(getResource("acid_version_validation/no_orc_acid_version_in_metadata/00000_0").toURI()); - Location tablePath = Location.of(tableFile.getParentFile().toURI().toString()); + TrinoFileSystemFactory fileSystemFactory = new MemoryFileSystemFactory(); + Location fileLocation = copyResource(fileSystemFactory, "acid_version_validation/no_orc_acid_version_in_metadata/00000_0"); - Optional acidInfo = AcidInfo.builder(tablePath) + Optional acidInfo = AcidInfo.builder(fileLocation.parentDirectory()) .setOrcAcidVersionValidated(false) .build(); - assertThatThrownBy(() -> readFile(Map.of(), OptionalLong.empty(), acidInfo, tableFile.getPath(), 730)) + assertThatThrownBy(() -> readFile(fileSystemFactory, Map.of(), OptionalLong.empty(), acidInfo, fileLocation, 730)) .hasMessageMatching("Hive transactional tables are supported since Hive 3.0. Expected `hive.acid.version` in ORC metadata" + " in .*/acid_version_validation/no_orc_acid_version_in_metadata/00000_0 to be >=2 but was ." + " If you have upgraded from an older version of Hive, make sure a major compaction has been run at least once after the upgrade."); @@ -159,17 +166,19 @@ public void testReadWithAcidVersionValidationNoVersionInMetadata() public void testFullFileReadOriginalFilesTable() throws Exception { - File tableFile = new File(getResource("fullacidNationTableWithOriginalFiles/000000_0").toURI()); - Location tablePath = Location.of(tableFile.toURI().toString()).parentDirectory(); + TrinoFileSystemFactory fileSystemFactory = new MemoryFileSystemFactory(); + Location fileLocation = copyResource(fileSystemFactory, "fullacidNationTableWithOriginalFiles/000000_0"); + Location deleteDeltaLocation = copyResource(fileSystemFactory, "fullacidNationTableWithOriginalFiles/delete_delta_10000001_10000001_0000/bucket_00000"); + Location tablePath = fileLocation.parentDirectory(); AcidInfo acidInfo = AcidInfo.builder(tablePath) - .addDeleteDelta(tablePath.appendPath(deleteDeltaSubdir(10000001, 10000001, 0))) - .addOriginalFile(tablePath.appendPath("000000_0"), 1780, 0) + .addDeleteDelta(deleteDeltaLocation.parentDirectory()) + .addOriginalFile(fileLocation, 1780, 0) .setOrcAcidVersionValidated(true) .buildWithRequiredOriginalFiles(0); List expected = expectedResult(OptionalLong.empty(), nationKey -> nationKey == 24, 1); - List result = readFile(ALL_COLUMNS, OptionalLong.empty(), Optional.of(acidInfo), tablePath + "/000000_0", 1780); + List result = readFile(fileSystemFactory, ALL_COLUMNS, OptionalLong.empty(), Optional.of(acidInfo), fileLocation, 1780); assertEquals(result.size(), expected.size()); int deletedRowKey = 24; @@ -179,6 +188,7 @@ public void testFullFileReadOriginalFilesTable() } private static void assertRead(Map columns, OptionalLong nationKeyPredicate, Optional acidInfo, LongPredicate deletedRows) + throws IOException { List actual = readFile(columns, nationKeyPredicate, acidInfo); @@ -203,18 +213,16 @@ private static List expectedResult(OptionalLong nationKeyPredicate, Long } private static List readFile(Map columns, OptionalLong nationKeyPredicate, Optional acidInfo) + throws IOException { // This file has the contains the TPC-H nation table which each row repeated 1000 times - try { - File testFile = new File(getResource("nationFile25kRowsSortedOnNationKey/bucket_00000").toURI()); - return readFile(columns, nationKeyPredicate, acidInfo, testFile.toURI().getPath(), testFile.length()); - } - catch (URISyntaxException e) { - throw new RuntimeException(e); - } + TrinoFileSystemFactory fileSystemFactory = new MemoryFileSystemFactory(); + Location fileLocation = copyResource(fileSystemFactory, "nationFile25kRowsSortedOnNationKey/bucket_00000"); + long fileLength = fileSystemFactory.create(ConnectorIdentity.ofUser("test")).newInputFile(fileLocation).length(); + return readFile(fileSystemFactory, columns, nationKeyPredicate, acidInfo, fileLocation, fileLength); } - private static List readFile(Map columns, OptionalLong nationKeyPredicate, Optional acidInfo, String filePath, long fileSize) + private static List readFile(TrinoFileSystemFactory fileSystemFactory, Map columns, OptionalLong nationKeyPredicate, Optional acidInfo, Location location, long fileSize) { TupleDomain tupleDomain = TupleDomain.all(); if (nationKeyPredicate.isPresent()) { @@ -229,12 +237,19 @@ private static List readFile(Map columns, Optiona .map(HiveColumnHandle::getName) .collect(toImmutableList()); - Optional pageSourceWithProjections = PAGE_SOURCE_FACTORY.createPageSource( + HivePageSourceFactory pageSourceFactory = new OrcPageSourceFactory( + new OrcReaderConfig(), + fileSystemFactory, + new FileFormatDataSourceStats(), + new HiveConfig()); + + Optional pageSourceWithProjections = pageSourceFactory.createPageSource( SESSION, - Location.of(filePath), + location, 0, fileSize, fileSize, + 12345, createSchema(), columnHandles, tupleDomain, @@ -312,13 +327,9 @@ private static HiveColumnHandle toHiveColumnHandle(NationColumn nationColumn, in Optional.empty()); } - private static Properties createSchema() + private static Schema createSchema() { - Properties schema = new Properties(); - schema.setProperty(SERIALIZATION_LIB, ORC.getSerde()); - schema.setProperty(FILE_INPUT_FORMAT, ORC.getInputFormat()); - schema.setProperty(TABLE_IS_TRANSACTIONAL, "true"); - return schema; + return new Schema(ORC.getSerde(), true, ImmutableMap.of()); } private static void assertEqualsByColumns(Set columns, List actualRows, List expectedRows) @@ -333,4 +344,15 @@ private static void assertEqualsByColumns(Set columns, List"); } } + + private static Location copyResource(TrinoFileSystemFactory fileSystemFactory, String resourceName) + throws IOException + { + Location location = Location.of("memory:///" + resourceName); + TrinoFileSystem fileSystem = fileSystemFactory.create(ConnectorIdentity.ofUser("test")); + try (OutputStream outputStream = fileSystem.newOutputFile(location).create()) { + Resources.copy(getResource(resourceName), outputStream); + } + return location; + } } diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/orc/TestOrcPredicates.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/orc/TestOrcPredicates.java index 383bf51e08c7..c34d24e31987 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/orc/TestOrcPredicates.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/orc/TestOrcPredicates.java @@ -17,56 +17,62 @@ import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; import io.trino.filesystem.Location; +import io.trino.filesystem.TrinoFileSystemFactory; +import io.trino.filesystem.TrinoInputFile; +import io.trino.filesystem.memory.MemoryFileSystemFactory; +import io.trino.hive.thrift.metastore.hive_metastoreConstants; import io.trino.orc.OrcReaderOptions; import io.trino.orc.OrcWriterOptions; import io.trino.plugin.hive.AbstractTestHiveFileFormats; import io.trino.plugin.hive.FileFormatDataSourceStats; +import io.trino.plugin.hive.FileWriter; import io.trino.plugin.hive.HiveColumnHandle; +import io.trino.plugin.hive.HiveColumnProjectionInfo; import io.trino.plugin.hive.HiveCompressionCodec; import io.trino.plugin.hive.HiveConfig; import io.trino.plugin.hive.HivePageSourceProvider; -import io.trino.plugin.hive.HivePartitionKey; +import io.trino.plugin.hive.HiveType; import io.trino.plugin.hive.NodeVersion; -import io.trino.plugin.hive.TableToPartitionMapping; +import io.trino.plugin.hive.Schema; +import io.trino.plugin.hive.WriterKind; +import io.trino.plugin.hive.util.HiveTypeTranslator; +import io.trino.plugin.hive.util.SerdeConstants; import io.trino.spi.Page; +import io.trino.spi.block.Block; +import io.trino.spi.block.IntArrayBlock; +import io.trino.spi.block.LongArrayBlock; +import io.trino.spi.block.RowBlock; +import io.trino.spi.block.RunLengthEncodedBlock; import io.trino.spi.connector.ConnectorPageSource; import io.trino.spi.connector.ConnectorSession; import io.trino.spi.predicate.Domain; import io.trino.spi.predicate.TupleDomain; -import org.apache.hadoop.mapred.FileSplit; +import io.trino.spi.type.RowType; import org.testng.annotations.Test; -import java.io.File; +import java.io.IOException; import java.time.Instant; -import java.util.HashSet; import java.util.List; +import java.util.Map; import java.util.Optional; import java.util.OptionalInt; -import java.util.Properties; -import java.util.Set; import java.util.stream.Collectors; -import static com.google.common.base.Preconditions.checkState; -import static com.google.common.collect.ImmutableList.toImmutableList; -import static io.trino.hadoop.ConfigurationInstantiator.newEmptyConfiguration; +import static io.trino.plugin.hive.HiveColumnHandle.ColumnType.REGULAR; +import static io.trino.plugin.hive.HiveColumnHandle.createBaseColumn; import static io.trino.plugin.hive.HivePageSourceProvider.ColumnMapping.buildColumnMappings; import static io.trino.plugin.hive.HiveStorageFormat.ORC; -import static io.trino.plugin.hive.HiveTestUtils.HDFS_FILE_SYSTEM_FACTORY; import static io.trino.plugin.hive.HiveTestUtils.getHiveSession; import static io.trino.plugin.hive.acid.AcidTransaction.NO_ACID_TRANSACTION; +import static io.trino.plugin.hive.util.SerdeConstants.LIST_COLUMNS; +import static io.trino.plugin.hive.util.SerdeConstants.LIST_COLUMN_TYPES; import static io.trino.spi.type.BigintType.BIGINT; -import static io.trino.testing.StructuralTestUtil.rowBlockOf; +import static io.trino.spi.type.IntegerType.INTEGER; +import static io.trino.spi.type.RowType.field; import static io.trino.type.InternalTypeManager.TESTING_TYPE_MANAGER; -import static java.lang.String.format; import static java.util.stream.Collectors.toList; -import static org.apache.hadoop.hive.metastore.api.hive_metastoreConstants.FILE_INPUT_FORMAT; -import static org.apache.hadoop.hive.serde.serdeConstants.SERIALIZATION_LIB; -import static org.apache.hadoop.hive.serde2.objectinspector.ObjectInspectorFactory.getStandardStructObjectInspector; -import static org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorFactory.javaIntObjectInspector; -import static org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorFactory.javaLongObjectInspector; +import static org.assertj.core.api.Assertions.assertThat; import static org.joda.time.DateTimeZone.UTC; -import static org.testng.Assert.assertEquals; -import static org.testng.Assert.assertTrue; public class TestOrcPredicates extends AbstractTestHiveFileFormats @@ -74,14 +80,30 @@ public class TestOrcPredicates private static final int NUM_ROWS = 50000; private static final FileFormatDataSourceStats STATS = new FileFormatDataSourceStats(); - // Prepare test columns - private static final TestColumn columnPrimitiveInteger = new TestColumn("column_primitive_integer", javaIntObjectInspector, 3, 3); - private static final TestColumn columnStruct = new TestColumn( + private static final HiveColumnHandle INTEGER_COLUMN = createBaseColumn("column_primitive_integer", 0, HiveType.HIVE_INT, INTEGER, REGULAR, Optional.empty()); + private static final HiveColumnHandle STRUCT_COLUMN = createBaseColumn( "column1_struct", - getStandardStructObjectInspector(ImmutableList.of("field0", "field1"), ImmutableList.of(javaLongObjectInspector, javaLongObjectInspector)), - new Long[] {4L, 5L}, - rowBlockOf(ImmutableList.of(BIGINT, BIGINT), 4L, 5L)); - private static final TestColumn columnPrimitiveBigInt = new TestColumn("column_primitive_bigint", javaLongObjectInspector, 6L, 6L); + 1, + HiveTypeTranslator.toHiveType(RowType.rowType(field("field0", BIGINT), field("field1", BIGINT))), + RowType.rowType(field("field0", BIGINT), field("field1", BIGINT)), + REGULAR, + Optional.empty()); + private static final HiveColumnHandle BIGINT_COLUMN = createBaseColumn("column_primitive_bigint", 2, HiveType.HIVE_LONG, BIGINT, REGULAR, Optional.empty()); + private static final List COLUMNS = ImmutableList.of(INTEGER_COLUMN, STRUCT_COLUMN, BIGINT_COLUMN); + + private static final HiveColumnHandle STRUCT_FIELD1_COLUMN = new HiveColumnHandle( + STRUCT_COLUMN.getBaseColumnName(), + STRUCT_COLUMN.getBaseHiveColumnIndex(), + STRUCT_COLUMN.getBaseHiveType(), + STRUCT_COLUMN.getBaseType(), + Optional.of(new HiveColumnProjectionInfo( + ImmutableList.of(1), + ImmutableList.of("field1"), + HiveType.HIVE_LONG, + BIGINT)), + STRUCT_COLUMN.getColumnType(), + STRUCT_COLUMN.getComment()); + private static final List PROJECTED_COLUMNS = ImmutableList.of(BIGINT_COLUMN, STRUCT_FIELD1_COLUMN); @Test public void testOrcPredicates() @@ -94,152 +116,129 @@ public void testOrcPredicates() private void testOrcPredicates(ConnectorSession session) throws Exception { - List columnsToWrite = ImmutableList.of(columnPrimitiveInteger, columnStruct, columnPrimitiveBigInt); - - File file = File.createTempFile("test", "orc_predicate"); - file.delete(); - try { - // Write data - OrcFileWriterFactory writerFactory = new OrcFileWriterFactory(TESTING_TYPE_MANAGER, new NodeVersion("test"), STATS, new OrcWriterOptions(), HDFS_FILE_SYSTEM_FACTORY); - FileSplit split = createTestFileTrino(file.getAbsolutePath(), ORC, HiveCompressionCodec.NONE, columnsToWrite, session, NUM_ROWS, writerFactory); - - TupleDomain testingPredicate; - - // Verify predicates on base column - List columnsToRead = columnsToWrite; - // All rows returned for a satisfying predicate - testingPredicate = TupleDomain.withColumnDomains(ImmutableMap.of(columnPrimitiveBigInt, Domain.singleValue(BIGINT, 6L))); - assertFilteredRows(testingPredicate, columnsToRead, session, split, NUM_ROWS); - // No rows returned for a mismatched predicate - testingPredicate = TupleDomain.withColumnDomains(ImmutableMap.of(columnPrimitiveBigInt, Domain.singleValue(BIGINT, 1L))); - assertFilteredRows(testingPredicate, columnsToRead, session, split, 0); - - // Verify predicates on projected column - TestColumn projectedColumn = new TestColumn( - columnStruct.getBaseName(), - columnStruct.getBaseObjectInspector(), - ImmutableList.of("field1"), - ImmutableList.of(1), - javaLongObjectInspector, - 5L, - 5L, - false); - - columnsToRead = ImmutableList.of(columnPrimitiveBigInt, projectedColumn); - // All rows returned for a satisfying predicate - testingPredicate = TupleDomain.withColumnDomains(ImmutableMap.of(projectedColumn, Domain.singleValue(BIGINT, 5L))); - assertFilteredRows(testingPredicate, columnsToRead, session, split, NUM_ROWS); - // No rows returned for a mismatched predicate - testingPredicate = TupleDomain.withColumnDomains(ImmutableMap.of(projectedColumn, Domain.singleValue(BIGINT, 6L))); - assertFilteredRows(testingPredicate, columnsToRead, session, split, 0); - } - finally { - file.delete(); - } + TrinoFileSystemFactory fileSystemFactory = new MemoryFileSystemFactory(); + Location location = Location.of("memory:///test"); + writeTestFile(session, fileSystemFactory, location); + + // Verify predicates on base column + // All rows returned for a satisfying predicate + assertFilteredRows(fileSystemFactory, location, TupleDomain.withColumnDomains(ImmutableMap.of(BIGINT_COLUMN, Domain.singleValue(BIGINT, 6L))), COLUMNS, session, NUM_ROWS); + // No rows returned for a mismatched predicate + assertFilteredRows(fileSystemFactory, location, TupleDomain.withColumnDomains(ImmutableMap.of(BIGINT_COLUMN, Domain.singleValue(BIGINT, 1L))), COLUMNS, session, 0); + + // Verify predicates on projected column + // All rows returned for a satisfying predicate + assertFilteredRows(fileSystemFactory, location, TupleDomain.withColumnDomains(ImmutableMap.of(STRUCT_FIELD1_COLUMN, Domain.singleValue(BIGINT, 5L))), PROJECTED_COLUMNS, session, NUM_ROWS); + // No rows returned for a mismatched predicate + assertFilteredRows(fileSystemFactory, location, TupleDomain.withColumnDomains(ImmutableMap.of(STRUCT_FIELD1_COLUMN, Domain.singleValue(BIGINT, 6L))), PROJECTED_COLUMNS, session, 0); } - private void assertFilteredRows( - TupleDomain effectivePredicate, - List columnsToRead, + private static void assertFilteredRows( + TrinoFileSystemFactory fileSystemFactory, + Location location, + TupleDomain effectivePredicate, + List columnsToRead, ConnectorSession session, - FileSplit split, int expectedRows) + throws IOException { - ConnectorPageSource pageSource = createPageSource(effectivePredicate, columnsToRead, session, split); - - int filteredRows = 0; - while (!pageSource.isFinished()) { - Page page = pageSource.getNextPage(); - if (page != null) { - filteredRows += page.getPositionCount(); + try (ConnectorPageSource pageSource = createPageSource(fileSystemFactory, location, effectivePredicate, columnsToRead, session)) { + int filteredRows = 0; + while (!pageSource.isFinished()) { + Page page = pageSource.getNextPage(); + if (page != null) { + filteredRows += page.getPositionCount(); + } } + assertThat(filteredRows).isEqualTo(expectedRows); } - - assertEquals(filteredRows, expectedRows); } - private ConnectorPageSource createPageSource( - TupleDomain effectivePredicate, - List columnsToRead, - ConnectorSession session, - FileSplit split) + private static ConnectorPageSource createPageSource( + TrinoFileSystemFactory fileSystemFactory, + Location location, + TupleDomain effectivePredicate, + List columns, + ConnectorSession session) + throws IOException { - OrcPageSourceFactory readerFactory = new OrcPageSourceFactory(new OrcReaderOptions(), HDFS_FILE_SYSTEM_FACTORY, STATS, UTC); - - Properties splitProperties = new Properties(); - splitProperties.setProperty(FILE_INPUT_FORMAT, ORC.getInputFormat()); - splitProperties.setProperty(SERIALIZATION_LIB, ORC.getSerde()); - - // Use full columns in split properties - ImmutableList.Builder splitPropertiesColumnNames = ImmutableList.builder(); - ImmutableList.Builder splitPropertiesColumnTypes = ImmutableList.builder(); - Set baseColumnNames = new HashSet<>(); - for (TestColumn columnToRead : columnsToRead) { - String name = columnToRead.getBaseName(); - if (!baseColumnNames.contains(name) && !columnToRead.isPartitionKey()) { - baseColumnNames.add(name); - splitPropertiesColumnNames.add(name); - splitPropertiesColumnTypes.add(columnToRead.getBaseObjectInspector().getTypeName()); - } - } - - splitProperties.setProperty("columns", splitPropertiesColumnNames.build().stream().collect(Collectors.joining(","))); - splitProperties.setProperty("columns.types", splitPropertiesColumnTypes.build().stream().collect(Collectors.joining(","))); - - List partitionKeys = columnsToRead.stream() - .filter(TestColumn::isPartitionKey) - .map(input -> new HivePartitionKey(input.getName(), (String) input.getWriteValue())) - .collect(toList()); - - String partitionName = String.join("/", partitionKeys.stream() - .map(partitionKey -> format("%s=%s", partitionKey.getName(), partitionKey.getValue())) - .collect(toImmutableList())); - - List columnHandles = getColumnHandles(columnsToRead); - - TupleDomain predicate = effectivePredicate.transformKeys(testColumn -> { - Optional handle = columnHandles.stream() - .filter(column -> testColumn.getName().equals(column.getName())) - .findFirst(); - - checkState(handle.isPresent(), "Predicate on invalid column"); - return handle.get(); - }); + OrcPageSourceFactory readerFactory = new OrcPageSourceFactory(new OrcReaderOptions(), fileSystemFactory, STATS, UTC); + TrinoInputFile inputFile = fileSystemFactory.create(session).newInputFile(location); + long length = inputFile.length(); List columnMappings = buildColumnMappings( - partitionName, - partitionKeys, - columnHandles, + "", + ImmutableList.of(), + columns, ImmutableList.of(), - TableToPartitionMapping.empty(), - split.getPath().toString(), + ImmutableMap.of(), + location.toString(), OptionalInt.empty(), - split.getLength(), + length, Instant.now().toEpochMilli()); - Optional pageSource = HivePageSourceProvider.createHivePageSource( - ImmutableSet.of(readerFactory), - ImmutableSet.of(), - newEmptyConfiguration(), - session, - Location.of(split.getPath().toString()), - OptionalInt.empty(), - split.getStart(), - split.getLength(), - split.getLength(), - splitProperties, - predicate, - columnHandles, - TESTING_TYPE_MANAGER, - Optional.empty(), - Optional.empty(), - false, - Optional.empty(), - false, - NO_ACID_TRANSACTION, - columnMappings); - - assertTrue(pageSource.isPresent()); - return pageSource.get(); + return HivePageSourceProvider.createHivePageSource( + ImmutableSet.of(readerFactory), + session, + location, + OptionalInt.empty(), + 0, + length, + length, + inputFile.lastModified().toEpochMilli(), + new Schema( + ORC.getSerde(), + false, + ImmutableMap.builder() + .put(LIST_COLUMNS, COLUMNS.stream().map(HiveColumnHandle::getName).collect(Collectors.joining(","))) + .put(LIST_COLUMN_TYPES, COLUMNS.stream().map(HiveColumnHandle::getHiveType).map(HiveType::toString).collect(Collectors.joining(","))) + .buildOrThrow()), + effectivePredicate, + TESTING_TYPE_MANAGER, + Optional.empty(), + Optional.empty(), + Optional.empty(), + false, + NO_ACID_TRANSACTION, + columnMappings) + .orElseThrow(); + } + + private static void writeTestFile(ConnectorSession session, TrinoFileSystemFactory fileSystemFactory, Location location) + { + FileWriter fileWriter = new OrcFileWriterFactory(TESTING_TYPE_MANAGER, new NodeVersion("test"), STATS, new OrcWriterOptions(), fileSystemFactory) + .createFileWriter( + location, + COLUMNS.stream().map(HiveColumnHandle::getName).collect(toList()), + ORC.toStorageFormat(), + HiveCompressionCodec.NONE, + getTableProperties(), + session, + OptionalInt.empty(), + NO_ACID_TRANSACTION, + false, + WriterKind.INSERT) + .orElseThrow(); + + fileWriter.appendRows(new Page( + RunLengthEncodedBlock.create(new IntArrayBlock(1, Optional.empty(), new int[] {3}), NUM_ROWS), + RunLengthEncodedBlock.create( + RowBlock.fromFieldBlocks(1, new Block[] { + new LongArrayBlock(1, Optional.empty(), new long[] {4}), + new LongArrayBlock(1, Optional.empty(), new long[] {5})}), + NUM_ROWS), + RunLengthEncodedBlock.create(new LongArrayBlock(1, Optional.empty(), new long[] {6}), NUM_ROWS))); + + fileWriter.commit(); + } + + private static Map getTableProperties() + { + return ImmutableMap.builder() + .put(hive_metastoreConstants.FILE_INPUT_FORMAT, ORC.getInputFormat()) + .put(SerdeConstants.SERIALIZATION_LIB, ORC.getSerde()) + .put(LIST_COLUMNS, COLUMNS.stream().map(HiveColumnHandle::getName).collect(Collectors.joining(","))) + .put(LIST_COLUMN_TYPES, COLUMNS.stream().map(HiveColumnHandle::getHiveType).map(HiveType::toString).collect(Collectors.joining(","))) + .buildOrThrow(); } } diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/orc/TestOrcWriterOptions.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/orc/TestOrcWriterOptions.java index a6b71af7b07a..b8fe121f7a6d 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/orc/TestOrcWriterOptions.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/orc/TestOrcWriterOptions.java @@ -13,6 +13,7 @@ */ package io.trino.plugin.hive.orc; +import com.google.shaded.common.shaded.collect.Maps; import io.airlift.units.DataSize; import io.trino.orc.OrcWriterOptions; import org.testng.annotations.DataProvider; @@ -80,7 +81,7 @@ public void testOrcWriterOptionsFromTableProperties() Properties tableProperties = new Properties(); tableProperties.setProperty(ORC_BLOOM_FILTER_COLUMNS_KEY, "column_a, column_b"); tableProperties.setProperty(ORC_BLOOM_FILTER_FPP_KEY, "0.5"); - OrcWriterOptions orcWriterOptions = getOrcWriterOptions(tableProperties, new OrcWriterOptions()); + OrcWriterOptions orcWriterOptions = getOrcWriterOptions(Maps.fromProperties(tableProperties), new OrcWriterOptions()); assertThat(orcWriterOptions.getBloomFilterFpp()).isEqualTo(0.5); assertThat(orcWriterOptions.isBloomFilterColumn("column_a")).isTrue(); assertThat(orcWriterOptions.isBloomFilterColumn("column_b")).isTrue(); @@ -93,7 +94,7 @@ public void testOrcWriterOptionsWithInvalidFPPValue() Properties tableProperties = new Properties(); tableProperties.setProperty(ORC_BLOOM_FILTER_COLUMNS_KEY, "column_with_bloom_filter"); tableProperties.setProperty(ORC_BLOOM_FILTER_FPP_KEY, "abc"); - assertThatThrownBy(() -> getOrcWriterOptions(tableProperties, new OrcWriterOptions())) + assertThatThrownBy(() -> getOrcWriterOptions(Maps.fromProperties(tableProperties), new OrcWriterOptions())) .hasMessage("Invalid value for orc_bloom_filter_fpp property: abc"); } @@ -103,7 +104,7 @@ public void testOrcBloomFilterWithInvalidRange(String fpp) Properties tableProperties = new Properties(); tableProperties.setProperty(ORC_BLOOM_FILTER_COLUMNS_KEY, "column_with_bloom_filter"); tableProperties.setProperty(ORC_BLOOM_FILTER_FPP_KEY, fpp); - assertThatThrownBy(() -> getOrcWriterOptions(tableProperties, new OrcWriterOptions())) + assertThatThrownBy(() -> getOrcWriterOptions(Maps.fromProperties(tableProperties), new OrcWriterOptions())) .hasMessage("bloomFilterFpp should be > 0.0 & < 1.0"); } diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/parquet/ParquetTester.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/parquet/ParquetTester.java index 5042ab9eb738..00beb1b93a3d 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/parquet/ParquetTester.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/parquet/ParquetTester.java @@ -29,7 +29,6 @@ import io.trino.parquet.writer.ParquetWriterOptions; import io.trino.plugin.hive.FileFormatDataSourceStats; import io.trino.plugin.hive.HiveConfig; -import io.trino.plugin.hive.HiveFormatsConfig; import io.trino.plugin.hive.HiveSessionProperties; import io.trino.plugin.hive.HiveStorageFormat; import io.trino.plugin.hive.benchmark.FileFormat; @@ -158,10 +157,10 @@ public class ParquetTester private static final int MAX_PRECISION_INT64 = toIntExact(maxPrecision(8)); private static final ConnectorSession SESSION = getHiveSession( - createHiveConfig(false), new ParquetReaderConfig().setOptimizedReaderEnabled(false)); + createHiveConfig(false), new ParquetReaderConfig()); private static final ConnectorSession SESSION_OPTIMIZED_READER = getHiveSession( - createHiveConfig(false), new ParquetReaderConfig().setOptimizedReaderEnabled(true)); + createHiveConfig(false), new ParquetReaderConfig()); private static final ConnectorSession SESSION_USE_NAME = getHiveSession(createHiveConfig(true)); @@ -452,12 +451,10 @@ void assertMaxReadBytes( new HiveConfig() .setHiveStorageFormat(HiveStorageFormat.PARQUET) .setUseParquetColumnNames(false), - new HiveFormatsConfig(), new OrcReaderConfig(), new OrcWriterConfig(), new ParquetReaderConfig() - .setMaxReadBlockSize(maxReadBlockSize) - .setOptimizedReaderEnabled(optimizedReaderEnabled), + .setMaxReadBlockSize(maxReadBlockSize), new ParquetWriterConfig()); ConnectorSession session = TestingConnectorSession.builder() .setPropertyMetadata(hiveSessionProperties.getSessionProperties()) @@ -793,7 +790,6 @@ private static void writeParquetColumnTrino( .build(), compressionCodec, "test-version", - false, Optional.of(DateTimeZone.getDefault()), Optional.of(new ParquetWriteValidationBuilder(types, columnNames))); diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/parquet/TestBloomFilterStore.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/parquet/TestBloomFilterStore.java index 0e3816d7ca35..362c522f1d0b 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/parquet/TestBloomFilterStore.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/parquet/TestBloomFilterStore.java @@ -18,6 +18,8 @@ import io.trino.filesystem.local.LocalInputFile; import io.trino.parquet.BloomFilterStore; import io.trino.parquet.ParquetReaderOptions; +import io.trino.parquet.metadata.ColumnChunkMetadata; +import io.trino.parquet.metadata.ParquetMetadata; import io.trino.parquet.predicate.TupleDomainParquetPredicate; import io.trino.parquet.reader.MetadataReader; import io.trino.plugin.hive.FileFormatDataSourceStats; @@ -32,8 +34,6 @@ import org.apache.parquet.column.ColumnDescriptor; import org.apache.parquet.column.values.bloomfilter.BloomFilter; import org.apache.parquet.format.CompressionCodec; -import org.apache.parquet.hadoop.metadata.ColumnChunkMetaData; -import org.apache.parquet.hadoop.metadata.ParquetMetadata; import org.apache.parquet.schema.PrimitiveType; import org.apache.parquet.schema.Types; import org.joda.time.DateTimeZone; @@ -311,7 +311,7 @@ private static BloomFilterStore generateBloomFilterStore(ParquetTester.TempFile TrinoParquetDataSource dataSource = new TrinoParquetDataSource(inputFile, new ParquetReaderOptions(), new FileFormatDataSourceStats()); ParquetMetadata parquetMetadata = MetadataReader.readFooter(dataSource, Optional.empty()); - ColumnChunkMetaData columnChunkMetaData = getOnlyElement(getOnlyElement(parquetMetadata.getBlocks()).getColumns()); + ColumnChunkMetadata columnChunkMetaData = getOnlyElement(getOnlyElement(parquetMetadata.getBlocks()).columns()); return new BloomFilterStore(dataSource, getOnlyElement(parquetMetadata.getBlocks()), Set.of(columnChunkMetaData.getPath())); } diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/parquet/TestOnlyNulls.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/parquet/TestOnlyNulls.java index c432a3974f8b..5f578caa077c 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/parquet/TestOnlyNulls.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/parquet/TestOnlyNulls.java @@ -13,17 +13,12 @@ */ package io.trino.plugin.hive.parquet; +import com.google.common.collect.ImmutableList; import com.google.common.io.Resources; -import io.trino.filesystem.Location; import io.trino.plugin.hive.HiveColumnHandle; import io.trino.plugin.hive.HiveConfig; -import io.trino.plugin.hive.HivePageSourceFactory; -import io.trino.plugin.hive.HiveStorageFormat; import io.trino.plugin.hive.HiveType; -import io.trino.plugin.hive.acid.AcidTransaction; -import io.trino.plugin.hive.benchmark.StandardFileFormats; import io.trino.spi.connector.ConnectorPageSource; -import io.trino.spi.predicate.Domain; import io.trino.spi.predicate.TupleDomain; import io.trino.spi.type.Type; import io.trino.testing.MaterializedResult; @@ -34,17 +29,17 @@ import java.util.List; import java.util.Map; import java.util.Optional; -import java.util.OptionalInt; -import java.util.Properties; import static io.trino.plugin.hive.HiveColumnHandle.ColumnType.REGULAR; import static io.trino.plugin.hive.HiveColumnHandle.createBaseColumn; -import static io.trino.plugin.hive.HiveTestUtils.HDFS_ENVIRONMENT; +import static io.trino.plugin.hive.HiveTestUtils.SESSION; import static io.trino.plugin.hive.HiveTestUtils.getHiveSession; +import static io.trino.plugin.hive.parquet.ParquetUtil.createPageSource; +import static io.trino.spi.predicate.Domain.notNull; +import static io.trino.spi.predicate.Domain.onlyNull; import static io.trino.spi.type.IntegerType.INTEGER; import static io.trino.testing.MaterializedResult.materializeSourceDataStream; import static java.util.Collections.singletonList; -import static org.apache.hadoop.hive.serde.serdeConstants.SERIALIZATION_LIB; import static org.assertj.core.api.Assertions.assertThat; public class TestOnlyNulls @@ -61,13 +56,13 @@ public void testOnlyNulls() HiveColumnHandle column = createBaseColumn(columnName, 0, HiveType.toHiveType(columnType), columnType, REGULAR, Optional.empty()); // match not null - try (ConnectorPageSource pageSource = createPageSource(parquetFile, column, TupleDomain.withColumnDomains(Map.of(column, Domain.notNull(columnType))))) { + try (ConnectorPageSource pageSource = createPageSource(SESSION, parquetFile, ImmutableList.of(column), TupleDomain.withColumnDomains(Map.of(column, notNull(columnType))))) { MaterializedResult result = materializeSourceDataStream(getHiveSession(new HiveConfig()), pageSource, List.of(columnType)).toTestTypes(); assertThat(result.getMaterializedRows()).isEmpty(); } // match null - try (ConnectorPageSource pageSource = createPageSource(parquetFile, column, TupleDomain.withColumnDomains(Map.of(column, Domain.onlyNull(columnType))))) { + try (ConnectorPageSource pageSource = createPageSource(SESSION, parquetFile, ImmutableList.of(column), TupleDomain.withColumnDomains(Map.of(column, onlyNull(columnType))))) { MaterializedResult result = materializeSourceDataStream(getHiveSession(new HiveConfig()), pageSource, List.of(columnType)).toTestTypes(); assertThat(result.getMaterializedRows()) @@ -78,28 +73,4 @@ public void testOnlyNulls() new MaterializedRow(singletonList(null)))); } } - - private static ConnectorPageSource createPageSource(File parquetFile, HiveColumnHandle column, TupleDomain domain) - { - HivePageSourceFactory pageSourceFactory = StandardFileFormats.TRINO_PARQUET.getHivePageSourceFactory(HDFS_ENVIRONMENT).orElseThrow(); - - Properties schema = new Properties(); - schema.setProperty(SERIALIZATION_LIB, HiveStorageFormat.PARQUET.getSerde()); - - return pageSourceFactory.createPageSource( - getHiveSession(new HiveConfig()), - Location.of(parquetFile.getPath()), - 0, - parquetFile.length(), - parquetFile.length(), - schema, - List.of(column), - domain, - Optional.empty(), - OptionalInt.empty(), - false, - AcidTransaction.NO_ACID_TRANSACTION) - .orElseThrow() - .get(); - } } diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/parquet/TestParquetReaderConfig.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/parquet/TestParquetReaderConfig.java index afa1f038ffb1..3d3315a865f0 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/parquet/TestParquetReaderConfig.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/parquet/TestParquetReaderConfig.java @@ -37,9 +37,8 @@ public void testDefaults() .setMaxMergeDistance(DataSize.of(1, MEGABYTE)) .setMaxBufferSize(DataSize.of(8, MEGABYTE)) .setUseColumnIndex(true) - .setOptimizedReaderEnabled(true) - .setOptimizedNestedReaderEnabled(true) - .setUseBloomFilter(true)); + .setUseBloomFilter(true) + .setSmallFileThreshold(DataSize.of(3, MEGABYTE))); } @Test @@ -52,9 +51,8 @@ public void testExplicitPropertyMappings() .put("parquet.max-buffer-size", "1431kB") .put("parquet.max-merge-distance", "342kB") .put("parquet.use-column-index", "false") - .put("parquet.optimized-reader.enabled", "false") - .put("parquet.optimized-nested-reader.enabled", "false") .put("parquet.use-bloom-filter", "false") + .put("parquet.small-file-threshold", "1kB") .buildOrThrow(); ParquetReaderConfig expected = new ParquetReaderConfig() @@ -64,9 +62,8 @@ public void testExplicitPropertyMappings() .setMaxBufferSize(DataSize.of(1431, KILOBYTE)) .setMaxMergeDistance(DataSize.of(342, KILOBYTE)) .setUseColumnIndex(false) - .setOptimizedReaderEnabled(false) - .setOptimizedNestedReaderEnabled(false) - .setUseBloomFilter(false); + .setUseBloomFilter(false) + .setSmallFileThreshold(DataSize.of(1, KILOBYTE)); assertFullMapping(properties, expected); } diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/parquet/TestReadingTimeLogicalAnnotation.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/parquet/TestReadingTimeLogicalAnnotation.java index dd1500a3b8ea..a932e1988cb9 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/parquet/TestReadingTimeLogicalAnnotation.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/parquet/TestReadingTimeLogicalAnnotation.java @@ -14,13 +14,20 @@ package io.trino.plugin.hive.parquet; import com.google.common.io.Resources; +import io.trino.filesystem.Location; +import io.trino.filesystem.TrinoFileSystem; +import io.trino.filesystem.TrinoFileSystemFactory; import io.trino.plugin.hive.HiveQueryRunner; +import io.trino.spi.security.ConnectorIdentity; import io.trino.sql.query.QueryAssertions; import io.trino.testing.DistributedQueryRunner; import org.testng.annotations.Test; -import java.io.File; +import java.io.OutputStream; +import java.net.URL; +import java.util.UUID; +import static io.trino.plugin.hive.TestingHiveUtils.getConnectorService; import static io.trino.spi.type.BigintType.BIGINT; import static io.trino.testing.MaterializedResult.resultBuilder; import static java.lang.String.format; @@ -32,17 +39,27 @@ public class TestReadingTimeLogicalAnnotation public void testReadingTimeLogicalAnnotationAsBigInt() throws Exception { - File parquetFile = new File(Resources.getResource("parquet_file_with_time_logical_annotation").toURI()); try (DistributedQueryRunner queryRunner = HiveQueryRunner.builder().build(); QueryAssertions assertions = new QueryAssertions(queryRunner)) { + URL resourceLocation = Resources.getResource("parquet_file_with_time_logical_annotation/time-micros.parquet"); + + TrinoFileSystem fileSystem = getConnectorService(queryRunner, TrinoFileSystemFactory.class) + .create(ConnectorIdentity.ofUser("test")); + + Location tempDir = Location.of("local:///temp_" + UUID.randomUUID()); + fileSystem.createDirectory(tempDir); + Location dataFile = tempDir.appendPath("data.parquet"); + try (OutputStream out = fileSystem.newOutputFile(dataFile).create()) { + Resources.copy(resourceLocation, out); + } + queryRunner.execute(format(""" CREATE TABLE table_with_time_logical_annotation ( "opens" row(member0 bigint, member_1 varchar)) WITH ( external_location = '%s', format = 'PARQUET') - """, - parquetFile.getAbsolutePath())); + """.formatted(dataFile.parentDirectory()))); assertThat(assertions.query("SELECT opens.member0 FROM table_with_time_logical_annotation GROUP BY 1 ORDER BY 1 LIMIT 5")) .matches(resultBuilder(queryRunner.getDefaultSession(), BIGINT) diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/parquet/TestTimestamp.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/parquet/TestTimestamp.java index 93e96e0027a7..e559f7531821 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/parquet/TestTimestamp.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/parquet/TestTimestamp.java @@ -19,7 +19,6 @@ import com.google.common.collect.Range; import io.trino.plugin.hive.HiveConfig; import io.trino.plugin.hive.HiveTimestampPrecision; -import io.trino.plugin.hive.benchmark.StandardFileFormats; import io.trino.spi.Page; import io.trino.spi.block.Block; import io.trino.spi.connector.ConnectorPageSource; @@ -44,7 +43,6 @@ import static com.google.common.collect.Iterables.limit; import static com.google.common.collect.Iterables.transform; import static io.trino.hadoop.ConfigurationInstantiator.newEmptyConfiguration; -import static io.trino.plugin.hive.HiveTestUtils.HDFS_ENVIRONMENT; import static io.trino.plugin.hive.HiveTestUtils.getHiveSession; import static io.trino.spi.type.BigintType.BIGINT; import static io.trino.spi.type.TimestampType.createTimestampType; @@ -96,6 +94,7 @@ public void testTimestampBackedByInt64(HiveTimestampPrecision timestamp) private static void testRoundTrip(MessageType parquetSchema, Iterable writeValues, HiveTimestampPrecision timestamp) throws Exception { + DateTimeZone dateTimeZone = DateTimeZone.forID("UTC"); Iterable timestampReadValues = transform(writeValues, value -> { if (value == null) { return null; @@ -128,21 +127,17 @@ private static void testRoundTrip(MessageType parquetSchema, Iterable writ false, DateTimeZone.getDefault()); - ConnectorSession session = getHiveSession(new HiveConfig(), new ParquetReaderConfig().setOptimizedReaderEnabled(false)); - testReadingAs(createTimestampType(timestamp.getPrecision()), session, tempFile, columnNames, timestampReadValues); - testReadingAs(BIGINT, session, tempFile, columnNames, writeValues); - - session = getHiveSession(new HiveConfig(), new ParquetReaderConfig().setOptimizedReaderEnabled(true)); - testReadingAs(createTimestampType(timestamp.getPrecision()), session, tempFile, columnNames, timestampReadValues); - testReadingAs(BIGINT, session, tempFile, columnNames, writeValues); + ConnectorSession session = getHiveSession(new HiveConfig()); + testReadingAs(createTimestampType(timestamp.getPrecision()), session, tempFile, columnNames, timestampReadValues, dateTimeZone); + testReadingAs(BIGINT, session, tempFile, columnNames, writeValues, dateTimeZone); } } - private static void testReadingAs(Type type, ConnectorSession session, ParquetTester.TempFile tempFile, List columnNames, Iterable expectedValues) + private static void testReadingAs(Type type, ConnectorSession session, ParquetTester.TempFile tempFile, List columnNames, Iterable expectedValues, DateTimeZone dateTimeZone) throws IOException { Iterator expected = expectedValues.iterator(); - try (ConnectorPageSource pageSource = StandardFileFormats.TRINO_PARQUET.createFileFormatReader(session, HDFS_ENVIRONMENT, tempFile.getFile(), columnNames, ImmutableList.of(type))) { + try (ConnectorPageSource pageSource = ParquetUtil.createPageSource(session, tempFile.getFile(), columnNames, ImmutableList.of(type), dateTimeZone)) { // skip a page to exercise the decoder's skip() logic Page firstPage = pageSource.getNextPage(); assertTrue(firstPage.getPositionCount() > 0, "Expected first page to have at least 1 row"); diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/parquet/TestTimestampMicros.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/parquet/TestTimestampMicros.java index 5d9f20ad13aa..da72b3934b8e 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/parquet/TestTimestampMicros.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/parquet/TestTimestampMicros.java @@ -14,15 +14,8 @@ package io.trino.plugin.hive.parquet; import com.google.common.io.Resources; -import io.trino.filesystem.Location; import io.trino.plugin.hive.HiveConfig; -import io.trino.plugin.hive.HivePageSourceFactory; -import io.trino.plugin.hive.HiveStorageFormat; import io.trino.plugin.hive.HiveTimestampPrecision; -import io.trino.plugin.hive.HiveType; -import io.trino.plugin.hive.ReaderPageSource; -import io.trino.plugin.hive.acid.AcidTransaction; -import io.trino.plugin.hive.benchmark.StandardFileFormats; import io.trino.spi.connector.ConnectorPageSource; import io.trino.spi.connector.ConnectorSession; import io.trino.spi.predicate.TupleDomain; @@ -37,19 +30,17 @@ import java.time.ZoneId; import java.util.List; import java.util.Optional; -import java.util.OptionalInt; -import java.util.Properties; import static io.trino.plugin.hive.HiveColumnHandle.ColumnType.REGULAR; import static io.trino.plugin.hive.HiveColumnHandle.createBaseColumn; -import static io.trino.plugin.hive.HiveTestUtils.HDFS_ENVIRONMENT; +import static io.trino.plugin.hive.HiveTestUtils.SESSION; import static io.trino.plugin.hive.HiveTestUtils.getHiveSession; import static io.trino.plugin.hive.HiveType.HIVE_TIMESTAMP; +import static io.trino.plugin.hive.parquet.ParquetUtil.createPageSource; import static io.trino.spi.type.TimestampType.createTimestampType; import static io.trino.spi.type.TimestampWithTimeZoneType.createTimestampWithTimeZoneType; import static io.trino.testing.DataProviders.cartesianProduct; import static io.trino.testing.MaterializedResult.materializeSourceDataStream; -import static org.apache.hadoop.hive.serde.serdeConstants.SERIALIZATION_LIB; import static org.assertj.core.api.Assertions.assertThat; public class TestTimestampMicros @@ -60,12 +51,12 @@ public void testTimestampMicros(HiveTimestampPrecision timestampPrecision, Local { ConnectorSession session = getHiveSession( new HiveConfig().setTimestampPrecision(timestampPrecision), - new ParquetReaderConfig().setOptimizedReaderEnabled(useOptimizedParquetReader)); + new ParquetReaderConfig()); File parquetFile = new File(Resources.getResource("issue-5483.parquet").toURI()); Type columnType = createTimestampType(timestampPrecision.getPrecision()); - try (ConnectorPageSource pageSource = createPageSource(session, parquetFile, "created", HIVE_TIMESTAMP, columnType)) { + try (ConnectorPageSource pageSource = createPageSource(SESSION, parquetFile, List.of(createBaseColumn("created", 0, HIVE_TIMESTAMP, columnType, REGULAR, Optional.empty())), TupleDomain.all())) { MaterializedResult result = materializeSourceDataStream(session, pageSource, List.of(columnType)).toTestTypes(); assertThat(result.getMaterializedRows()) .containsOnly(new MaterializedRow(List.of(expected))); @@ -78,12 +69,12 @@ public void testTimestampMicrosAsTimestampWithTimeZone(HiveTimestampPrecision ti { ConnectorSession session = getHiveSession( new HiveConfig().setTimestampPrecision(timestampPrecision), - new ParquetReaderConfig().setOptimizedReaderEnabled(useOptimizedParquetReader)); + new ParquetReaderConfig()); File parquetFile = new File(Resources.getResource("issue-5483.parquet").toURI()); Type columnType = createTimestampWithTimeZoneType(timestampPrecision.getPrecision()); - try (ConnectorPageSource pageSource = createPageSource(session, parquetFile, "created", HIVE_TIMESTAMP, columnType)) { + try (ConnectorPageSource pageSource = createPageSource(SESSION, parquetFile, List.of(createBaseColumn("created", 0, HIVE_TIMESTAMP, columnType, REGULAR, Optional.empty())), TupleDomain.all())) { MaterializedResult result = materializeSourceDataStream(session, pageSource, List.of(columnType)).toTestTypes(); assertThat(result.getMaterializedRows()) .containsOnly(new MaterializedRow(List.of(expected.atZone(ZoneId.of("UTC"))))); @@ -101,36 +92,4 @@ public static Object[][] testTimestampMicrosDataProvider() }, new Object[][] {{true}, {false}}); } - - private ConnectorPageSource createPageSource(ConnectorSession session, File parquetFile, String columnName, HiveType columnHiveType, Type columnType) - { - // TODO after https://github.com/trinodb/trino/pull/5283, replace the method with - // return FileFormat.PRESTO_PARQUET.createFileFormatReader(session, HDFS_ENVIRONMENT, parquetFile, columnNames, columnTypes); - - HivePageSourceFactory pageSourceFactory = StandardFileFormats.TRINO_PARQUET.getHivePageSourceFactory(HDFS_ENVIRONMENT).orElseThrow(); - - Properties schema = new Properties(); - schema.setProperty(SERIALIZATION_LIB, HiveStorageFormat.PARQUET.getSerde()); - - ReaderPageSource pageSourceWithProjections = pageSourceFactory.createPageSource( - session, - Location.of(parquetFile.getPath()), - 0, - parquetFile.length(), - parquetFile.length(), - schema, - List.of(createBaseColumn(columnName, 0, columnHiveType, columnType, REGULAR, Optional.empty())), - TupleDomain.all(), - Optional.empty(), - OptionalInt.empty(), - false, - AcidTransaction.NO_ACID_TRANSACTION) - .orElseThrow(); - - pageSourceWithProjections.getReaderColumns().ifPresent(projections -> { - throw new IllegalStateException("Unexpected projections: " + projections); - }); - - return pageSourceWithProjections.get(); - } } diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/s3/S3HiveQueryRunner.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/s3/S3HiveQueryRunner.java index 05a332b5335c..3bfbb3211eb6 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/s3/S3HiveQueryRunner.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/s3/S3HiveQueryRunner.java @@ -14,11 +14,11 @@ package io.trino.plugin.hive.s3; import com.google.common.collect.ImmutableMap; -import com.google.common.net.HostAndPort; import com.google.errorprone.annotations.CanIgnoreReturnValue; import io.airlift.log.Logger; import io.airlift.units.Duration; import io.trino.plugin.hive.HiveQueryRunner; +import io.trino.plugin.hive.containers.Hive3MinioDataLake; import io.trino.plugin.hive.containers.HiveMinioDataLake; import io.trino.plugin.hive.metastore.thrift.BridgingHiveMetastore; import io.trino.plugin.hive.metastore.thrift.TestingTokenAwareMetastoreClientFactory; @@ -26,6 +26,7 @@ import io.trino.testing.DistributedQueryRunner; import io.trino.tpch.TpchTable; +import java.net.URI; import java.util.Locale; import java.util.Map; @@ -33,6 +34,7 @@ import static io.trino.plugin.hive.TestingThriftHiveMetastoreBuilder.testingThriftHiveMetastoreBuilder; import static io.trino.plugin.hive.security.HiveSecurityModule.ALLOW_ALL; import static io.trino.testing.containers.Minio.MINIO_ACCESS_KEY; +import static io.trino.testing.containers.Minio.MINIO_REGION; import static io.trino.testing.containers.Minio.MINIO_SECRET_KEY; import static java.util.Objects.requireNonNull; @@ -51,7 +53,7 @@ public static DistributedQueryRunner create( } public static DistributedQueryRunner create( - HostAndPort hiveMetastoreEndpoint, + URI hiveMetastoreEndpoint, String s3Endpoint, String s3AccessKey, String s3SecretKey, @@ -74,6 +76,7 @@ public static Builder builder(HiveMinioDataLake hiveMinioDataLake) return builder() .setHiveMetastoreEndpoint(hiveMinioDataLake.getHiveHadoop().getHiveMetastoreEndpoint()) .setS3Endpoint("http://" + hiveMinioDataLake.getMinio().getMinioApiEndpoint()) + .setS3Region(MINIO_REGION) .setS3AccessKey(MINIO_ACCESS_KEY) .setS3SecretKey(MINIO_SECRET_KEY) .setBucketName(hiveMinioDataLake.getBucketName()); @@ -87,16 +90,17 @@ public static Builder builder() public static class Builder extends HiveQueryRunner.Builder { - private HostAndPort hiveMetastoreEndpoint; + private URI hiveMetastoreEndpoint; private Duration thriftMetastoreTimeout = TestingTokenAwareMetastoreClientFactory.TIMEOUT; private ThriftMetastoreConfig thriftMetastoreConfig = new ThriftMetastoreConfig(); + private String s3Region; private String s3Endpoint; private String s3AccessKey; private String s3SecretKey; private String bucketName; @CanIgnoreReturnValue - public Builder setHiveMetastoreEndpoint(HostAndPort hiveMetastoreEndpoint) + public Builder setHiveMetastoreEndpoint(URI hiveMetastoreEndpoint) { this.hiveMetastoreEndpoint = requireNonNull(hiveMetastoreEndpoint, "hiveMetastoreEndpoint is null"); return this; @@ -116,6 +120,13 @@ public Builder setThriftMetastoreConfig(ThriftMetastoreConfig thriftMetastoreCon return this; } + @CanIgnoreReturnValue + public Builder setS3Region(String s3Region) + { + this.s3Region = requireNonNull(s3Region, "s3Region is null"); + return this; + } + @CanIgnoreReturnValue public Builder setS3Endpoint(String s3Endpoint) { @@ -156,15 +167,17 @@ public DistributedQueryRunner build() String lowerCaseS3Endpoint = s3Endpoint.toLowerCase(Locale.ENGLISH); checkArgument(lowerCaseS3Endpoint.startsWith("http://") || lowerCaseS3Endpoint.startsWith("https://"), "Expected http URI for S3 endpoint; got %s", s3Endpoint); - addHiveProperty("hive.s3.endpoint", s3Endpoint); - addHiveProperty("hive.s3.aws-access-key", s3AccessKey); - addHiveProperty("hive.s3.aws-secret-key", s3SecretKey); - addHiveProperty("hive.s3.path-style-access", "true"); + addHiveProperty("fs.native-s3.enabled", "true"); + addHiveProperty("s3.region", s3Region); + addHiveProperty("s3.endpoint", s3Endpoint); + addHiveProperty("s3.aws-access-key", s3AccessKey); + addHiveProperty("s3.aws-secret-key", s3SecretKey); + addHiveProperty("s3.path-style-access", "true"); setMetastore(distributedQueryRunner -> new BridgingHiveMetastore( testingThriftHiveMetastoreBuilder() .metastoreClient(hiveMetastoreEndpoint, thriftMetastoreTimeout) .thriftMetastoreConfig(thriftMetastoreConfig) - .build())); + .build(distributedQueryRunner::registerResource))); setInitialSchemasLocationBase("s3a://" + bucketName); // cannot use s3:// as Hive metastore is not configured to accept it return super.build(); } @@ -173,7 +186,7 @@ public DistributedQueryRunner build() public static void main(String[] args) throws Exception { - HiveMinioDataLake hiveMinioDataLake = new HiveMinioDataLake("tpch"); + HiveMinioDataLake hiveMinioDataLake = new Hive3MinioDataLake("tpch"); hiveMinioDataLake.start(); DistributedQueryRunner queryRunner = S3HiveQueryRunner.builder(hiveMinioDataLake) diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/s3/TestHiveS3MinioQueries.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/s3/TestHiveS3MinioQueries.java index 1c48f13ae3f6..f0b762e51b1b 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/s3/TestHiveS3MinioQueries.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/s3/TestHiveS3MinioQueries.java @@ -14,59 +14,37 @@ package io.trino.plugin.hive.s3; import com.google.common.collect.ImmutableMap; -import io.trino.plugin.hive.HiveQueryRunner; -import io.trino.plugin.hive.NodeVersion; -import io.trino.plugin.hive.metastore.HiveMetastoreConfig; -import io.trino.plugin.hive.metastore.file.FileHiveMetastore; -import io.trino.plugin.hive.metastore.file.FileHiveMetastoreConfig; +import io.trino.plugin.hive.containers.Hive3MinioDataLake; import io.trino.testing.AbstractTestQueryFramework; import io.trino.testing.DataProviders; import io.trino.testing.QueryRunner; -import io.trino.testing.containers.Minio; import org.testng.annotations.AfterClass; import org.testng.annotations.Test; -import java.io.File; import java.util.regex.Matcher; import java.util.regex.Pattern; import static com.google.common.base.Verify.verify; -import static io.trino.plugin.hive.HiveTestUtils.HDFS_ENVIRONMENT; import static io.trino.testing.TestingNames.randomNameSuffix; -import static io.trino.testing.containers.Minio.MINIO_ACCESS_KEY; -import static io.trino.testing.containers.Minio.MINIO_SECRET_KEY; import static java.nio.charset.StandardCharsets.UTF_8; import static org.assertj.core.api.Assertions.assertThat; public class TestHiveS3MinioQueries extends AbstractTestQueryFramework { - private Minio minio; + private Hive3MinioDataLake hiveMinioDataLake; + private String bucketName; @Override protected QueryRunner createQueryRunner() throws Exception { - minio = closeAfterClass(Minio.builder().build()); - minio.start(); + this.bucketName = "test-hive-minio-queries-" + randomNameSuffix(); + this.hiveMinioDataLake = closeAfterClass(new Hive3MinioDataLake(bucketName)); + this.hiveMinioDataLake.start(); - return HiveQueryRunner.builder() - .setMetastore(queryRunner -> { - File baseDir = queryRunner.getCoordinator().getBaseDataDir().resolve("hive_data").toFile(); - return new FileHiveMetastore( - new NodeVersion("testversion"), - HDFS_ENVIRONMENT, - new HiveMetastoreConfig().isHideDeltaLakeTables(), - new FileHiveMetastoreConfig() - .setCatalogDirectory(baseDir.toURI().toString()) - .setDisableLocationChecks(true) // matches Glue behavior - .setMetastoreUser("test")); - }) + return S3HiveQueryRunner.builder(hiveMinioDataLake) .setHiveProperties(ImmutableMap.builder() - .put("hive.s3.aws-access-key", MINIO_ACCESS_KEY) - .put("hive.s3.aws-secret-key", MINIO_SECRET_KEY) - .put("hive.s3.endpoint", minio.getMinioAddress()) - .put("hive.s3.path-style-access", "true") .put("hive.non-managed-table-writes-enabled", "true") .buildOrThrow()) .build(); @@ -75,15 +53,14 @@ protected QueryRunner createQueryRunner() @AfterClass(alwaysRun = true) public void cleanUp() { - minio = null; // closed by closeAfterClass } @Test(dataProviderClass = DataProviders.class, dataProvider = "trueFalse") public void testTableLocationTopOfTheBucket(boolean locationWithTrailingSlash) { String bucketName = "test-bucket-" + randomNameSuffix(); - minio.createBucket(bucketName); - minio.writeFile("We are\nawesome at\nmultiple slashes.".getBytes(UTF_8), bucketName, "a_file"); + hiveMinioDataLake.getMinio().createBucket(bucketName); + hiveMinioDataLake.getMinio().writeFile("We are\nawesome at\nmultiple slashes.".getBytes(UTF_8), bucketName, "a_file"); String location = "s3://%s%s".formatted(bucketName, locationWithTrailingSlash ? "/" : ""); String tableName = "test_table_top_of_bucket_%s_%s".formatted(locationWithTrailingSlash, randomNameSuffix()); diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/s3/TestS3SelectQueries.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/s3/TestS3SelectQueries.java deleted file mode 100644 index c15f91c8f28e..000000000000 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/s3/TestS3SelectQueries.java +++ /dev/null @@ -1,189 +0,0 @@ -/* - * 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.hive.s3; - -import com.google.common.collect.ImmutableList; -import com.google.common.collect.ImmutableMap; -import io.trino.Session; -import io.trino.plugin.hive.HiveQueryRunner; -import io.trino.plugin.hive.NodeVersion; -import io.trino.plugin.hive.metastore.HiveMetastoreConfig; -import io.trino.plugin.hive.metastore.file.FileHiveMetastore; -import io.trino.plugin.hive.metastore.file.FileHiveMetastoreConfig; -import io.trino.testing.AbstractTestQueryFramework; -import io.trino.testing.MaterializedResult; -import io.trino.testing.QueryRunner; -import io.trino.testing.sql.TestTable; -import org.intellij.lang.annotations.Language; -import org.testng.annotations.DataProvider; -import org.testng.annotations.Parameters; -import org.testng.annotations.Test; - -import java.io.File; -import java.util.List; - -import static io.trino.plugin.hive.HiveTestUtils.HDFS_ENVIRONMENT; -import static io.trino.testing.TestingNames.randomNameSuffix; -import static java.util.Objects.requireNonNull; -import static org.assertj.core.api.Assertions.assertThat; -import static org.testng.Assert.assertEquals; - -// The test requires AWS credentials be provided via one of the ways used by the DefaultAWSCredentialsProviderChain. -public class TestS3SelectQueries - extends AbstractTestQueryFramework -{ - private final String bucket; - private final String bucketEndpoint; - - @Parameters({"s3.bucket", "s3.bucket-endpoint"}) - public TestS3SelectQueries(String bucket, String bucketEndpoint) - { - this.bucket = requireNonNull(bucket, "bucket is null"); - this.bucketEndpoint = requireNonNull(bucketEndpoint, "bucketEndpoint is null"); - } - - @Override - protected QueryRunner createQueryRunner() - throws Exception - { - ImmutableMap.Builder hiveProperties = ImmutableMap.builder(); - hiveProperties.put("hive.s3.endpoint", bucketEndpoint); - hiveProperties.put("hive.non-managed-table-writes-enabled", "true"); - hiveProperties.put("hive.s3select-pushdown.experimental-textfile-pushdown-enabled", "true"); - return HiveQueryRunner.builder() - .setHiveProperties(hiveProperties.buildOrThrow()) - .setInitialTables(ImmutableList.of()) - .setMetastore(queryRunner -> { - File baseDir = queryRunner.getCoordinator().getBaseDataDir().resolve("hive_data").toFile(); - return new FileHiveMetastore( - new NodeVersion("testversion"), - HDFS_ENVIRONMENT, - new HiveMetastoreConfig().isHideDeltaLakeTables(), - new FileHiveMetastoreConfig() - .setCatalogDirectory(baseDir.toURI().toString()) - .setMetastoreUser("test") - .setDisableLocationChecks(true)); - }) - .build(); - } - - @Test(dataProvider = "s3SelectFileFormats") - public void testS3SelectPushdown(String tableProperties) - { - Session usingAppendInserts = Session.builder(getSession()) - .setCatalogSessionProperty("hive", "insert_existing_partitions_behavior", "APPEND") - .build(); - List values = ImmutableList.of( - "1, true, 11, 111, 1111, 11111, 'one', DATE '2020-01-01'", - "2, true, 22, 222, 2222, 22222, 'two', DATE '2020-02-02'", - "3, NULL, NULL, NULL, NULL, NULL, NULL, NULL", - "4, false, 44, 444, 4444, 44444, '', DATE '2020-04-04'"); - try (TestTable table = new TestTable( - sql -> getQueryRunner().execute(usingAppendInserts, sql), - "hive.%s.test_s3_select_pushdown".formatted(HiveQueryRunner.TPCH_SCHEMA), - "(id INT, bool_t BOOLEAN, tiny_t TINYINT, small_t SMALLINT, int_t INT, big_t BIGINT, string_t VARCHAR, date_t DATE) " + - "WITH (external_location = 's3://" + bucket + "/test_s3_select_pushdown/test_table_" + randomNameSuffix() + "', " + tableProperties + ")", values)) { - assertS3SelectQuery("SELECT id FROM " + table.getName() + " WHERE bool_t = true", "VALUES 1, 2"); - assertS3SelectQuery("SELECT id FROM " + table.getName() + " WHERE bool_t = false", "VALUES 4"); - assertS3SelectQuery("SELECT id FROM " + table.getName() + " WHERE bool_t IS NULL", "VALUES 3"); - assertS3SelectQuery("SELECT id FROM " + table.getName() + " WHERE bool_t IS NOT NULL", "VALUES 1, 2, 4"); - - assertS3SelectQuery("SELECT id FROM " + table.getName() + " WHERE tiny_t = 22", "VALUES 2"); - assertS3SelectQuery("SELECT id FROM " + table.getName() + " WHERE tiny_t != 22", "VALUES 1, 4"); - assertS3SelectQuery("SELECT id FROM " + table.getName() + " WHERE tiny_t > 22", "VALUES 4"); - assertS3SelectQuery("SELECT id FROM " + table.getName() + " WHERE tiny_t >= 22", "VALUES 2, 4"); - assertS3SelectQuery("SELECT id FROM " + table.getName() + " WHERE tiny_t = 22 OR tiny_t = 44", "VALUES 2, 4"); - assertS3SelectQuery("SELECT id FROM " + table.getName() + " WHERE tiny_t IS NULL OR tiny_t >= 22", "VALUES 2, 3, 4"); - assertS3SelectQuery("SELECT id FROM " + table.getName() + " WHERE tiny_t IS NULL", "VALUES 3"); - assertS3SelectQuery("SELECT id FROM " + table.getName() + " WHERE tiny_t IS NOT NULL", "VALUES 1, 2, 4"); - - assertS3SelectQuery("SELECT id FROM " + table.getName() + " WHERE small_t = 222", "VALUES 2"); - assertS3SelectQuery("SELECT id FROM " + table.getName() + " WHERE small_t != 222", "VALUES 1, 4"); - assertS3SelectQuery("SELECT id FROM " + table.getName() + " WHERE small_t > 222", "VALUES 4"); - assertS3SelectQuery("SELECT id FROM " + table.getName() + " WHERE small_t >= 222", "VALUES 2, 4"); - assertS3SelectQuery("SELECT id FROM " + table.getName() + " WHERE small_t = 222 OR small_t = 444", "VALUES 2, 4"); - assertS3SelectQuery("SELECT id FROM " + table.getName() + " WHERE small_t IS NULL OR small_t >= 222", "VALUES 2, 3, 4"); - assertS3SelectQuery("SELECT id FROM " + table.getName() + " WHERE small_t IS NULL", "VALUES 3"); - assertS3SelectQuery("SELECT id FROM " + table.getName() + " WHERE small_t IS NOT NULL", "VALUES 1, 2, 4"); - - assertS3SelectQuery("SELECT id FROM " + table.getName() + " WHERE int_t = 2222", "VALUES 2"); - assertS3SelectQuery("SELECT id FROM " + table.getName() + " WHERE int_t != 2222", "VALUES 1, 4"); - assertS3SelectQuery("SELECT id FROM " + table.getName() + " WHERE int_t > 2222", "VALUES 4"); - assertS3SelectQuery("SELECT id FROM " + table.getName() + " WHERE int_t >= 2222", "VALUES 2, 4"); - assertS3SelectQuery("SELECT id FROM " + table.getName() + " WHERE int_t = 2222 OR int_t = 4444", "VALUES 2, 4"); - assertS3SelectQuery("SELECT id FROM " + table.getName() + " WHERE int_t IS NULL OR int_t >= 2222", "VALUES 2, 3, 4"); - assertS3SelectQuery("SELECT id FROM " + table.getName() + " WHERE int_t IS NULL", "VALUES 3"); - assertS3SelectQuery("SELECT id FROM " + table.getName() + " WHERE int_t IS NOT NULL", "VALUES 1, 2, 4"); - - assertS3SelectQuery("SELECT id FROM " + table.getName() + " WHERE big_t = 22222", "VALUES 2"); - assertS3SelectQuery("SELECT id FROM " + table.getName() + " WHERE big_t != 22222", "VALUES 1, 4"); - assertS3SelectQuery("SELECT id FROM " + table.getName() + " WHERE big_t > 22222", "VALUES 4"); - assertS3SelectQuery("SELECT id FROM " + table.getName() + " WHERE big_t >= 22222", "VALUES 2, 4"); - assertS3SelectQuery("SELECT id FROM " + table.getName() + " WHERE big_t = 22222 OR big_t = 44444", "VALUES 2, 4"); - assertS3SelectQuery("SELECT id FROM " + table.getName() + " WHERE big_t IS NULL OR big_t >= 22222", "VALUES 2, 3, 4"); - assertS3SelectQuery("SELECT id FROM " + table.getName() + " WHERE big_t IS NULL", "VALUES 3"); - assertS3SelectQuery("SELECT id FROM " + table.getName() + " WHERE big_t IS NOT NULL", "VALUES 1, 2, 4"); - - assertS3SelectQuery("SELECT id FROM " + table.getName() + " WHERE string_t = 'two'", "VALUES 2"); - assertS3SelectQuery("SELECT id FROM " + table.getName() + " WHERE string_t != 'two'", "VALUES 1, 4"); - assertS3SelectQuery("SELECT id FROM " + table.getName() + " WHERE string_t < 'two'", "VALUES 1, 4"); - assertS3SelectQuery("SELECT id FROM " + table.getName() + " WHERE string_t <= 'two'", "VALUES 1, 2, 4"); - assertS3SelectQuery("SELECT id FROM " + table.getName() + " WHERE string_t = 'two' OR string_t = ''", "VALUES 2, 4"); - assertS3SelectQuery("SELECT id FROM " + table.getName() + " WHERE string_t IS NULL OR string_t >= 'two'", "VALUES 2, 3"); - assertS3SelectQuery("SELECT id FROM " + table.getName() + " WHERE string_t IS NULL", "VALUES 3"); - assertS3SelectQuery("SELECT id FROM " + table.getName() + " WHERE string_t IS NOT NULL", "VALUES 1, 2, 4"); - - assertS3SelectQuery("SELECT id FROM " + table.getName() + " WHERE date_t = DATE '2020-02-02'", "VALUES 2"); - assertS3SelectQuery("SELECT id FROM " + table.getName() + " WHERE date_t != DATE '2020-02-02'", "VALUES 1, 4"); - assertS3SelectQuery("SELECT id FROM " + table.getName() + " WHERE date_t > DATE '2020-02-02'", "VALUES 4"); - assertS3SelectQuery("SELECT id FROM " + table.getName() + " WHERE date_t <= DATE '2020-02-02'", "VALUES 1, 2"); - assertS3SelectQuery("SELECT id FROM " + table.getName() + " WHERE date_t = DATE '2020-02-02' OR date_t = DATE '2020-04-04'", "VALUES 2, 4"); - assertS3SelectQuery("SELECT id FROM " + table.getName() + " WHERE date_t IS NULL OR date_t >= DATE '2020-02-02'", "VALUES 2, 3, 4"); - assertS3SelectQuery("SELECT id FROM " + table.getName() + " WHERE date_t IS NULL", "VALUES 3"); - assertS3SelectQuery("SELECT id FROM " + table.getName() + " WHERE date_t IS NOT NULL", "VALUES 1, 2, 4"); - } - } - - private void assertS3SelectQuery(@Language("SQL") String query, @Language("SQL") String expectedValues) - { - Session withS3SelectPushdown = Session.builder(getSession()) - .setCatalogSessionProperty("hive", "s3_select_pushdown_enabled", "true") - .setCatalogSessionProperty("hive", "json_native_reader_enabled", "false") - .setCatalogSessionProperty("hive", "text_file_native_reader_enabled", "false") - .build(); - - MaterializedResult expectedResult = computeActual(expectedValues); - assertQueryStats( - withS3SelectPushdown, - query, - statsWithPushdown -> { - long inputPositionsWithPushdown = statsWithPushdown.getPhysicalInputPositions(); - assertQueryStats( - getSession(), - query, - statsWithoutPushdown -> assertThat(statsWithoutPushdown.getPhysicalInputPositions()).isGreaterThan(inputPositionsWithPushdown), - results -> assertEquals(results.getOnlyColumnAsSet(), expectedResult.getOnlyColumnAsSet())); - }, - results -> assertEquals(results.getOnlyColumnAsSet(), expectedResult.getOnlyColumnAsSet())); - } - - @DataProvider - public static Object[][] s3SelectFileFormats() - { - return new Object[][] { - {"format = 'JSON'"}, - {"format = 'TEXTFILE', textfile_field_separator=',', textfile_field_separator_escape='|', null_format='~'"} - }; - } -} diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/s3/TestS3WrongRegionPicked.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/s3/TestS3WrongRegionPicked.java deleted file mode 100644 index 390fee01923b..000000000000 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/s3/TestS3WrongRegionPicked.java +++ /dev/null @@ -1,47 +0,0 @@ -/* - * 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.hive.s3; - -import com.google.common.collect.ImmutableMap; -import io.trino.plugin.hive.containers.HiveMinioDataLake; -import io.trino.testing.QueryRunner; -import org.testng.annotations.Test; - -import static io.trino.testing.TestingNames.randomNameSuffix; -import static org.assertj.core.api.Assertions.assertThatThrownBy; - -public class TestS3WrongRegionPicked -{ - @Test - public void testS3WrongRegionSelection() - throws Exception - { - // Bucket names are global so a unique one needs to be used. - String bucketName = "test-bucket" + randomNameSuffix(); - - try (HiveMinioDataLake dataLake = new HiveMinioDataLake(bucketName)) { - dataLake.start(); - try (QueryRunner queryRunner = S3HiveQueryRunner.builder(dataLake) - .setHiveProperties(ImmutableMap.of("hive.s3.region", "eu-central-1")) // Different than the default one - .build()) { - String tableName = "s3_region_test_" + randomNameSuffix(); - queryRunner.execute("CREATE TABLE default." + tableName + " (a int) WITH (external_location = 's3://" + bucketName + "/" + tableName + "')"); - assertThatThrownBy(() -> queryRunner.execute("SELECT * FROM default." + tableName)) - .rootCause() - .hasMessageContaining("Status Code: 400") - .hasMessageContaining("Error Code: AuthorizationHeaderMalformed"); // That is how Minio reacts to bad region - } - } - } -} diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/s3/TestTrinoS3FileSystemAccessOperations.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/s3/TestTrinoS3FileSystemAccessOperations.java deleted file mode 100644 index a065406c8d39..000000000000 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/s3/TestTrinoS3FileSystemAccessOperations.java +++ /dev/null @@ -1,213 +0,0 @@ -/* - * 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.hive.s3; - -import com.google.common.collect.HashMultiset; -import com.google.common.collect.ImmutableMap; -import com.google.common.collect.ImmutableMultiset; -import com.google.common.collect.Multiset; -import io.opentelemetry.api.OpenTelemetry; -import io.opentelemetry.sdk.OpenTelemetrySdk; -import io.opentelemetry.sdk.testing.exporter.InMemorySpanExporter; -import io.opentelemetry.sdk.trace.SdkTracerProvider; -import io.opentelemetry.sdk.trace.data.SpanData; -import io.opentelemetry.sdk.trace.export.SimpleSpanProcessor; -import io.trino.hdfs.HdfsConfig; -import io.trino.hdfs.HdfsEnvironment; -import io.trino.hdfs.authentication.NoHdfsAuthentication; -import io.trino.plugin.hive.HiveQueryRunner; -import io.trino.plugin.hive.NodeVersion; -import io.trino.plugin.hive.metastore.HiveMetastoreConfig; -import io.trino.plugin.hive.metastore.file.FileHiveMetastore; -import io.trino.plugin.hive.metastore.file.FileHiveMetastoreConfig; -import io.trino.testing.AbstractTestQueryFramework; -import io.trino.testing.DistributedQueryRunner; -import io.trino.testing.QueryRunner; -import io.trino.testing.containers.Minio; -import org.intellij.lang.annotations.Language; -import org.testng.annotations.AfterClass; -import org.testng.annotations.DataProvider; -import org.testng.annotations.Test; - -import java.io.File; -import java.util.Arrays; - -import static com.google.common.base.Preconditions.checkArgument; -import static io.trino.plugin.hive.HiveQueryRunner.TPCH_SCHEMA; -import static io.trino.plugin.hive.HiveTestUtils.HDFS_CONFIGURATION; -import static io.trino.plugin.hive.util.MultisetAssertions.assertMultisetsEqual; -import static io.trino.testing.DataProviders.toDataProvider; -import static io.trino.testing.TestingNames.randomNameSuffix; -import static io.trino.testing.containers.Minio.MINIO_ACCESS_KEY; -import static io.trino.testing.containers.Minio.MINIO_SECRET_KEY; -import static java.util.stream.Collectors.toCollection; - -@Test(singleThreaded = true) // S3 request counters shares mutable state so can't be run from many threads simultaneously -public class TestTrinoS3FileSystemAccessOperations - extends AbstractTestQueryFramework -{ - private static final String BUCKET = "test-bucket"; - - private Minio minio; - private InMemorySpanExporter spanExporter; - - @Override - protected QueryRunner createQueryRunner() - throws Exception - { - minio = closeAfterClass(Minio.builder().build()); - minio.start(); - minio.createBucket(BUCKET); - - spanExporter = closeAfterClass(InMemorySpanExporter.create()); - - SdkTracerProvider tracerProvider = SdkTracerProvider.builder() - .addSpanProcessor(SimpleSpanProcessor.create(spanExporter)) - .build(); - - OpenTelemetry openTelemetry = OpenTelemetrySdk.builder() - .setTracerProvider(tracerProvider) - .build(); - - return HiveQueryRunner.builder() - .setMetastore(distributedQueryRunner -> { - File baseDir = distributedQueryRunner.getCoordinator().getBaseDataDir().resolve("hive_data").toFile(); - return new FileHiveMetastore( - new NodeVersion("testversion"), - new HdfsEnvironment(HDFS_CONFIGURATION, new HdfsConfig(), new NoHdfsAuthentication()), - new HiveMetastoreConfig().isHideDeltaLakeTables(), - new FileHiveMetastoreConfig() - .setCatalogDirectory(baseDir.toURI().toString()) - .setDisableLocationChecks(true) // matches Glue behavior - .setMetastoreUser("test")); - }) - .setHiveProperties(ImmutableMap.builder() - .put("hive.s3.aws-access-key", MINIO_ACCESS_KEY) - .put("hive.s3.aws-secret-key", MINIO_SECRET_KEY) - .put("hive.s3.endpoint", minio.getMinioAddress()) - .put("hive.s3.path-style-access", "true") - .put("hive.non-managed-table-writes-enabled", "true") - .buildOrThrow()) - .setOpenTelemetry(openTelemetry) - .setInitialSchemasLocationBase("s3://" + BUCKET) - .build(); - } - - @AfterClass(alwaysRun = true) - public void tearDown() - { - // closed by closeAfterClass - spanExporter = null; - minio = null; - } - - @Test(dataProvider = "storageFormats") - public void testSelectWithFilter(StorageFormat format) - { - assertUpdate("DROP TABLE IF EXISTS test_select_from_where"); - String tableLocation = randomTableLocation("test_select_from_where"); - - assertUpdate("CREATE TABLE test_select_from_where WITH (format = '" + format + "', external_location = '" + tableLocation + "') AS SELECT 2 AS age", 1); - - assertFileSystemAccesses("SELECT * FROM test_select_from_where WHERE age = 2", - ImmutableMultiset.builder() - // TODO https://github.com/trinodb/trino/issues/18334 Reduce GetObject call for Parquet format - .addCopies("S3.GetObject", occurrences(format, 1, 2)) - .add("S3.ListObjectsV2") - .addCopies("S3.GetObjectMetadata", occurrences(format, 1, 0)) - .build()); - - assertUpdate("DROP TABLE test_select_from_where"); - } - - @Test(dataProvider = "storageFormats") - public void testSelectPartitionTable(StorageFormat format) - { - assertUpdate("DROP TABLE IF EXISTS test_select_from_partition"); - String tableLocation = randomTableLocation("test_select_from_partition"); - - assertUpdate("CREATE TABLE test_select_from_partition (data int, key varchar)" + - "WITH (partitioned_by = ARRAY['key'], format = '" + format + "', external_location = '" + tableLocation + "')"); - assertUpdate("INSERT INTO test_select_from_partition VALUES (1, 'part1'), (2, 'part2')", 2); - - assertFileSystemAccesses("SELECT * FROM test_select_from_partition", - ImmutableMultiset.builder() - // TODO https://github.com/trinodb/trino/issues/18334 Reduce GetObject call for Parquet format - .addCopies("S3.GetObject", occurrences(format, 2, 4)) - .addCopies("S3.ListObjectsV2", 2) - .addCopies("S3.GetObjectMetadata", occurrences(format, 2, 0)) - .build()); - - assertFileSystemAccesses("SELECT * FROM test_select_from_partition WHERE key = 'part1'", - ImmutableMultiset.builder() - // TODO https://github.com/trinodb/trino/issues/18334 Reduce GetObject call for Parquet format - .addCopies("S3.GetObject", occurrences(format, 1, 2)) - .add("S3.ListObjectsV2") - .addCopies("S3.GetObjectMetadata", occurrences(format, 1, 0)) - .build()); - - assertUpdate("INSERT INTO test_select_from_partition VALUES (11, 'part1')", 1); - assertFileSystemAccesses("SELECT * FROM test_select_from_partition WHERE key = 'part1'", - ImmutableMultiset.builder() - // TODO https://github.com/trinodb/trino/issues/18334 Reduce GetObject call for Parquet format - .addCopies("S3.GetObject", occurrences(format, 2, 4)) - .addCopies("S3.ListObjectsV2", 1) - .addCopies("S3.GetObjectMetadata", occurrences(format, 2, 0)) - .build()); - - assertUpdate("DROP TABLE test_select_from_partition"); - } - - private static String randomTableLocation(String tableName) - { - return "s3://%s/%s/%s-%s".formatted(BUCKET, TPCH_SCHEMA, tableName, randomNameSuffix()); - } - - private void assertFileSystemAccesses(@Language("SQL") String query, Multiset expectedAccesses) - { - DistributedQueryRunner queryRunner = getDistributedQueryRunner(); - spanExporter.reset(); - queryRunner.executeWithQueryId(queryRunner.getDefaultSession(), query); - assertMultisetsEqual(getOperations(), expectedAccesses); - } - - private Multiset getOperations() - { - return spanExporter.getFinishedSpanItems().stream() - .map(SpanData::getName) - .collect(toCollection(HashMultiset::create)); - } - - @DataProvider - public static Object[][] storageFormats() - { - return Arrays.stream(StorageFormat.values()) - .collect(toDataProvider()); - } - - private static int occurrences(StorageFormat tableType, int orcValue, int parquetValue) - { - checkArgument(!(orcValue == parquetValue), "No need to use Occurrences when ORC and Parquet"); - return switch (tableType) { - case ORC -> orcValue; - case PARQUET -> parquetValue; - }; - } - - enum StorageFormat - { - ORC, - PARQUET, - } -} diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/s3select/TestIonSqlQueryBuilder.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/s3select/TestIonSqlQueryBuilder.java deleted file mode 100644 index bb2974310153..000000000000 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/s3select/TestIonSqlQueryBuilder.java +++ /dev/null @@ -1,175 +0,0 @@ -/* - * 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.hive.s3select; - -import com.google.common.collect.ImmutableList; -import com.google.common.collect.ImmutableMap; -import io.airlift.slice.Slices; -import io.trino.plugin.hive.HiveColumnHandle; -import io.trino.plugin.hive.HiveType; -import io.trino.spi.predicate.Domain; -import io.trino.spi.predicate.Range; -import io.trino.spi.predicate.SortedRangeSet; -import io.trino.spi.predicate.TupleDomain; -import io.trino.spi.predicate.ValueSet; -import io.trino.spi.type.DecimalType; -import io.trino.spi.type.TypeManager; -import io.trino.util.DateTimeUtils; -import org.testng.annotations.Test; - -import java.util.List; -import java.util.Optional; - -import static io.trino.plugin.hive.HiveColumnHandle.ColumnType.REGULAR; -import static io.trino.plugin.hive.HiveColumnHandle.createBaseColumn; -import static io.trino.plugin.hive.HiveTestUtils.longDecimal; -import static io.trino.plugin.hive.HiveTestUtils.shortDecimal; -import static io.trino.plugin.hive.HiveType.HIVE_DATE; -import static io.trino.plugin.hive.HiveType.HIVE_DOUBLE; -import static io.trino.plugin.hive.HiveType.HIVE_INT; -import static io.trino.plugin.hive.HiveType.HIVE_STRING; -import static io.trino.plugin.hive.HiveType.HIVE_TIMESTAMP; -import static io.trino.spi.predicate.TupleDomain.withColumnDomains; -import static io.trino.spi.predicate.ValueSet.ofRanges; -import static io.trino.spi.type.BigintType.BIGINT; -import static io.trino.spi.type.DateType.DATE; -import static io.trino.spi.type.DoubleType.DOUBLE; -import static io.trino.spi.type.IntegerType.INTEGER; -import static io.trino.spi.type.TimestampType.TIMESTAMP_MILLIS; -import static io.trino.spi.type.VarcharType.VARCHAR; -import static io.trino.type.InternalTypeManager.TESTING_TYPE_MANAGER; -import static org.testng.Assert.assertEquals; - -public class TestIonSqlQueryBuilder -{ - @Test - public void testBuildSQL() - { - List columns = ImmutableList.of( - createBaseColumn("n_nationkey", 0, HIVE_INT, INTEGER, REGULAR, Optional.empty()), - createBaseColumn("n_name", 1, HIVE_STRING, VARCHAR, REGULAR, Optional.empty()), - createBaseColumn("n_regionkey", 2, HIVE_INT, INTEGER, REGULAR, Optional.empty())); - - // CSV - IonSqlQueryBuilder queryBuilder = new IonSqlQueryBuilder(TESTING_TYPE_MANAGER, S3SelectDataType.CSV, Optional.empty()); - assertEquals(queryBuilder.buildSql(columns, TupleDomain.all()), - "SELECT s._1, s._2, s._3 FROM S3Object s"); - - TupleDomain tupleDomain = withColumnDomains(ImmutableMap.of( - columns.get(2), Domain.create(SortedRangeSet.copyOf(BIGINT, ImmutableList.of(Range.equal(BIGINT, 3L))), false))); - assertEquals(queryBuilder.buildSql(columns, tupleDomain), - "SELECT s._1, s._2, s._3 FROM S3Object s WHERE (s._3 != '' AND CAST(s._3 AS INT) = 3)"); - - // JSON - queryBuilder = new IonSqlQueryBuilder(TESTING_TYPE_MANAGER, S3SelectDataType.JSON, Optional.empty()); - assertEquals(queryBuilder.buildSql(columns, TupleDomain.all()), - "SELECT s.n_nationkey, s.n_name, s.n_regionkey FROM S3Object s"); - assertEquals(queryBuilder.buildSql(columns, tupleDomain), - "SELECT s.n_nationkey, s.n_name, s.n_regionkey FROM S3Object s WHERE (s.n_regionkey IS NOT NULL AND CAST(s.n_regionkey AS INT) = 3)"); - } - - @Test - public void testEmptyColumns() - { - // CSV - IonSqlQueryBuilder queryBuilder = new IonSqlQueryBuilder(TESTING_TYPE_MANAGER, S3SelectDataType.CSV, Optional.empty()); - assertEquals(queryBuilder.buildSql(ImmutableList.of(), TupleDomain.all()), "SELECT ' ' FROM S3Object s"); - - // JSON - queryBuilder = new IonSqlQueryBuilder(TESTING_TYPE_MANAGER, S3SelectDataType.JSON, Optional.empty()); - assertEquals(queryBuilder.buildSql(ImmutableList.of(), TupleDomain.all()), "SELECT ' ' FROM S3Object s"); - } - - @Test - public void testDecimalColumns() - { - TypeManager typeManager = TESTING_TYPE_MANAGER; - List columns = ImmutableList.of( - createBaseColumn("quantity", 0, HiveType.valueOf("decimal(20,0)"), DecimalType.createDecimalType(), REGULAR, Optional.empty()), - createBaseColumn("extendedprice", 1, HiveType.valueOf("decimal(20,2)"), DecimalType.createDecimalType(), REGULAR, Optional.empty()), - createBaseColumn("discount", 2, HiveType.valueOf("decimal(10,2)"), DecimalType.createDecimalType(), REGULAR, Optional.empty())); - DecimalType decimalType = DecimalType.createDecimalType(10, 2); - TupleDomain tupleDomain = withColumnDomains( - ImmutableMap.of( - columns.get(0), Domain.create(ofRanges(Range.lessThan(DecimalType.createDecimalType(20, 0), longDecimal("50"))), false), - columns.get(1), Domain.create(ofRanges(Range.equal(HiveType.valueOf("decimal(20,2)").getType(typeManager), longDecimal("0.05"))), false), - columns.get(2), Domain.create(ofRanges(Range.range(decimalType, shortDecimal("0.0"), true, shortDecimal("0.02"), true)), false))); - - // CSV - IonSqlQueryBuilder queryBuilder = new IonSqlQueryBuilder(typeManager, S3SelectDataType.CSV, Optional.empty()); - assertEquals(queryBuilder.buildSql(columns, tupleDomain), "SELECT s._1, s._2, s._3 FROM S3Object s"); - - // JSON - queryBuilder = new IonSqlQueryBuilder(typeManager, S3SelectDataType.JSON, Optional.empty()); - assertEquals(queryBuilder.buildSql(columns, tupleDomain), "SELECT s.quantity, s.extendedprice, s.discount FROM S3Object s"); - } - - @Test - public void testDateColumn() - { - List columns = ImmutableList.of( - createBaseColumn("t1", 0, HIVE_TIMESTAMP, TIMESTAMP_MILLIS, REGULAR, Optional.empty()), - createBaseColumn("t2", 1, HIVE_DATE, DATE, REGULAR, Optional.empty())); - TupleDomain tupleDomain = withColumnDomains(ImmutableMap.of( - columns.get(1), Domain.create(SortedRangeSet.copyOf(DATE, ImmutableList.of(Range.equal(DATE, (long) DateTimeUtils.parseDate("2001-08-22")))), false))); - - // CSV - IonSqlQueryBuilder queryBuilder = new IonSqlQueryBuilder(TESTING_TYPE_MANAGER, S3SelectDataType.CSV, Optional.empty()); - assertEquals(queryBuilder.buildSql(columns, tupleDomain), "SELECT s._1, s._2 FROM S3Object s WHERE (s._2 != '' AND s._2 = '2001-08-22')"); - - // JSON - queryBuilder = new IonSqlQueryBuilder(TESTING_TYPE_MANAGER, S3SelectDataType.JSON, Optional.empty()); - assertEquals(queryBuilder.buildSql(columns, tupleDomain), "SELECT s.t1, s.t2 FROM S3Object s WHERE (s.t2 IS NOT NULL AND s.t2 = '2001-08-22')"); - } - - @Test - public void testNotPushDoublePredicates() - { - List columns = ImmutableList.of( - createBaseColumn("quantity", 0, HIVE_INT, INTEGER, REGULAR, Optional.empty()), - createBaseColumn("extendedprice", 1, HIVE_DOUBLE, DOUBLE, REGULAR, Optional.empty()), - createBaseColumn("discount", 2, HIVE_DOUBLE, DOUBLE, REGULAR, Optional.empty())); - TupleDomain tupleDomain = withColumnDomains( - ImmutableMap.of( - columns.get(0), Domain.create(ofRanges(Range.lessThan(BIGINT, 50L)), false), - columns.get(1), Domain.create(ofRanges(Range.equal(DOUBLE, 0.05)), false), - columns.get(2), Domain.create(ofRanges(Range.range(DOUBLE, 0.0, true, 0.02, true)), false))); - - // CSV - IonSqlQueryBuilder queryBuilder = new IonSqlQueryBuilder(TESTING_TYPE_MANAGER, S3SelectDataType.CSV, Optional.empty()); - assertEquals(queryBuilder.buildSql(columns, tupleDomain), "SELECT s._1, s._2, s._3 FROM S3Object s WHERE (s._1 != '' AND CAST(s._1 AS INT) < 50)"); - - // JSON - queryBuilder = new IonSqlQueryBuilder(TESTING_TYPE_MANAGER, S3SelectDataType.JSON, Optional.empty()); - assertEquals(queryBuilder.buildSql(columns, tupleDomain), "SELECT s.quantity, s.extendedprice, s.discount FROM S3Object s WHERE (s.quantity IS NOT NULL AND CAST(s.quantity AS INT) < 50)"); - } - - @Test - public void testStringEscaping() - { - List columns = ImmutableList.of( - createBaseColumn("string", 0, HIVE_STRING, VARCHAR, REGULAR, Optional.empty())); - TupleDomain tupleDomain = withColumnDomains(ImmutableMap.of( - columns.get(0), - Domain.create(ValueSet.of(VARCHAR, Slices.utf8Slice("value with a ' quote")), false))); - - // CSV - IonSqlQueryBuilder queryBuilder = new IonSqlQueryBuilder(TESTING_TYPE_MANAGER, S3SelectDataType.CSV, Optional.empty()); - assertEquals(queryBuilder.buildSql(columns, tupleDomain), "SELECT s._1 FROM S3Object s WHERE (s._1 != '' AND s._1 = 'value with a '' quote')"); - - // JSON - queryBuilder = new IonSqlQueryBuilder(TESTING_TYPE_MANAGER, S3SelectDataType.JSON, Optional.empty()); - assertEquals(queryBuilder.buildSql(columns, tupleDomain), "SELECT s.string FROM S3Object s WHERE (s.string IS NOT NULL AND s.string = 'value with a '' quote')"); - } -} diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/s3select/TestS3SelectPushdown.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/s3select/TestS3SelectPushdown.java deleted file mode 100644 index ceb44e81b4ad..000000000000 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/s3select/TestS3SelectPushdown.java +++ /dev/null @@ -1,297 +0,0 @@ -/* - * 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.hive.s3select; - -import io.trino.plugin.hive.metastore.Column; -import io.trino.plugin.hive.metastore.Partition; -import io.trino.plugin.hive.metastore.Storage; -import io.trino.plugin.hive.metastore.StorageFormat; -import io.trino.plugin.hive.metastore.Table; -import io.trino.spi.connector.ConnectorSession; -import io.trino.testing.TestingConnectorSession; -import org.apache.hadoop.hive.serde2.lazy.LazySimpleSerDe; -import org.apache.hadoop.mapred.TextInputFormat; -import org.apache.hive.hcatalog.data.JsonSerDe; -import org.testng.annotations.AfterClass; -import org.testng.annotations.BeforeClass; -import org.testng.annotations.Test; - -import java.util.List; -import java.util.Map; -import java.util.Optional; -import java.util.OptionalLong; -import java.util.Properties; - -import static io.trino.hive.thrift.metastore.hive_metastoreConstants.FILE_INPUT_FORMAT; -import static io.trino.plugin.hive.HiveMetadata.SKIP_FOOTER_COUNT_KEY; -import static io.trino.plugin.hive.HiveMetadata.SKIP_HEADER_COUNT_KEY; -import static io.trino.plugin.hive.HiveStorageFormat.ORC; -import static io.trino.plugin.hive.HiveStorageFormat.TEXTFILE; -import static io.trino.plugin.hive.HiveType.HIVE_BINARY; -import static io.trino.plugin.hive.HiveType.HIVE_BOOLEAN; -import static io.trino.plugin.hive.metastore.MetastoreUtil.getHiveSchema; -import static io.trino.plugin.hive.metastore.StorageFormat.fromHiveStorageFormat; -import static io.trino.plugin.hive.s3select.S3SelectPushdown.isCompressionCodecSupported; -import static io.trino.plugin.hive.s3select.S3SelectPushdown.isSplittable; -import static io.trino.plugin.hive.s3select.S3SelectPushdown.shouldEnablePushdownForTable; -import static io.trino.spi.session.PropertyMetadata.booleanProperty; -import static java.util.Collections.emptyList; -import static java.util.Collections.emptyMap; -import static java.util.Collections.singletonList; -import static org.apache.hadoop.hive.serde.serdeConstants.SERIALIZATION_LIB; -import static org.testng.Assert.assertFalse; -import static org.testng.Assert.assertTrue; - -public class TestS3SelectPushdown -{ - private static final String S3_SELECT_PUSHDOWN_ENABLED = "s3_select_pushdown_enabled"; - - private ConnectorSession session; - private Table table; - private Partition partition; - private Storage storage; - private Column column; - private Properties schema; - - @BeforeClass - public void setUp() - { - session = TestingConnectorSession.builder() - .setPropertyMetadata(List.of(booleanProperty( - S3_SELECT_PUSHDOWN_ENABLED, - "S3 Select pushdown enabled", - true, - false))) - .setPropertyValues(Map.of(S3_SELECT_PUSHDOWN_ENABLED, true)) - .build(); - - column = new Column("column", HIVE_BOOLEAN, Optional.empty()); - - storage = Storage.builder() - .setStorageFormat(fromHiveStorageFormat(TEXTFILE)) - .setLocation("location") - .build(); - - partition = new Partition( - "db", - "table", - emptyList(), - storage, - singletonList(column), - emptyMap()); - - table = new Table( - "db", - "table", - Optional.of("owner"), - "type", - storage, - singletonList(column), - emptyList(), - emptyMap(), - Optional.empty(), - Optional.empty(), - OptionalLong.empty()); - - schema = getHiveSchema(partition, table); - } - - @Test - public void testIsCompressionCodecSupported() - { - assertTrue(isCompressionCodecSupported(schema, "s3://fakeBucket/fakeObject.gz")); - assertTrue(isCompressionCodecSupported(schema, "s3://fakeBucket/fakeObject")); - assertFalse(isCompressionCodecSupported(schema, "s3://fakeBucket/fakeObject.lz4")); - assertFalse(isCompressionCodecSupported(schema, "s3://fakeBucket/fakeObject.snappy")); - assertTrue(isCompressionCodecSupported(schema, "s3://fakeBucket/fakeObject.bz2")); - } - - @Test - public void testShouldEnableSelectPushdown() - { - assertTrue(shouldEnablePushdownForTable(session, table, "s3://fakeBucket/fakeObject", Optional.empty())); - assertTrue(shouldEnablePushdownForTable(session, table, "s3://fakeBucket/fakeObject", Optional.of(partition))); - } - - @Test - public void testShouldNotEnableSelectPushdownWhenDisabledOnSession() - { - ConnectorSession testSession = TestingConnectorSession.builder() - .setPropertyMetadata(List.of(booleanProperty( - S3_SELECT_PUSHDOWN_ENABLED, - "S3 Select pushdown enabled", - false, - false))) - .setPropertyValues(Map.of(S3_SELECT_PUSHDOWN_ENABLED, false)) - .build(); - assertFalse(shouldEnablePushdownForTable(testSession, table, "", Optional.empty())); - } - - @Test - public void testShouldNotEnableSelectPushdownWhenIsNotS3StoragePath() - { - assertFalse(shouldEnablePushdownForTable(session, table, null, Optional.empty())); - assertFalse(shouldEnablePushdownForTable(session, table, "", Optional.empty())); - assertFalse(shouldEnablePushdownForTable(session, table, "s3:/invalid", Optional.empty())); - assertFalse(shouldEnablePushdownForTable(session, table, "s3:/invalid", Optional.of(partition))); - } - - @Test - public void testShouldNotEnableSelectPushdownWhenIsNotSupportedSerde() - { - Storage newStorage = Storage.builder() - .setStorageFormat(fromHiveStorageFormat(ORC)) - .setLocation("location") - .build(); - Table newTable = new Table( - "db", - "table", - Optional.of("owner"), - "type", - newStorage, - singletonList(column), - emptyList(), - emptyMap(), - Optional.empty(), - Optional.empty(), - OptionalLong.empty()); - - assertFalse(shouldEnablePushdownForTable(session, newTable, "s3://fakeBucket/fakeObject", Optional.empty())); - - Partition newPartition = new Partition("db", - "table", - emptyList(), - newStorage, - singletonList(column), - emptyMap()); - assertFalse(shouldEnablePushdownForTable(session, newTable, "s3://fakeBucket/fakeObject", Optional.of(newPartition))); - } - - @Test - public void testShouldNotEnableSelectPushdownWhenIsNotSupportedInputFormat() - { - Storage newStorage = Storage.builder() - .setStorageFormat(StorageFormat.create(LazySimpleSerDe.class.getName(), "inputFormat", "outputFormat")) - .setLocation("location") - .build(); - Table newTable = new Table("db", - "table", - Optional.of("owner"), - "type", - newStorage, - singletonList(column), - emptyList(), - emptyMap(), - Optional.empty(), - Optional.empty(), - OptionalLong.empty()); - assertFalse(shouldEnablePushdownForTable(session, newTable, "s3://fakeBucket/fakeObject", Optional.empty())); - - Partition newPartition = new Partition("db", - "table", - emptyList(), - newStorage, - singletonList(column), - emptyMap()); - assertFalse(shouldEnablePushdownForTable(session, newTable, "s3://fakeBucket/fakeObject", Optional.of(newPartition))); - - newStorage = Storage.builder() - .setStorageFormat(StorageFormat.create(LazySimpleSerDe.class.getName(), TextInputFormat.class.getName(), "outputFormat")) - .setLocation("location") - .build(); - newTable = new Table("db", - "table", - Optional.of("owner"), - "type", - newStorage, - singletonList(column), - emptyList(), - Map.of(SKIP_HEADER_COUNT_KEY, "1"), - Optional.empty(), - Optional.empty(), - OptionalLong.empty()); - assertFalse(shouldEnablePushdownForTable(session, newTable, "s3://fakeBucket/fakeObject", Optional.empty())); - - newTable = new Table("db", - "table", - Optional.of("owner"), - "type", - newStorage, - singletonList(column), - emptyList(), - Map.of(SKIP_FOOTER_COUNT_KEY, "1"), - Optional.empty(), - Optional.empty(), - OptionalLong.empty()); - assertFalse(shouldEnablePushdownForTable(session, newTable, "s3://fakeBucket/fakeObject", Optional.empty())); - } - - @Test - public void testShouldNotEnableSelectPushdownWhenColumnTypesAreNotSupported() - { - Column newColumn = new Column("column", HIVE_BINARY, Optional.empty()); - Table newTable = new Table("db", - "table", - Optional.of("owner"), - "type", - storage, - singletonList(newColumn), - emptyList(), - emptyMap(), - Optional.empty(), - Optional.empty(), - OptionalLong.empty()); - assertFalse(shouldEnablePushdownForTable(session, newTable, "s3://fakeBucket/fakeObject", Optional.empty())); - - Partition newPartition = new Partition("db", - "table", - emptyList(), - storage, - singletonList(newColumn), - emptyMap()); - assertFalse(shouldEnablePushdownForTable(session, newTable, "s3://fakeBucket/fakeObject", Optional.of(newPartition))); - } - - @Test - public void testShouldEnableSplits() - { - // Uncompressed CSV - assertTrue(isSplittable(true, schema, "s3://fakeBucket/fakeObject.csv")); - // Pushdown disabled - assertTrue(isSplittable(false, schema, "s3://fakeBucket/fakeObject.csv")); - // JSON - Properties jsonSchema = new Properties(); - jsonSchema.setProperty(FILE_INPUT_FORMAT, TextInputFormat.class.getName()); - jsonSchema.setProperty(SERIALIZATION_LIB, JsonSerDe.class.getName()); - assertTrue(isSplittable(true, jsonSchema, "s3://fakeBucket/fakeObject.json")); - } - - @Test - public void testShouldNotEnableSplits() - { - // Compressed file - assertFalse(isSplittable(true, schema, "s3://fakeBucket/fakeObject.gz")); - } - - @AfterClass(alwaysRun = true) - public void tearDown() - { - session = null; - table = null; - partition = null; - storage = null; - column = null; - schema = null; - } -} diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/s3select/TestS3SelectRecordCursorProvider.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/s3select/TestS3SelectRecordCursorProvider.java deleted file mode 100644 index b7f263575ff3..000000000000 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/s3select/TestS3SelectRecordCursorProvider.java +++ /dev/null @@ -1,214 +0,0 @@ -/* - * 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.hive.s3select; - -import com.google.common.collect.ImmutableList; -import com.google.common.collect.ImmutableMap; -import io.trino.filesystem.Location; -import io.trino.hadoop.ConfigurationInstantiator; -import io.trino.plugin.hive.HiveColumnHandle; -import io.trino.plugin.hive.HiveConfig; -import io.trino.plugin.hive.HiveRecordCursorProvider.ReaderRecordCursorWithProjections; -import io.trino.plugin.hive.TestBackgroundHiveSplitLoader.TestingHdfsEnvironment; -import io.trino.spi.predicate.Domain; -import io.trino.spi.predicate.Range; -import io.trino.spi.predicate.SortedRangeSet; -import io.trino.spi.predicate.TupleDomain; -import org.apache.hadoop.hive.serde2.lazy.LazySimpleSerDe; -import org.apache.hive.hcatalog.data.JsonSerDe; -import org.testng.annotations.Test; - -import java.util.ArrayList; -import java.util.List; -import java.util.Optional; -import java.util.Properties; -import java.util.function.Function; -import java.util.stream.Collectors; - -import static io.trino.plugin.hive.HiveColumnHandle.ColumnType.REGULAR; -import static io.trino.plugin.hive.HiveColumnHandle.createBaseColumn; -import static io.trino.plugin.hive.HiveTestUtils.SESSION; -import static io.trino.plugin.hive.HiveType.HIVE_INT; -import static io.trino.plugin.hive.HiveType.HIVE_STRING; -import static io.trino.spi.predicate.TupleDomain.withColumnDomains; -import static io.trino.spi.type.BigintType.BIGINT; -import static io.trino.spi.type.DateType.DATE; -import static io.trino.spi.type.IntegerType.INTEGER; -import static io.trino.spi.type.VarcharType.VARCHAR; -import static io.trino.type.InternalTypeManager.TESTING_TYPE_MANAGER; -import static org.apache.hadoop.hive.serde.serdeConstants.LIST_COLUMNS; -import static org.apache.hadoop.hive.serde.serdeConstants.LIST_COLUMN_TYPES; -import static org.apache.hadoop.hive.serde.serdeConstants.SERIALIZATION_LIB; -import static org.testng.Assert.assertTrue; - -public class TestS3SelectRecordCursorProvider -{ - private static final HiveColumnHandle ARTICLE_COLUMN = createBaseColumn("article", 1, HIVE_STRING, VARCHAR, REGULAR, Optional.empty()); - private static final HiveColumnHandle AUTHOR_COLUMN = createBaseColumn("author", 1, HIVE_STRING, VARCHAR, REGULAR, Optional.empty()); - private static final HiveColumnHandle DATE_ARTICLE_COLUMN = createBaseColumn("date_pub", 1, HIVE_INT, DATE, REGULAR, Optional.empty()); - private static final HiveColumnHandle QUANTITY_COLUMN = createBaseColumn("quantity", 1, HIVE_INT, INTEGER, REGULAR, Optional.empty()); - - @Test - public void shouldReturnSelectRecordCursor() - { - List readerColumns = new ArrayList<>(); - TupleDomain effectivePredicate = TupleDomain.all(); - Optional recordCursor = - getRecordCursor(effectivePredicate, readerColumns, true); - assertTrue(recordCursor.isPresent()); - } - - @Test - public void shouldReturnSelectRecordCursorWhenEffectivePredicateExists() - { - TupleDomain effectivePredicate = withColumnDomains(ImmutableMap.of(QUANTITY_COLUMN, - Domain.create(SortedRangeSet.copyOf(BIGINT, ImmutableList.of(Range.equal(BIGINT, 3L))), false))); - Optional recordCursor = - getRecordCursor(effectivePredicate, getAllColumns(), true); - assertTrue(recordCursor.isPresent()); - } - - @Test - public void shouldReturnSelectRecordCursorWhenProjectionExists() - { - TupleDomain effectivePredicate = TupleDomain.all(); - List readerColumns = ImmutableList.of(QUANTITY_COLUMN, AUTHOR_COLUMN, ARTICLE_COLUMN); - Optional recordCursor = - getRecordCursor(effectivePredicate, readerColumns, true); - assertTrue(recordCursor.isPresent()); - } - - @Test - public void shouldNotReturnSelectRecordCursorWhenPushdownIsDisabled() - { - List readerColumns = new ArrayList<>(); - TupleDomain effectivePredicate = TupleDomain.all(); - Optional recordCursor = - getRecordCursor(effectivePredicate, readerColumns, false); - assertTrue(recordCursor.isEmpty()); - } - - @Test - public void shouldNotReturnSelectRecordCursorWhenQueryIsNotFiltering() - { - TupleDomain effectivePredicate = TupleDomain.all(); - Optional recordCursor = - getRecordCursor(effectivePredicate, getAllColumns(), true); - assertTrue(recordCursor.isEmpty()); - } - - @Test - public void shouldNotReturnSelectRecordCursorWhenProjectionOrderIsDifferent() - { - TupleDomain effectivePredicate = TupleDomain.all(); - List readerColumns = ImmutableList.of(DATE_ARTICLE_COLUMN, QUANTITY_COLUMN, ARTICLE_COLUMN, AUTHOR_COLUMN); - Optional recordCursor = - getRecordCursor(effectivePredicate, readerColumns, true); - assertTrue(recordCursor.isEmpty()); - } - - @Test - public void testDisableExperimentalFeatures() - { - List readerColumns = new ArrayList<>(); - TupleDomain effectivePredicate = TupleDomain.all(); - S3SelectRecordCursorProvider s3SelectRecordCursorProvider = new S3SelectRecordCursorProvider( - new TestingHdfsEnvironment(new ArrayList<>()), - new TrinoS3ClientFactory(new HiveConfig()), - new HiveConfig().setS3SelectExperimentalPushdownEnabled(false)); - - Optional csvRecordCursor = s3SelectRecordCursorProvider.createRecordCursor( - ConfigurationInstantiator.newEmptyConfiguration(), - SESSION, - Location.of("s3://fakeBucket/fakeObject.gz"), - 0, - 10, - 10, - createTestingSchema(LazySimpleSerDe.class.getName()), - readerColumns, - effectivePredicate, - TESTING_TYPE_MANAGER, - true); - assertTrue(csvRecordCursor.isEmpty()); - - Optional jsonRecordCursor = s3SelectRecordCursorProvider.createRecordCursor( - ConfigurationInstantiator.newEmptyConfiguration(), - SESSION, - Location.of("s3://fakeBucket/fakeObject.gz"), - 0, - 10, - 10, - createTestingSchema(JsonSerDe.class.getName()), - readerColumns, - effectivePredicate, - TESTING_TYPE_MANAGER, - true); - assertTrue(jsonRecordCursor.isPresent()); - } - - private static Optional getRecordCursor(TupleDomain effectivePredicate, - List readerColumns, - boolean s3SelectPushdownEnabled) - { - S3SelectRecordCursorProvider s3SelectRecordCursorProvider = new S3SelectRecordCursorProvider( - new TestingHdfsEnvironment(new ArrayList<>()), - new TrinoS3ClientFactory(new HiveConfig()), - new HiveConfig().setS3SelectExperimentalPushdownEnabled(true)); - - return s3SelectRecordCursorProvider.createRecordCursor( - ConfigurationInstantiator.newEmptyConfiguration(), - SESSION, - Location.of("s3://fakeBucket/fakeObject.gz"), - 0, - 10, - 10, - createTestingSchema(), - readerColumns, - effectivePredicate, - TESTING_TYPE_MANAGER, - s3SelectPushdownEnabled); - } - - private static Properties createTestingSchema() - { - return createTestingSchema(LazySimpleSerDe.class.getName()); - } - - private static Properties createTestingSchema(String serdeClassName) - { - List schemaColumns = getAllColumns(); - Properties schema = new Properties(); - String columnNames = buildPropertyFromColumns(schemaColumns, HiveColumnHandle::getName); - String columnTypeNames = buildPropertyFromColumns(schemaColumns, column -> column.getHiveType().getTypeInfo().getTypeName()); - schema.setProperty(LIST_COLUMNS, columnNames); - schema.setProperty(LIST_COLUMN_TYPES, columnTypeNames); - schema.setProperty(SERIALIZATION_LIB, serdeClassName); - return schema; - } - - private static String buildPropertyFromColumns(List columns, Function mapper) - { - if (columns.isEmpty()) { - return ""; - } - return columns.stream() - .map(mapper) - .collect(Collectors.joining(",")); - } - - private static List getAllColumns() - { - return ImmutableList.of(ARTICLE_COLUMN, AUTHOR_COLUMN, DATE_ARTICLE_COLUMN, QUANTITY_COLUMN); - } -} diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/statistics/TestMetastoreHiveStatisticsProvider.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/statistics/TestMetastoreHiveStatisticsProvider.java index d24bf491c6c9..7c53c1a47a85 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/statistics/TestMetastoreHiveStatisticsProvider.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/statistics/TestMetastoreHiveStatisticsProvider.java @@ -146,12 +146,12 @@ public void testValidatePartitionStatistics() .setColumnStatistics(ImmutableMap.of(COLUMN, HiveColumnStatistics.builder().setMaxValueSizeInBytes(-1).build())) .build(), invalidColumnStatistics("maxValueSizeInBytes must be greater than or equal to zero: -1")); - assertInvalidStatistics( - PartitionStatistics.builder() - .setBasicStatistics(new HiveBasicStatistics(0, 0, 0, 0)) - .setColumnStatistics(ImmutableMap.of(COLUMN, HiveColumnStatistics.builder().setTotalSizeInBytes(-1).build())) - .build(), - invalidColumnStatistics("totalSizeInBytes must be greater than or equal to zero: -1")); + //assertInvalidStatistics( + // PartitionStatistics.builder() + // .setBasicStatistics(new HiveBasicStatistics(0, 0, 0, 0)) + // .setColumnStatistics(ImmutableMap.of(COLUMN, HiveColumnStatistics.builder().setTotalSizeInBytes(-1).build())) + // .build(), + // invalidColumnStatistics("totalSizeInBytes must be greater than or equal to zero: -1")); assertInvalidStatistics( PartitionStatistics.builder() .setBasicStatistics(new HiveBasicStatistics(0, 0, 0, 0)) @@ -164,24 +164,24 @@ public void testValidatePartitionStatistics() .setColumnStatistics(ImmutableMap.of(COLUMN, HiveColumnStatistics.builder().setNullsCount(1).build())) .build(), invalidColumnStatistics("nullsCount must be less than or equal to rowCount. nullsCount: 1. rowCount: 0.")); - assertInvalidStatistics( - PartitionStatistics.builder() - .setBasicStatistics(new HiveBasicStatistics(0, 0, 0, 0)) - .setColumnStatistics(ImmutableMap.of(COLUMN, HiveColumnStatistics.builder().setDistinctValuesCount(-1).build())) - .build(), - invalidColumnStatistics("distinctValuesCount must be greater than or equal to zero: -1")); - assertInvalidStatistics( - PartitionStatistics.builder() - .setBasicStatistics(new HiveBasicStatistics(0, 0, 0, 0)) - .setColumnStatistics(ImmutableMap.of(COLUMN, HiveColumnStatistics.builder().setDistinctValuesCount(1).build())) - .build(), - invalidColumnStatistics("distinctValuesCount must be less than or equal to rowCount. distinctValuesCount: 1. rowCount: 0.")); - assertInvalidStatistics( - PartitionStatistics.builder() - .setBasicStatistics(new HiveBasicStatistics(0, 1, 0, 0)) - .setColumnStatistics(ImmutableMap.of(COLUMN, HiveColumnStatistics.builder().setDistinctValuesCount(1).setNullsCount(1).build())) - .build(), - invalidColumnStatistics("distinctValuesCount must be less than or equal to nonNullsCount. distinctValuesCount: 1. nonNullsCount: 0.")); + //assertInvalidStatistics( + // PartitionStatistics.builder() + // .setBasicStatistics(new HiveBasicStatistics(0, 0, 0, 0)) + // .setColumnStatistics(ImmutableMap.of(COLUMN, HiveColumnStatistics.builder().setDistinctValuesCount(-1).build())) + // .build(), + // invalidColumnStatistics("distinctValuesCount must be greater than or equal to zero: -1")); + //assertInvalidStatistics( + // PartitionStatistics.builder() + // .setBasicStatistics(new HiveBasicStatistics(0, 0, 0, 0)) + // .setColumnStatistics(ImmutableMap.of(COLUMN, HiveColumnStatistics.builder().setDistinctValuesCount(1).build())) + // .build(), + // invalidColumnStatistics("distinctValuesCount must be less than or equal to rowCount. distinctValuesCount: 1. rowCount: 0.")); + //assertInvalidStatistics( + // PartitionStatistics.builder() + // .setBasicStatistics(new HiveBasicStatistics(0, 1, 0, 0)) + // .setColumnStatistics(ImmutableMap.of(COLUMN, HiveColumnStatistics.builder().setDistinctValuesCount(1).setNullsCount(1).build())) + // .build(), + // invalidColumnStatistics("distinctValuesCount must be less than or equal to nonNullsCount. distinctValuesCount: 1. nonNullsCount: 0.")); assertInvalidStatistics( PartitionStatistics.builder() .setBasicStatistics(new HiveBasicStatistics(0, 0, 0, 0)) @@ -516,23 +516,23 @@ public void testCreateDataColumnStatistics() @Test public void testCalculateDistinctValuesCount() { - assertEquals(calculateDistinctValuesCount(ImmutableList.of()), Estimate.unknown()); - assertEquals(calculateDistinctValuesCount(ImmutableList.of(HiveColumnStatistics.empty())), Estimate.unknown()); - assertEquals(calculateDistinctValuesCount(ImmutableList.of(HiveColumnStatistics.empty(), HiveColumnStatistics.empty())), Estimate.unknown()); - assertEquals(calculateDistinctValuesCount(ImmutableList.of(distinctValuesCount(1))), Estimate.of(1)); - assertEquals(calculateDistinctValuesCount(ImmutableList.of(distinctValuesCount(1), distinctValuesCount(2))), Estimate.of(2)); - assertEquals(calculateDistinctValuesCount(ImmutableList.of(distinctValuesCount(1), HiveColumnStatistics.empty())), Estimate.of(1)); - assertEquals(calculateDistinctValuesCount(ImmutableList.of(createBooleanColumnStatistics(OptionalLong.empty(), OptionalLong.empty(), OptionalLong.empty()))), Estimate.unknown()); - assertEquals(calculateDistinctValuesCount(ImmutableList.of(createBooleanColumnStatistics(OptionalLong.of(1), OptionalLong.of(0), OptionalLong.empty()))), Estimate.of(1)); - assertEquals(calculateDistinctValuesCount(ImmutableList.of(createBooleanColumnStatistics(OptionalLong.of(10), OptionalLong.empty(), OptionalLong.empty()))), Estimate.unknown()); - assertEquals(calculateDistinctValuesCount(ImmutableList.of(createBooleanColumnStatistics(OptionalLong.of(10), OptionalLong.of(10), OptionalLong.empty()))), Estimate.of(2)); - assertEquals(calculateDistinctValuesCount(ImmutableList.of(createBooleanColumnStatistics(OptionalLong.empty(), OptionalLong.of(10), OptionalLong.empty()))), Estimate.unknown()); - assertEquals(calculateDistinctValuesCount(ImmutableList.of(createBooleanColumnStatistics(OptionalLong.of(0), OptionalLong.of(10), OptionalLong.empty()))), Estimate.of(1)); - assertEquals(calculateDistinctValuesCount(ImmutableList.of(createBooleanColumnStatistics(OptionalLong.of(0), OptionalLong.of(0), OptionalLong.empty()))), Estimate.of(0)); - assertEquals( - calculateDistinctValuesCount(ImmutableList.of( - createBooleanColumnStatistics(OptionalLong.of(0), OptionalLong.of(10), OptionalLong.empty()), - createBooleanColumnStatistics(OptionalLong.of(1), OptionalLong.of(10), OptionalLong.empty()))), + assertEquals(calculateDistinctValuesCount(COLUMN, ImmutableList.of()), Estimate.unknown()); + assertEquals(calculateDistinctValuesCount(COLUMN, ImmutableList.of(PartitionStatistics.empty())), Estimate.unknown()); + assertEquals(calculateDistinctValuesCount(COLUMN, ImmutableList.of(PartitionStatistics.empty(), PartitionStatistics.empty())), Estimate.unknown()); + assertEquals(calculateDistinctValuesCount(COLUMN, ImmutableList.of(distinctValuesCount(1))), Estimate.of(1)); + assertEquals(calculateDistinctValuesCount(COLUMN, ImmutableList.of(distinctValuesCount(1), distinctValuesCount(2))), Estimate.of(2)); + assertEquals(calculateDistinctValuesCount(COLUMN, ImmutableList.of(distinctValuesCount(1), PartitionStatistics.empty())), Estimate.of(1)); + assertEquals(calculateDistinctValuesCount(COLUMN, ImmutableList.of(booleanDistinctValuesCount(OptionalLong.empty(), OptionalLong.empty(), OptionalLong.empty()))), Estimate.unknown()); + assertEquals(calculateDistinctValuesCount(COLUMN, ImmutableList.of(booleanDistinctValuesCount(OptionalLong.of(1), OptionalLong.of(0), OptionalLong.empty()))), Estimate.of(1)); + assertEquals(calculateDistinctValuesCount(COLUMN, ImmutableList.of(booleanDistinctValuesCount(OptionalLong.of(10), OptionalLong.empty(), OptionalLong.empty()))), Estimate.unknown()); + assertEquals(calculateDistinctValuesCount(COLUMN, ImmutableList.of(booleanDistinctValuesCount(OptionalLong.of(10), OptionalLong.of(10), OptionalLong.empty()))), Estimate.of(2)); + assertEquals(calculateDistinctValuesCount(COLUMN, ImmutableList.of(booleanDistinctValuesCount(OptionalLong.empty(), OptionalLong.of(10), OptionalLong.empty()))), Estimate.unknown()); + assertEquals(calculateDistinctValuesCount(COLUMN, ImmutableList.of(booleanDistinctValuesCount(OptionalLong.of(0), OptionalLong.of(10), OptionalLong.empty()))), Estimate.of(1)); + assertEquals(calculateDistinctValuesCount(COLUMN, ImmutableList.of(booleanDistinctValuesCount(OptionalLong.of(0), OptionalLong.of(0), OptionalLong.empty()))), Estimate.of(0)); + assertEquals( + calculateDistinctValuesCount(COLUMN, ImmutableList.of( + booleanDistinctValuesCount(OptionalLong.of(0), OptionalLong.of(10), OptionalLong.empty()), + booleanDistinctValuesCount(OptionalLong.of(1), OptionalLong.of(10), OptionalLong.empty()))), Estimate.of(2)); } @@ -557,7 +557,7 @@ public void testCalculateDataSize() assertEquals(calculateDataSize(COLUMN, ImmutableList.of(rowsCount(1000)), 1000), Estimate.unknown()); assertEquals(calculateDataSize(COLUMN, ImmutableList.of(dataSize(1000)), 1000), Estimate.unknown()); assertEquals(calculateDataSize(COLUMN, ImmutableList.of(dataSize(1000), rowsCount(1000)), 1000), Estimate.unknown()); - assertEquals(calculateDataSize(COLUMN, ImmutableList.of(rowsCountAndDataSize(500, 1000)), 2000), Estimate.of(4000)); + assertEquals(calculateDataSize(COLUMN, ImmutableList.of(rowsCountAndDataSize(500, 2)), 2000), Estimate.of(4000)); assertEquals(calculateDataSize(COLUMN, ImmutableList.of(rowsCountAndDataSize(0, 0)), 2000), Estimate.unknown()); assertEquals(calculateDataSize(COLUMN, ImmutableList.of(rowsCountAndDataSize(0, 0)), 0), Estimate.zero()); assertEquals(calculateDataSize(COLUMN, ImmutableList.of(rowsCountAndDataSize(1000, 0)), 2000), Estimate.of(0)); @@ -565,8 +565,8 @@ public void testCalculateDataSize() calculateDataSize( COLUMN, ImmutableList.of( - rowsCountAndDataSize(500, 1000), - rowsCountAndDataSize(1000, 5000)), + rowsCountAndDataSize(500, 2), + rowsCountAndDataSize(1000, 5)), 5000), Estimate.of(20000)); assertEquals( @@ -574,9 +574,9 @@ public void testCalculateDataSize() COLUMN, ImmutableList.of( dataSize(1000), - rowsCountAndDataSize(500, 1000), + rowsCountAndDataSize(500, 2), rowsCount(3000), - rowsCountAndDataSize(1000, 5000)), + rowsCountAndDataSize(1000, 5)), 5000), Estimate.of(20000)); } @@ -628,7 +628,7 @@ public void testGetTableStatistics() String partitionName = "p1=string1/p2=1234"; PartitionStatistics statistics = PartitionStatistics.builder() .setBasicStatistics(new HiveBasicStatistics(OptionalLong.empty(), OptionalLong.of(1000), OptionalLong.empty(), OptionalLong.empty())) - .setColumnStatistics(ImmutableMap.of(COLUMN, createIntegerColumnStatistics(OptionalLong.of(-100), OptionalLong.of(100), OptionalLong.of(500), OptionalLong.of(300)))) + .setColumnStatistics(ImmutableMap.of(COLUMN, createIntegerColumnStatistics(OptionalLong.of(-100), OptionalLong.of(100), OptionalLong.of(500), OptionalLong.of(301)))) .build(); MetastoreHiveStatisticsProvider statisticsProvider = new MetastoreHiveStatisticsProvider((session, table, hivePartitions, columns) -> ImmutableMap.of(partitionName, statistics)); HiveColumnHandle columnHandle = createBaseColumn(COLUMN, 2, HIVE_LONG, BIGINT, REGULAR, Optional.empty()); @@ -677,7 +677,7 @@ public void testGetTableStatisticsUnpartitioned() { PartitionStatistics statistics = PartitionStatistics.builder() .setBasicStatistics(new HiveBasicStatistics(OptionalLong.empty(), OptionalLong.of(1000), OptionalLong.empty(), OptionalLong.empty())) - .setColumnStatistics(ImmutableMap.of(COLUMN, createIntegerColumnStatistics(OptionalLong.of(-100), OptionalLong.of(100), OptionalLong.of(500), OptionalLong.of(300)))) + .setColumnStatistics(ImmutableMap.of(COLUMN, createIntegerColumnStatistics(OptionalLong.of(-100), OptionalLong.of(100), OptionalLong.of(500), OptionalLong.of(301)))) .build(); MetastoreHiveStatisticsProvider statisticsProvider = new MetastoreHiveStatisticsProvider((session, table, hivePartitions, columns) -> ImmutableMap.of(UNPARTITIONED_ID, statistics)); @@ -862,7 +862,7 @@ private static PartitionStatistics nullsCount(long nullsCount) private static PartitionStatistics dataSize(long dataSize) { - return new PartitionStatistics(HiveBasicStatistics.createEmptyStatistics(), ImmutableMap.of(COLUMN, HiveColumnStatistics.builder().setTotalSizeInBytes(dataSize).build())); + return new PartitionStatistics(HiveBasicStatistics.createEmptyStatistics(), ImmutableMap.of(COLUMN, HiveColumnStatistics.builder().setAverageColumnLength(dataSize).build())); } private static PartitionStatistics rowsCountAndNullsCount(long rowsCount, long nullsCount) @@ -872,18 +872,24 @@ private static PartitionStatistics rowsCountAndNullsCount(long rowsCount, long n ImmutableMap.of(COLUMN, HiveColumnStatistics.builder().setNullsCount(nullsCount).build())); } - private static PartitionStatistics rowsCountAndDataSize(long rowsCount, long dataSize) + private static PartitionStatistics rowsCountAndDataSize(long rowsCount, long averageColumnLength) { return new PartitionStatistics( new HiveBasicStatistics(0, rowsCount, 0, 0), - ImmutableMap.of(COLUMN, HiveColumnStatistics.builder().setTotalSizeInBytes(dataSize).build())); + ImmutableMap.of(COLUMN, HiveColumnStatistics.builder().setAverageColumnLength(averageColumnLength).build())); } - private static HiveColumnStatistics distinctValuesCount(long count) + private static PartitionStatistics distinctValuesCount(long count) { - return HiveColumnStatistics.builder() - .setDistinctValuesCount(count) - .build(); + return new PartitionStatistics(HiveBasicStatistics.createEmptyStatistics(), ImmutableMap.of(COLUMN, HiveColumnStatistics.builder() + .setDistinctValuesWithNullCount(count) + .build())); + } + + private static PartitionStatistics booleanDistinctValuesCount(OptionalLong trueCount, OptionalLong falseCount, OptionalLong nullsCount) + { + return new PartitionStatistics(HiveBasicStatistics.createEmptyStatistics(), + ImmutableMap.of(COLUMN, createBooleanColumnStatistics(trueCount, falseCount, nullsCount))); } private static HiveColumnStatistics integerRange(long min, long max) diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/util/TestHiveUtil.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/util/TestHiveUtil.java index 5ae8d4470c03..11c55eaf77b5 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/util/TestHiveUtil.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/util/TestHiveUtil.java @@ -13,17 +13,9 @@ */ package io.trino.plugin.hive.util; -import org.apache.hadoop.conf.Configuration; import org.apache.hadoop.hive.common.FileUtils; import org.apache.hadoop.hive.metastore.Warehouse; import org.apache.hadoop.hive.metastore.api.MetaException; -import org.apache.hadoop.hive.ql.io.SymlinkTextInputFormat; -import org.apache.hadoop.hive.ql.io.avro.AvroContainerInputFormat; -import org.apache.hadoop.hive.ql.io.parquet.MapredParquetInputFormat; -import org.apache.hadoop.hive.serde2.thrift.ThriftDeserializer; -import org.apache.hadoop.hive.serde2.thrift.test.IntString; -import org.apache.hadoop.mapred.TextInputFormat; -import org.apache.thrift.protocol.TBinaryProtocol; import org.joda.time.DateTime; import org.joda.time.DateTimeZone; import org.joda.time.format.DateTimeFormat; @@ -32,24 +24,12 @@ import java.util.AbstractList; import java.util.ArrayList; import java.util.List; -import java.util.Properties; - -import static io.airlift.testing.Assertions.assertInstanceOf; -import static io.trino.hadoop.ConfigurationInstantiator.newEmptyConfiguration; -import static io.trino.plugin.hive.HiveStorageFormat.AVRO; -import static io.trino.plugin.hive.HiveStorageFormat.PARQUET; -import static io.trino.plugin.hive.HiveStorageFormat.SEQUENCEFILE; -import static io.trino.plugin.hive.util.HiveReaderUtil.getDeserializer; -import static io.trino.plugin.hive.util.HiveReaderUtil.getInputFormat; + import static io.trino.plugin.hive.util.HiveUtil.escapeSchemaName; import static io.trino.plugin.hive.util.HiveUtil.escapeTableName; import static io.trino.plugin.hive.util.HiveUtil.parseHiveTimestamp; import static io.trino.plugin.hive.util.HiveUtil.toPartitionValues; import static io.trino.type.DateTimes.MICROSECONDS_PER_MILLISECOND; -import static org.apache.hadoop.hive.metastore.api.hive_metastoreConstants.FILE_INPUT_FORMAT; -import static org.apache.hadoop.hive.serde.serdeConstants.SERIALIZATION_CLASS; -import static org.apache.hadoop.hive.serde.serdeConstants.SERIALIZATION_FORMAT; -import static org.apache.hadoop.hive.serde.serdeConstants.SERIALIZATION_LIB; import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatThrownBy; import static org.testng.Assert.assertEquals; @@ -67,17 +47,6 @@ public void testParseHiveTimestamp() assertEquals(parse(time, "yyyy-MM-dd HH:mm:ss.SSSSSSSSS"), unixTime(time, 7)); } - @Test - public void testGetThriftDeserializer() - { - Properties schema = new Properties(); - schema.setProperty(SERIALIZATION_LIB, ThriftDeserializer.class.getName()); - schema.setProperty(SERIALIZATION_CLASS, IntString.class.getName()); - schema.setProperty(SERIALIZATION_FORMAT, TBinaryProtocol.class.getName()); - - assertInstanceOf(getDeserializer(newEmptyConfiguration(), schema), ThriftDeserializer.class); - } - @Test public void testToPartitionValues() throws MetaException @@ -90,37 +59,6 @@ public void testToPartitionValues() assertToPartitionValues("pk=__HIVE_DEFAULT_PARTITION__"); } - @Test - public void testGetInputFormat() - { - Configuration configuration = newEmptyConfiguration(); - - // LazySimpleSerDe is used by TEXTFILE and SEQUENCEFILE. getInputFormat should default to TEXTFILE - // per Hive spec. - Properties sequenceFileSchema = new Properties(); - sequenceFileSchema.setProperty(FILE_INPUT_FORMAT, SymlinkTextInputFormat.class.getName()); - sequenceFileSchema.setProperty(SERIALIZATION_LIB, SEQUENCEFILE.getSerde()); - assertInstanceOf(getInputFormat(configuration, sequenceFileSchema), TextInputFormat.class); - - Properties avroSymlinkSchema = new Properties(); - avroSymlinkSchema.setProperty(FILE_INPUT_FORMAT, SymlinkTextInputFormat.class.getName()); - avroSymlinkSchema.setProperty(SERIALIZATION_LIB, AVRO.getSerde()); - assertInstanceOf(getInputFormat(configuration, avroSymlinkSchema), AvroContainerInputFormat.class); - - Properties parquetSymlinkSchema = new Properties(); - parquetSymlinkSchema.setProperty(FILE_INPUT_FORMAT, SymlinkTextInputFormat.class.getName()); - parquetSymlinkSchema.setProperty(SERIALIZATION_LIB, PARQUET.getSerde()); - assertInstanceOf(getInputFormat(configuration, parquetSymlinkSchema), MapredParquetInputFormat.class); - - Properties parquetSchema = new Properties(); - parquetSchema.setProperty(FILE_INPUT_FORMAT, PARQUET.getInputFormat()); - assertInstanceOf(getInputFormat(configuration, parquetSchema), MapredParquetInputFormat.class); - - Properties legacyParquetSchema = new Properties(); - legacyParquetSchema.setProperty(FILE_INPUT_FORMAT, "parquet.hive.MapredParquetInputFormat"); - assertInstanceOf(getInputFormat(configuration, legacyParquetSchema), MapredParquetInputFormat.class); - } - @Test public void testUnescapePathName() { diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/util/TestHiveWriteUtils.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/util/TestHiveWriteUtils.java index 2bb269ed70c9..57c05f772c0d 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/util/TestHiveWriteUtils.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/util/TestHiveWriteUtils.java @@ -13,40 +13,68 @@ */ package io.trino.plugin.hive.util; -import io.trino.hdfs.HdfsContext; -import org.apache.hadoop.fs.Path; -import org.testng.annotations.Test; - -import static io.trino.plugin.hive.HiveTestUtils.HDFS_ENVIRONMENT; -import static io.trino.plugin.hive.util.HiveWriteUtils.isS3FileSystem; -import static io.trino.plugin.hive.util.HiveWriteUtils.isViewFileSystem; -import static io.trino.testing.TestingConnectorSession.SESSION; -import static io.trino.testing.TestingNames.randomNameSuffix; -import static org.testng.Assert.assertFalse; -import static org.testng.Assert.assertTrue; +import io.trino.spi.Page; +import io.trino.spi.PageBuilder; +import io.trino.spi.block.BlockBuilder; +import io.trino.spi.type.DecimalType; +import io.trino.spi.type.SqlDecimal; +import io.trino.spi.type.Type; +import org.apache.hadoop.hive.common.type.HiveDecimal; +import org.junit.jupiter.api.Test; + +import java.util.List; + +import static io.trino.plugin.hive.util.HiveWriteUtils.createPartitionValues; +import static io.trino.spi.type.DecimalType.createDecimalType; +import static io.trino.spi.type.Decimals.writeBigDecimal; +import static io.trino.spi.type.Decimals.writeShortDecimal; +import static io.trino.spi.type.SqlDecimal.decimal; +import static org.assertj.core.api.Assertions.assertThat; public class TestHiveWriteUtils { - private static final HdfsContext CONTEXT = new HdfsContext(SESSION); - private static final String RANDOM_SUFFIX = randomNameSuffix(); - @Test - public void testIsS3FileSystem() + public void testCreatePartitionValuesDecimal() { - assertTrue(isS3FileSystem(CONTEXT, HDFS_ENVIRONMENT, new Path("s3://test-bucket-%s/test-folder".formatted(RANDOM_SUFFIX)))); - assertFalse(isS3FileSystem(CONTEXT, HDFS_ENVIRONMENT, new Path("/test-dir-%s/test-folder".formatted(RANDOM_SUFFIX)))); + assertCreatePartitionValuesDecimal(10, 0, "12345", "12345"); + assertCreatePartitionValuesDecimal(10, 2, "123.45", "123.45"); + assertCreatePartitionValuesDecimal(10, 2, "12345.00", "12345"); + assertCreatePartitionValuesDecimal(5, 0, "12345", "12345"); + assertCreatePartitionValuesDecimal(38, 2, "12345.00", "12345"); + assertCreatePartitionValuesDecimal(38, 20, "12345.00000000000000000000", "12345"); + assertCreatePartitionValuesDecimal(38, 20, "12345.67898000000000000000", "12345.67898"); } - @Test - public void testIsViewFileSystem() + private static void assertCreatePartitionValuesDecimal(int precision, int scale, String decimalValue, String expectedValue) { - Path viewfsPath = new Path("viewfs://ns-default-%s/test-folder".formatted(RANDOM_SUFFIX)); - Path nonViewfsPath = new Path("hdfs://localhost/test-dir/test-folder"); + DecimalType decimalType = createDecimalType(precision, scale); + List types = List.of(decimalType); + SqlDecimal decimal = decimal(decimalValue, decimalType); + + // verify the test values are as expected + assertThat(decimal.toString()).isEqualTo(decimalValue); + assertThat(decimal.toBigDecimal().toString()).isEqualTo(decimalValue); - // ViewFS check requires the mount point config - HDFS_ENVIRONMENT.getConfiguration(CONTEXT, viewfsPath).set("fs.viewfs.mounttable.ns-default-%s.link./test-folder".formatted(RANDOM_SUFFIX), "hdfs://localhost/app"); + PageBuilder pageBuilder = new PageBuilder(types); + pageBuilder.declarePosition(); + writeDecimal(decimalType, decimal, pageBuilder.getBlockBuilder(0)); + Page page = pageBuilder.build(); - assertTrue(isViewFileSystem(CONTEXT, HDFS_ENVIRONMENT, viewfsPath)); - assertFalse(isViewFileSystem(CONTEXT, HDFS_ENVIRONMENT, nonViewfsPath)); + // verify the expected value against HiveDecimal + assertThat(HiveDecimal.create(decimal.toBigDecimal()).toString()) + .isEqualTo(expectedValue); + + assertThat(createPartitionValues(types, page, 0)) + .isEqualTo(List.of(expectedValue)); + } + + private static void writeDecimal(DecimalType decimalType, SqlDecimal decimal, BlockBuilder blockBuilder) + { + if (decimalType.isShort()) { + writeShortDecimal(blockBuilder, decimal.toBigDecimal().unscaledValue().longValue()); + } + else { + writeBigDecimal(decimalType, blockBuilder, decimal.toBigDecimal()); + } } } diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/util/TestStatistics.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/util/TestStatistics.java index 02af9c7e2511..f63882421af7 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/util/TestStatistics.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/util/TestStatistics.java @@ -44,9 +44,6 @@ import static io.trino.plugin.hive.HiveColumnStatisticType.MIN_VALUE; import static io.trino.plugin.hive.HiveColumnStatisticType.NUMBER_OF_DISTINCT_VALUES; import static io.trino.plugin.hive.HiveColumnStatisticType.NUMBER_OF_NON_NULL_VALUES; -import static io.trino.plugin.hive.metastore.HiveColumnStatistics.createBinaryColumnStatistics; -import static io.trino.plugin.hive.metastore.HiveColumnStatistics.createBooleanColumnStatistics; -import static io.trino.plugin.hive.metastore.HiveColumnStatistics.createIntegerColumnStatistics; import static io.trino.plugin.hive.util.Statistics.ReduceOperator.ADD; import static io.trino.plugin.hive.util.Statistics.ReduceOperator.SUBTRACT; import static io.trino.plugin.hive.util.Statistics.createHiveColumnStatistics; @@ -250,36 +247,23 @@ public void testMergeStringColumnStatistics() HiveColumnStatistics.builder().setMaxValueSizeInBytes(OptionalLong.of(2)).build(), HiveColumnStatistics.builder().setMaxValueSizeInBytes(OptionalLong.of(3)).build(), HiveColumnStatistics.builder().setMaxValueSizeInBytes(OptionalLong.of(3)).build()); - - assertMergeHiveColumnStatistics( - HiveColumnStatistics.builder().setTotalSizeInBytes(OptionalLong.empty()).build(), - HiveColumnStatistics.builder().setTotalSizeInBytes(OptionalLong.empty()).build(), - HiveColumnStatistics.builder().setTotalSizeInBytes(OptionalLong.empty()).build()); - assertMergeHiveColumnStatistics( - HiveColumnStatistics.builder().setTotalSizeInBytes(OptionalLong.of(1)).build(), - HiveColumnStatistics.builder().setTotalSizeInBytes(OptionalLong.empty()).build(), - HiveColumnStatistics.builder().setTotalSizeInBytes(OptionalLong.of(1)).build()); - assertMergeHiveColumnStatistics( - HiveColumnStatistics.builder().setTotalSizeInBytes(OptionalLong.of(2)).build(), - HiveColumnStatistics.builder().setTotalSizeInBytes(OptionalLong.of(3)).build(), - HiveColumnStatistics.builder().setTotalSizeInBytes(OptionalLong.of(5)).build()); } @Test public void testMergeGenericColumnStatistics() { assertMergeHiveColumnStatistics( - HiveColumnStatistics.builder().setDistinctValuesCount(OptionalLong.empty()).build(), - HiveColumnStatistics.builder().setDistinctValuesCount(OptionalLong.empty()).build(), - HiveColumnStatistics.builder().setDistinctValuesCount(OptionalLong.empty()).build()); + HiveColumnStatistics.builder().setDistinctValuesWithNullCount(OptionalLong.empty()).build(), + HiveColumnStatistics.builder().setDistinctValuesWithNullCount(OptionalLong.empty()).build(), + HiveColumnStatistics.builder().setDistinctValuesWithNullCount(OptionalLong.empty()).build()); assertMergeHiveColumnStatistics( - HiveColumnStatistics.builder().setDistinctValuesCount(OptionalLong.of(1)).build(), - HiveColumnStatistics.builder().setDistinctValuesCount(OptionalLong.empty()).build(), - HiveColumnStatistics.builder().setDistinctValuesCount(OptionalLong.empty()).build()); + HiveColumnStatistics.builder().setDistinctValuesWithNullCount(OptionalLong.of(1)).build(), + HiveColumnStatistics.builder().setDistinctValuesWithNullCount(OptionalLong.empty()).build(), + HiveColumnStatistics.builder().setDistinctValuesWithNullCount(OptionalLong.empty()).build()); assertMergeHiveColumnStatistics( - HiveColumnStatistics.builder().setDistinctValuesCount(OptionalLong.of(1)).build(), - HiveColumnStatistics.builder().setDistinctValuesCount(OptionalLong.of(2)).build(), - HiveColumnStatistics.builder().setDistinctValuesCount(OptionalLong.of(2)).build()); + HiveColumnStatistics.builder().setDistinctValuesWithNullCount(OptionalLong.of(1)).build(), + HiveColumnStatistics.builder().setDistinctValuesWithNullCount(OptionalLong.of(2)).build(), + HiveColumnStatistics.builder().setDistinctValuesWithNullCount(OptionalLong.of(2)).build()); assertMergeHiveColumnStatistics( HiveColumnStatistics.builder().setNullsCount(OptionalLong.empty()).build(), @@ -295,26 +279,6 @@ public void testMergeGenericColumnStatistics() HiveColumnStatistics.builder().setNullsCount(OptionalLong.of(3)).build()); } - @Test - public void testMergeHiveColumnStatisticsMap() - { - Map first = ImmutableMap.of( - "column1", createIntegerColumnStatistics(OptionalLong.of(1), OptionalLong.of(2), OptionalLong.of(3), OptionalLong.of(4)), - "column2", HiveColumnStatistics.createDoubleColumnStatistics(OptionalDouble.of(2), OptionalDouble.of(3), OptionalLong.of(4), OptionalLong.of(5)), - "column3", createBinaryColumnStatistics(OptionalLong.of(5), OptionalLong.of(5), OptionalLong.of(10)), - "column4", createBooleanColumnStatistics(OptionalLong.of(1), OptionalLong.of(2), OptionalLong.of(3))); - Map second = ImmutableMap.of( - "column5", createIntegerColumnStatistics(OptionalLong.of(1), OptionalLong.of(2), OptionalLong.of(3), OptionalLong.of(4)), - "column2", HiveColumnStatistics.createDoubleColumnStatistics(OptionalDouble.of(1), OptionalDouble.of(4), OptionalLong.of(4), OptionalLong.of(6)), - "column3", createBinaryColumnStatistics(OptionalLong.of(6), OptionalLong.of(5), OptionalLong.of(10)), - "column6", createBooleanColumnStatistics(OptionalLong.of(1), OptionalLong.of(2), OptionalLong.of(3))); - Map expected = ImmutableMap.of( - "column2", HiveColumnStatistics.createDoubleColumnStatistics(OptionalDouble.of(1), OptionalDouble.of(4), OptionalLong.of(8), OptionalLong.of(6)), - "column3", createBinaryColumnStatistics(OptionalLong.of(6), OptionalLong.of(10), OptionalLong.of(20))); - assertThat(merge(first, second)).isEqualTo(expected); - assertThat(merge(ImmutableMap.of(), ImmutableMap.of())).isEqualTo(ImmutableMap.of()); - } - @Test public void testFromComputedStatistics() { @@ -344,7 +308,7 @@ public void testFromComputedStatistics() HiveColumnStatistics.builder() .setIntegerStatistics(new IntegerStatistics(OptionalLong.of(1), OptionalLong.of(5))) .setNullsCount(0) - .setDistinctValuesCount(5) + .setDistinctValuesWithNullCount(5) .build()); assertThat(columnStatistics.get("b_column")).isEqualTo( HiveColumnStatistics.builder() diff --git a/plugin/trino-hive/src/test/resources/parquet_page_skipping/column_name_with_dot/data.parquet b/plugin/trino-hive/src/test/resources/parquet_page_skipping/column_name_with_dot/data.parquet new file mode 100644 index 000000000000..a086d155dffa Binary files /dev/null and b/plugin/trino-hive/src/test/resources/parquet_page_skipping/column_name_with_dot/data.parquet differ diff --git a/plugin/trino-hive/src/test/resources/parquet_page_skipping/lineitem_sorted_by_suppkey/data.parquet b/plugin/trino-hive/src/test/resources/parquet_page_skipping/lineitem_sorted_by_suppkey/data.parquet new file mode 100644 index 000000000000..076faaa38b5e Binary files /dev/null and b/plugin/trino-hive/src/test/resources/parquet_page_skipping/lineitem_sorted_by_suppkey/data.parquet differ diff --git a/plugin/trino-hive/src/test/resources/parquet_page_skipping/orders_sorted_by_totalprice/data.parquet b/plugin/trino-hive/src/test/resources/parquet_page_skipping/orders_sorted_by_totalprice/data.parquet new file mode 100644 index 000000000000..82fd24bf7c66 Binary files /dev/null and b/plugin/trino-hive/src/test/resources/parquet_page_skipping/orders_sorted_by_totalprice/data.parquet differ diff --git a/plugin/trino-hive/src/test/resources/parquet_page_skipping/random/data.parquet b/plugin/trino-hive/src/test/resources/parquet_page_skipping/random/data.parquet new file mode 100644 index 000000000000..dc1be1bbf044 Binary files /dev/null and b/plugin/trino-hive/src/test/resources/parquet_page_skipping/random/data.parquet differ diff --git a/plugin/trino-hive/src/test/resources/with_short_zone_id/data/data.orc b/plugin/trino-hive/src/test/resources/with_short_zone_id/data/data.orc new file mode 100755 index 000000000000..d59fd2486871 Binary files /dev/null and b/plugin/trino-hive/src/test/resources/with_short_zone_id/data/data.orc differ diff --git a/plugin/trino-hudi/pom.xml b/plugin/trino-hudi/pom.xml index 10bfc51f4372..4d2abdc5efce 100644 --- a/plugin/trino-hudi/pom.xml +++ b/plugin/trino-hudi/pom.xml @@ -142,11 +142,6 @@ parquet-column - - org.apache.parquet - parquet-hadoop - - org.weakref jmxutils diff --git a/plugin/trino-hudi/src/main/java/io/trino/plugin/hudi/HudiMetadata.java b/plugin/trino-hudi/src/main/java/io/trino/plugin/hudi/HudiMetadata.java index 709e7439862c..51ff19069ce7 100644 --- a/plugin/trino-hudi/src/main/java/io/trino/plugin/hudi/HudiMetadata.java +++ b/plugin/trino-hudi/src/main/java/io/trino/plugin/hudi/HudiMetadata.java @@ -22,6 +22,7 @@ import io.trino.plugin.hive.metastore.Column; import io.trino.plugin.hive.metastore.HiveMetastore; import io.trino.plugin.hive.metastore.Table; +import io.trino.plugin.hive.metastore.TableInfo; import io.trino.spi.TrinoException; import io.trino.spi.connector.ColumnHandle; import io.trino.spi.connector.ColumnMetadata; @@ -204,8 +205,8 @@ public List listTables(ConnectorSession session, Optional tableNames = ImmutableList.builder(); for (String schemaName : listSchemas(session, optionalSchemaName)) { - for (String tableName : metastore.getAllTables(schemaName)) { - tableNames.add(new SchemaTableName(schemaName, tableName)); + for (TableInfo tableInfo : metastore.getTables(schemaName)) { + tableNames.add(tableInfo.tableName()); } } return tableNames.build(); diff --git a/plugin/trino-hudi/src/main/java/io/trino/plugin/hudi/HudiPageSourceProvider.java b/plugin/trino-hudi/src/main/java/io/trino/plugin/hudi/HudiPageSourceProvider.java index 9e888c6a6a93..359535c675d1 100644 --- a/plugin/trino-hudi/src/main/java/io/trino/plugin/hudi/HudiPageSourceProvider.java +++ b/plugin/trino-hudi/src/main/java/io/trino/plugin/hudi/HudiPageSourceProvider.java @@ -23,9 +23,14 @@ import io.trino.parquet.ParquetDataSource; import io.trino.parquet.ParquetDataSourceId; import io.trino.parquet.ParquetReaderOptions; +import io.trino.parquet.metadata.BlockMetadata; +import io.trino.parquet.metadata.FileMetadata; +import io.trino.parquet.metadata.ParquetMetadata; +import io.trino.parquet.metadata.PrunedBlockMetadata; import io.trino.parquet.predicate.TupleDomainParquetPredicate; import io.trino.parquet.reader.MetadataReader; import io.trino.parquet.reader.ParquetReader; +import io.trino.parquet.reader.RowGroupInfo; import io.trino.plugin.hive.FileFormatDataSourceStats; import io.trino.plugin.hive.HiveColumnHandle; import io.trino.plugin.hive.HivePartitionKey; @@ -47,9 +52,6 @@ import io.trino.spi.type.Decimals; import io.trino.spi.type.TypeSignature; import org.apache.parquet.column.ColumnDescriptor; -import org.apache.parquet.hadoop.metadata.BlockMetaData; -import org.apache.parquet.hadoop.metadata.FileMetaData; -import org.apache.parquet.hadoop.metadata.ParquetMetadata; import org.apache.parquet.internal.filter2.columnindex.ColumnIndexStore; import org.apache.parquet.io.MessageColumnIO; import org.apache.parquet.schema.MessageType; @@ -72,12 +74,13 @@ import static io.trino.memory.context.AggregatedMemoryContext.newSimpleAggregatedMemoryContext; import static io.trino.parquet.ParquetTypeUtils.getColumnIO; import static io.trino.parquet.ParquetTypeUtils.getDescriptors; +import static io.trino.parquet.metadata.PrunedBlockMetadata.createPrunedColumnsMetadata; import static io.trino.parquet.predicate.PredicateUtils.buildPredicate; import static io.trino.parquet.predicate.PredicateUtils.predicateMatches; +import static io.trino.parquet.reader.TrinoColumnIndexStore.getColumnIndexStore; import static io.trino.plugin.hive.HivePageSourceProvider.projectBaseColumns; import static io.trino.plugin.hive.parquet.ParquetPageSourceFactory.ParquetReaderProvider; import static io.trino.plugin.hive.parquet.ParquetPageSourceFactory.createParquetPageSource; -import static io.trino.plugin.hive.parquet.ParquetPageSourceFactory.getColumnIndexStore; import static io.trino.plugin.hive.parquet.ParquetPageSourceFactory.getParquetMessageType; import static io.trino.plugin.hive.parquet.ParquetPageSourceFactory.getParquetTupleDomain; import static io.trino.plugin.hive.util.HiveUtil.makePartName; @@ -86,8 +89,6 @@ import static io.trino.plugin.hudi.HudiErrorCode.HUDI_CURSOR_ERROR; import static io.trino.plugin.hudi.HudiErrorCode.HUDI_INVALID_PARTITION_VALUE; import static io.trino.plugin.hudi.HudiErrorCode.HUDI_UNSUPPORTED_FILE_FORMAT; -import static io.trino.plugin.hudi.HudiSessionProperties.isParquetOptimizedNestedReaderEnabled; -import static io.trino.plugin.hudi.HudiSessionProperties.isParquetOptimizedReaderEnabled; import static io.trino.plugin.hudi.HudiSessionProperties.shouldUseParquetColumnNames; import static io.trino.plugin.hudi.HudiUtil.getHudiFileFormat; import static io.trino.spi.predicate.Utils.nativeValueToBlock; @@ -167,8 +168,7 @@ public ConnectorPageSource createPageSource( split, inputFile, dataSourceStats, - options.withBatchColumnReaders(isParquetOptimizedReaderEnabled(session)) - .withBatchNestedColumnReaders(isParquetOptimizedNestedReaderEnabled(session)), + options, timeZone); return new HudiPageSource( @@ -198,7 +198,7 @@ private static ConnectorPageSource createPageSource( try { dataSource = new TrinoParquetDataSource(inputFile, options, dataSourceStats); ParquetMetadata parquetMetadata = MetadataReader.readFooter(dataSource, Optional.empty()); - FileMetaData fileMetaData = parquetMetadata.getFileMetaData(); + FileMetadata fileMetaData = parquetMetadata.getFileMetaData(); MessageType fileSchema = fileMetaData.getSchema(); Optional message = getParquetMessageType(columns, useColumnNames, fileSchema); @@ -214,19 +214,20 @@ private static ConnectorPageSource createPageSource( TupleDomainParquetPredicate parquetPredicate = buildPredicate(requestedSchema, parquetTupleDomain, descriptorsByPath, timeZone); long nextStart = 0; - ImmutableList.Builder blocks = ImmutableList.builder(); + ImmutableList.Builder blocks = ImmutableList.builder(); ImmutableList.Builder blockStarts = ImmutableList.builder(); ImmutableList.Builder> columnIndexes = ImmutableList.builder(); - for (BlockMetaData block : parquetMetadata.getBlocks()) { - long firstDataPage = block.getColumns().get(0).getFirstDataPageOffset(); + for (BlockMetadata block : parquetMetadata.getBlocks()) { + long firstDataPage = block.columns().get(0).getFirstDataPageOffset(); Optional columnIndex = getColumnIndexStore(dataSource, block, descriptorsByPath, parquetTupleDomain, options); + PrunedBlockMetadata prunedBlock = createPrunedColumnsMetadata(block, dataSource.getId(), descriptorsByPath); if (start <= firstDataPage && firstDataPage < start + length - && predicateMatches(parquetPredicate, block, dataSource, descriptorsByPath, parquetTupleDomain, columnIndex, Optional.empty(), timeZone, DOMAIN_COMPACTION_THRESHOLD)) { + && predicateMatches(parquetPredicate, prunedBlock, dataSource, descriptorsByPath, parquetTupleDomain, columnIndex, Optional.empty(), timeZone, DOMAIN_COMPACTION_THRESHOLD)) { blocks.add(block); blockStarts.add(nextStart); columnIndexes.add(columnIndex); } - nextStart += block.getRowCount(); + nextStart += block.rowCount(); } Optional readerProjections = projectBaseColumns(columns); @@ -237,18 +238,22 @@ && predicateMatches(parquetPredicate, block, dataSource, descriptorsByPath, parq .orElse(columns); ParquetDataSourceId dataSourceId = dataSource.getId(); ParquetDataSource finalDataSource = dataSource; - ParquetReaderProvider parquetReaderProvider = fields -> new ParquetReader( + ImmutableList.Builder rowGroupInfoBuilder = ImmutableList.builder(); + for (BlockMetadata block : parquetMetadata.getBlocks()) { + rowGroupInfoBuilder.add(new RowGroupInfo(createPrunedColumnsMetadata(block, finalDataSource.getId(), descriptorsByPath), nextStart, Optional.empty())); + nextStart += block.rowCount(); + } + ParquetReaderProvider parquetReaderProvider = (fields, appendRowNumberColumn) -> new ParquetReader( Optional.ofNullable(fileMetaData.getCreatedBy()), fields, - blocks.build(), - blockStarts.build(), + appendRowNumberColumn, + rowGroupInfoBuilder.build(), finalDataSource, timeZone, newSimpleAggregatedMemoryContext(), options, exception -> handleException(dataSourceId, exception), Optional.of(parquetPredicate), - columnIndexes.build(), Optional.empty()); return createParquetPageSource(baseColumns, fileSchema, messageColumn, useColumnNames, parquetReaderProvider); } diff --git a/plugin/trino-hudi/src/main/java/io/trino/plugin/hudi/HudiSessionProperties.java b/plugin/trino-hudi/src/main/java/io/trino/plugin/hudi/HudiSessionProperties.java index bb5400b1e17e..9887fdc1125b 100644 --- a/plugin/trino-hudi/src/main/java/io/trino/plugin/hudi/HudiSessionProperties.java +++ b/plugin/trino-hudi/src/main/java/io/trino/plugin/hudi/HudiSessionProperties.java @@ -78,16 +78,6 @@ public HudiSessionProperties(HudiConfig hudiConfig, ParquetReaderConfig parquetR "Access parquet columns using names from the file. If disabled, then columns are accessed using index.", hudiConfig.getUseParquetColumnNames(), false), - booleanProperty( - PARQUET_OPTIMIZED_READER_ENABLED, - "Use optimized Parquet reader", - parquetReaderConfig.isOptimizedReaderEnabled(), - false), - booleanProperty( - PARQUET_OPTIMIZED_NESTED_READER_ENABLED, - "Use optimized Parquet reader for nested columns", - parquetReaderConfig.isOptimizedNestedReaderEnabled(), - false), booleanProperty( SIZE_BASED_SPLIT_WEIGHTS_ENABLED, format("If enabled, size-based splitting ensures that each batch of splits has enough data to process as defined by %s", STANDARD_SPLIT_WEIGHT_SIZE), diff --git a/plugin/trino-hudi/src/main/java/io/trino/plugin/hudi/InternalHudiConnectorFactory.java b/plugin/trino-hudi/src/main/java/io/trino/plugin/hudi/InternalHudiConnectorFactory.java index c8ea38715ca6..0af38058d2a9 100644 --- a/plugin/trino-hudi/src/main/java/io/trino/plugin/hudi/InternalHudiConnectorFactory.java +++ b/plugin/trino-hudi/src/main/java/io/trino/plugin/hudi/InternalHudiConnectorFactory.java @@ -78,7 +78,7 @@ public static Connector createConnector( new HdfsAuthenticationModule(), fileSystemFactory .map(factory -> (Module) binder -> binder.bind(TrinoFileSystemFactory.class).toInstance(factory)) - .orElseGet(FileSystemModule::new), + .orElseGet(() -> new FileSystemModule(catalogName, context.getNodeManager(), context.getOpenTelemetry(), false)), new MBeanServerModule(), binder -> { binder.bind(OpenTelemetry.class).toInstance(context.getOpenTelemetry()); diff --git a/plugin/trino-hudi/src/test/java/io/trino/plugin/hudi/S3HudiQueryRunner.java b/plugin/trino-hudi/src/test/java/io/trino/plugin/hudi/S3HudiQueryRunner.java index 8f29552f6c80..e223e3d1a73e 100644 --- a/plugin/trino-hudi/src/test/java/io/trino/plugin/hudi/S3HudiQueryRunner.java +++ b/plugin/trino-hudi/src/test/java/io/trino/plugin/hudi/S3HudiQueryRunner.java @@ -65,8 +65,7 @@ public static DistributedQueryRunner create( HiveMetastore metastore = new BridgingHiveMetastore( testingThriftHiveMetastoreBuilder() .metastoreClient(hiveMinioDataLake.getHiveHadoop().getHiveMetastoreEndpoint()) - .hdfsEnvironment(hdfsEnvironment) - .build()); + .build(x -> {})); Database database = Database.builder() .setDatabaseName(TPCH_SCHEMA) .setOwnerName(Optional.of("public")) diff --git a/plugin/trino-hudi/src/test/java/io/trino/plugin/hudi/TestHudiPartitionManager.java b/plugin/trino-hudi/src/test/java/io/trino/plugin/hudi/TestHudiPartitionManager.java index f2e8c1cca027..df4f84ceb78a 100644 --- a/plugin/trino-hudi/src/test/java/io/trino/plugin/hudi/TestHudiPartitionManager.java +++ b/plugin/trino-hudi/src/test/java/io/trino/plugin/hudi/TestHudiPartitionManager.java @@ -32,7 +32,6 @@ import static io.trino.plugin.hive.HiveType.HIVE_STRING; import static io.trino.plugin.hive.TableType.MANAGED_TABLE; import static io.trino.plugin.hive.metastore.StorageFormat.fromHiveStorageFormat; -import static io.trino.plugin.hive.util.HiveBucketing.BucketingVersion.BUCKETING_V1; import static io.trino.plugin.hudi.model.HudiTableType.COPY_ON_WRITE; import static org.testng.Assert.assertEquals; @@ -54,7 +53,6 @@ public class TestHudiPartitionManager Optional.of(LOCATION), Optional.of(new HiveBucketProperty( ImmutableList.of(BUCKET_COLUMN.getName()), - BUCKETING_V1, 2, ImmutableList.of())), false, diff --git a/plugin/trino-iceberg/pom.xml b/plugin/trino-iceberg/pom.xml index c37e9c99f077..b9f0d978ca2b 100644 --- a/plugin/trino-iceberg/pom.xml +++ b/plugin/trino-iceberg/pom.xml @@ -25,20 +25,10 @@ --> instances - 0.59.0 + 0.103.3 - - com.amazonaws - aws-java-sdk-core - - - - com.amazonaws - aws-java-sdk-glue - - com.fasterxml.jackson.core jackson-core @@ -58,6 +48,12 @@ com.google.guava guava + + + com.google.code.findbugs + jsr305 + + @@ -77,12 +73,12 @@ io.airlift - configuration + concurrent io.airlift - event + configuration @@ -132,7 +128,7 @@ io.trino - trino-hdfs + trino-filesystem-s3 @@ -194,6 +190,12 @@ org.apache.avro avro + + + org.slf4j + slf4j-api + + @@ -213,6 +215,11 @@ iceberg-api + + org.apache.iceberg + iceberg-aws + + org.apache.iceberg iceberg-core @@ -222,6 +229,20 @@ org.apache.iceberg iceberg-nessie ${dep.iceberg.version} + + + com.fasterxml.jackson.core + jackson-annotations + + + com.fasterxml.jackson.core + jackson-databind + + + org.slf4j + slf4j-api + + @@ -247,11 +268,12 @@ org.apache.parquet parquet-format-structures - - - - org.apache.parquet - parquet-hadoop + + + javax.annotation + javax.annotation-api + + @@ -263,12 +285,32 @@ org.projectnessie.nessie nessie-client ${dep.nessie.version} + + + com.fasterxml.jackson.core + jackson-annotations + + + com.fasterxml.jackson.core + jackson-databind + + org.projectnessie.nessie nessie-model ${dep.nessie.version} + + + com.fasterxml.jackson.core + jackson-annotations + + + com.fasterxml.jackson.core + jackson-databind + + @@ -281,6 +323,41 @@ jmxutils + + software.amazon.awssdk + auth + + + + software.amazon.awssdk + aws-core + + + + software.amazon.awssdk + glue + + + + software.amazon.awssdk + identity-spi + + + + software.amazon.awssdk + regions + + + + software.amazon.awssdk + sdk-core + + + + software.amazon.awssdk + sts + + com.fasterxml.jackson.core jackson-annotations @@ -317,18 +394,6 @@ provided - - com.amazonaws - aws-java-sdk-s3 - runtime - - - - io.airlift - concurrent - runtime - - io.airlift log-manager @@ -368,14 +433,12 @@ org.apache.httpcomponents.client5 httpclient5 - 5.2.1 runtime org.apache.httpcomponents.core5 httpcore5 - 5.2.1 runtime @@ -398,6 +461,12 @@ runtime + + software.amazon.awssdk + s3 + runtime + + io.airlift http-server @@ -454,6 +523,12 @@ io.trino trino-exchange-filesystem test + + + com.google.code.findbugs + jsr305 + + @@ -470,6 +545,12 @@ test + + io.trino + trino-hdfs + test + + io.trino trino-hive @@ -728,6 +809,7 @@ **/TestIcebergS3AndGlueMetastoreTest.java **/TestIcebergGcsConnectorSmokeTest.java **/TestIcebergAbfsConnectorSmokeTest.java + **/TestIcebergS3TablesConnectorSmokeTes.java diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/AllManifestsTable.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/AllManifestsTable.java new file mode 100644 index 000000000000..59c6381f0ee9 --- /dev/null +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/AllManifestsTable.java @@ -0,0 +1,99 @@ +/* + * 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.iceberg; + +import com.google.common.collect.ImmutableList; +import io.trino.plugin.iceberg.util.PageListBuilder; +import io.trino.spi.block.ArrayBlockBuilder; +import io.trino.spi.block.RowBlockBuilder; +import io.trino.spi.connector.ColumnMetadata; +import io.trino.spi.connector.ConnectorTableMetadata; +import io.trino.spi.connector.SchemaTableName; +import io.trino.spi.type.ArrayType; +import io.trino.spi.type.RowType; +import io.trino.spi.type.TimeZoneKey; +import org.apache.iceberg.StructLike; +import org.apache.iceberg.Table; + +import java.util.List; +import java.util.concurrent.ExecutorService; + +import static io.trino.spi.type.BigintType.BIGINT; +import static io.trino.spi.type.BooleanType.BOOLEAN; +import static io.trino.spi.type.IntegerType.INTEGER; +import static io.trino.spi.type.VarcharType.VARCHAR; +import static java.util.Objects.requireNonNull; +import static org.apache.iceberg.MetadataTableType.ALL_MANIFESTS; + +public class AllManifestsTable + extends BaseSystemTable +{ + public AllManifestsTable(SchemaTableName tableName, Table icebergTable, ExecutorService executor) + { + super(requireNonNull(icebergTable, "icebergTable is null"), + new ConnectorTableMetadata(requireNonNull(tableName, "tableName is null"), ImmutableList.builder() + .add(new ColumnMetadata("path", VARCHAR)) + .add(new ColumnMetadata("length", BIGINT)) + .add(new ColumnMetadata("partition_spec_id", INTEGER)) + .add(new ColumnMetadata("added_snapshot_id", BIGINT)) + .add(new ColumnMetadata("added_data_files_count", INTEGER)) + .add(new ColumnMetadata("existing_data_files_count", INTEGER)) + .add(new ColumnMetadata("deleted_data_files_count", INTEGER)) + .add(new ColumnMetadata("partition_summaries", new ArrayType(RowType.rowType( + RowType.field("contains_null", BOOLEAN), + RowType.field("contains_nan", BOOLEAN), + RowType.field("lower_bound", VARCHAR), + RowType.field("upper_bound", VARCHAR))))) + .build()), + ALL_MANIFESTS, + executor); + } + + @Override + protected void addRow(PageListBuilder pagesBuilder, Row row, TimeZoneKey timeZoneKey) + { + pagesBuilder.beginRow(); + pagesBuilder.appendVarchar(row.get("path", String.class)); + pagesBuilder.appendBigint(row.get("length", Long.class)); + pagesBuilder.appendInteger(row.get("partition_spec_id", Integer.class)); + pagesBuilder.appendBigint(row.get("added_snapshot_id", Long.class)); + pagesBuilder.appendInteger(row.get("added_data_files_count", Integer.class)); + pagesBuilder.appendInteger(row.get("existing_data_files_count", Integer.class)); + pagesBuilder.appendInteger(row.get("deleted_data_files_count", Integer.class)); + //noinspection unchecked + appendPartitionSummaries((ArrayBlockBuilder) pagesBuilder.nextColumn(), row.get("partition_summaries", List.class)); + pagesBuilder.endRow(); + } + + private static void appendPartitionSummaries(ArrayBlockBuilder arrayBuilder, List partitionSummaries) + { + arrayBuilder.buildEntry(elementBuilder -> { + for (StructLike partitionSummary : partitionSummaries) { + ((RowBlockBuilder) elementBuilder).buildEntry(fieldBuilders -> { + BOOLEAN.writeBoolean(fieldBuilders.get(0), partitionSummary.get(0, Boolean.class)); // required contains_null + Boolean containsNan = partitionSummary.get(1, Boolean.class); + if (containsNan == null) { + // This usually occurs when reading from V1 table, where contains_nan is not populated. + fieldBuilders.get(1).appendNull(); + } + else { + BOOLEAN.writeBoolean(fieldBuilders.get(1), containsNan); + } + VARCHAR.writeString(fieldBuilders.get(2), partitionSummary.get(2, String.class)); // optional lower_bound (human-readable) + VARCHAR.writeString(fieldBuilders.get(3), partitionSummary.get(3, String.class)); // optional upper_bound (human-readable) + }); + } + }); + } +} diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/BaseSystemTable.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/BaseSystemTable.java new file mode 100644 index 000000000000..adc5e5523ed4 --- /dev/null +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/BaseSystemTable.java @@ -0,0 +1,127 @@ +/* + * 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.iceberg; + +import com.google.common.collect.ImmutableMap; +import io.trino.plugin.iceberg.util.PageListBuilder; +import io.trino.spi.Page; +import io.trino.spi.connector.ConnectorPageSource; +import io.trino.spi.connector.ConnectorSession; +import io.trino.spi.connector.ConnectorTableMetadata; +import io.trino.spi.connector.ConnectorTransactionHandle; +import io.trino.spi.connector.FixedPageSource; +import io.trino.spi.connector.SystemTable; +import io.trino.spi.predicate.TupleDomain; +import io.trino.spi.type.TimeZoneKey; +import org.apache.iceberg.DataTask; +import org.apache.iceberg.FileScanTask; +import org.apache.iceberg.MetadataTableType; +import org.apache.iceberg.StructLike; +import org.apache.iceberg.Table; +import org.apache.iceberg.TableScan; +import org.apache.iceberg.io.CloseableIterable; + +import java.io.IOException; +import java.io.UncheckedIOException; +import java.util.List; +import java.util.Map; +import java.util.concurrent.ExecutorService; + +import static com.google.common.collect.ImmutableMap.toImmutableMap; +import static com.google.common.collect.Maps.immutableEntry; +import static com.google.common.collect.Streams.mapWithIndex; +import static java.util.Objects.requireNonNull; +import static org.apache.iceberg.MetadataTableUtils.createMetadataTableInstance; + +public abstract class BaseSystemTable + implements SystemTable +{ + private final Table icebergTable; + private final ConnectorTableMetadata tableMetadata; + private final MetadataTableType metadataTableType; + private final ExecutorService executor; + + BaseSystemTable(Table icebergTable, ConnectorTableMetadata tableMetadata, MetadataTableType metadataTableType, ExecutorService executor) + { + this.icebergTable = requireNonNull(icebergTable, "icebergTable is null"); + this.tableMetadata = requireNonNull(tableMetadata, "tableMetadata is null"); + this.metadataTableType = requireNonNull(metadataTableType, "metadataTableType is null"); + this.executor = requireNonNull(executor, "executor is null"); + } + + @Override + public Distribution getDistribution() + { + return Distribution.SINGLE_COORDINATOR; + } + + @Override + public ConnectorTableMetadata getTableMetadata() + { + return tableMetadata; + } + + @Override + public ConnectorPageSource pageSource(ConnectorTransactionHandle transactionHandle, ConnectorSession session, TupleDomain constraint) + { + return new FixedPageSource(buildPages(tableMetadata, session, icebergTable, metadataTableType)); + } + + private List buildPages(ConnectorTableMetadata tableMetadata, ConnectorSession session, Table icebergTable, MetadataTableType metadataTableType) + { + PageListBuilder pagesBuilder = PageListBuilder.forTable(tableMetadata); + + TableScan tableScan = createMetadataTableInstance(icebergTable, metadataTableType).newScan().planWith(executor); + TimeZoneKey timeZoneKey = session.getTimeZoneKey(); + + Map columnNameToPosition = mapWithIndex(tableScan.schema().columns().stream(), + (column, position) -> immutableEntry(column.name(), Long.valueOf(position).intValue())) + .collect(toImmutableMap(Map.Entry::getKey, Map.Entry::getValue)); + + try (CloseableIterable fileScanTasks = tableScan.planFiles()) { + fileScanTasks.forEach(fileScanTask -> addRows((DataTask) fileScanTask, pagesBuilder, timeZoneKey, columnNameToPosition)); + } + catch (IOException e) { + throw new UncheckedIOException(e); + } + + return pagesBuilder.build(); + } + + private void addRows(DataTask dataTask, PageListBuilder pagesBuilder, TimeZoneKey timeZoneKey, Map columnNameToPositionInSchema) + { + try (CloseableIterable dataRows = dataTask.rows()) { + dataRows.forEach(dataTaskRow -> addRow(pagesBuilder, new Row(dataTaskRow, columnNameToPositionInSchema), timeZoneKey)); + } + catch (IOException e) { + throw new UncheckedIOException(e); + } + } + + protected abstract void addRow(PageListBuilder pagesBuilder, Row row, TimeZoneKey timeZoneKey); + + public record Row(StructLike structLike, Map columnNameToPositionInSchema) + { + public Row + { + requireNonNull(structLike, "structLike is null"); + columnNameToPositionInSchema = ImmutableMap.copyOf(columnNameToPositionInSchema); + } + + public T get(String columnName, Class javaClass) + { + return structLike.get(columnNameToPositionInSchema.get(columnName), javaClass); + } + } +} diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/ColumnIdentity.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/ColumnIdentity.java index 161120adf541..3272f2042c7c 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/ColumnIdentity.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/ColumnIdentity.java @@ -18,6 +18,7 @@ import com.fasterxml.jackson.annotation.JsonProperty; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; +import io.airlift.slice.SizeOf; import org.apache.iceberg.types.Types; import java.util.List; @@ -27,6 +28,9 @@ import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.collect.ImmutableList.toImmutableList; import static com.google.common.collect.Iterables.getOnlyElement; +import static io.airlift.slice.SizeOf.estimatedSizeOf; +import static io.airlift.slice.SizeOf.instanceSize; +import static io.airlift.slice.SizeOf.sizeOf; import static io.trino.plugin.iceberg.ColumnIdentity.TypeCategory.ARRAY; import static io.trino.plugin.iceberg.ColumnIdentity.TypeCategory.MAP; import static io.trino.plugin.iceberg.ColumnIdentity.TypeCategory.PRIMITIVE; @@ -36,6 +40,8 @@ public class ColumnIdentity { + private static final int INSTANCE_SIZE = instanceSize(ColumnIdentity.class); + private final int id; private final String name; private final TypeCategory typeCategory; @@ -134,6 +140,16 @@ public String toString() return id + ":" + name; } + public long getRetainedSizeInBytes() + { + // type is not accounted for as the instances are cached (by TypeRegistry) and shared + return INSTANCE_SIZE + + sizeOf(id) + + estimatedSizeOf(name) + + estimatedSizeOf(children, SizeOf::sizeOf, ColumnIdentity::getRetainedSizeInBytes) + + estimatedSizeOf(childFieldIdToIndex, SizeOf::sizeOf, SizeOf::sizeOf); + } + public enum TypeCategory { PRIMITIVE, diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/CommitTaskData.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/CommitTaskData.java index 9870f0b03502..644d8f47b646 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/CommitTaskData.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/CommitTaskData.java @@ -17,6 +17,7 @@ import com.fasterxml.jackson.annotation.JsonProperty; import org.apache.iceberg.FileContent; +import java.util.List; import java.util.Optional; import static java.util.Objects.requireNonNull; @@ -31,6 +32,7 @@ public class CommitTaskData private final Optional partitionDataJson; private final FileContent content; private final Optional referencedDataFile; + private final Optional> fileSplitOffsets; @JsonCreator public CommitTaskData( @@ -41,7 +43,8 @@ public CommitTaskData( @JsonProperty("partitionSpecJson") String partitionSpecJson, @JsonProperty("partitionDataJson") Optional partitionDataJson, @JsonProperty("content") FileContent content, - @JsonProperty("referencedDataFile") Optional referencedDataFile) + @JsonProperty("referencedDataFile") Optional referencedDataFile, + @JsonProperty("fileSplitOffsets") Optional> fileSplitOffsets) { this.path = requireNonNull(path, "path is null"); this.fileFormat = requireNonNull(fileFormat, "fileFormat is null"); @@ -51,6 +54,7 @@ public CommitTaskData( this.partitionDataJson = requireNonNull(partitionDataJson, "partitionDataJson is null"); this.content = requireNonNull(content, "content is null"); this.referencedDataFile = requireNonNull(referencedDataFile, "referencedDataFile is null"); + this.fileSplitOffsets = requireNonNull(fileSplitOffsets, "fileSplitOffsets is null"); } @JsonProperty @@ -100,4 +104,10 @@ public Optional getReferencedDataFile() { return referencedDataFile; } + + @JsonProperty + public Optional> getFileSplitOffsets() + { + return fileSplitOffsets; + } } diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/CreateTableException.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/CreateTableException.java new file mode 100644 index 000000000000..66dfccc6a424 --- /dev/null +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/CreateTableException.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.iceberg; + +import io.trino.spi.TrinoException; +import io.trino.spi.connector.SchemaTableName; +import org.apache.iceberg.exceptions.CleanableFailure; + +import static io.trino.plugin.iceberg.IcebergErrorCode.ICEBERG_COMMIT_ERROR; +import static java.lang.String.format; + +public class CreateTableException + extends TrinoException + implements CleanableFailure +{ + public CreateTableException(Throwable throwable, SchemaTableName tableName) + { + super(ICEBERG_COMMIT_ERROR, format("Failed to create table %s: %s", tableName, throwable.getMessage()), throwable); + } +} diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/EntriesTable.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/EntriesTable.java new file mode 100644 index 000000000000..4d76b6191b7d --- /dev/null +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/EntriesTable.java @@ -0,0 +1,356 @@ +/* + * 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.iceberg; + +import com.google.common.collect.ImmutableList; +import io.trino.plugin.iceberg.util.PageListBuilder; +import io.trino.spi.block.ArrayBlockBuilder; +import io.trino.spi.block.MapBlockBuilder; +import io.trino.spi.block.RowBlockBuilder; +import io.trino.spi.connector.ColumnMetadata; +import io.trino.spi.connector.ConnectorTableMetadata; +import io.trino.spi.connector.SchemaTableName; +import io.trino.spi.type.ArrayType; +import io.trino.spi.type.RowType; +import io.trino.spi.type.TimeZoneKey; +import io.trino.spi.type.TypeManager; +import io.trino.spi.type.TypeSignature; +import jakarta.annotation.Nullable; +import org.apache.iceberg.MetadataTableType; +import org.apache.iceberg.MetricsUtil.ReadableMetricsStruct; +import org.apache.iceberg.PartitionField; +import org.apache.iceberg.Table; +import org.apache.iceberg.transforms.Transforms; +import org.apache.iceberg.types.Conversions; +import org.apache.iceberg.types.Type; +import org.apache.iceberg.types.Types; +import org.apache.iceberg.util.StructProjection; + +import java.nio.ByteBuffer; +import java.util.Comparator; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.concurrent.ExecutorService; + +import static com.google.common.base.Preconditions.checkArgument; +import static com.google.common.collect.ImmutableList.toImmutableList; +import static io.airlift.slice.Slices.wrappedBuffer; +import static io.trino.plugin.iceberg.FilesTable.getIcebergIdToTypeMapping; +import static io.trino.plugin.iceberg.IcebergTypes.convertIcebergValueToTrino; +import static io.trino.plugin.iceberg.IcebergUtil.getPartitionColumnType; +import static io.trino.plugin.iceberg.IcebergUtil.partitionTypes; +import static io.trino.plugin.iceberg.IcebergUtil.primitiveFieldTypes; +import static io.trino.plugin.iceberg.PartitionsTable.getAllPartitionFields; +import static io.trino.spi.type.BigintType.BIGINT; +import static io.trino.spi.type.IntegerType.INTEGER; +import static io.trino.spi.type.StandardTypes.JSON; +import static io.trino.spi.type.TypeSignature.mapType; +import static io.trino.spi.type.TypeUtils.writeNativeValue; +import static io.trino.spi.type.VarbinaryType.VARBINARY; +import static io.trino.spi.type.VarcharType.VARCHAR; +import static java.util.Objects.requireNonNull; +import static org.apache.iceberg.MetadataColumns.DELETE_FILE_PATH; +import static org.apache.iceberg.MetadataColumns.DELETE_FILE_POS; +import static org.apache.iceberg.MetadataTableType.ALL_ENTRIES; +import static org.apache.iceberg.MetadataTableType.ENTRIES; + +// https://iceberg.apache.org/docs/latest/spark-queries/#all-entries +// https://iceberg.apache.org/docs/latest/spark-queries/#entries +public class EntriesTable + extends BaseSystemTable +{ + private final Map idToTypeMapping; + private final List primitiveFields; + private final Optional partitionColumn; + private final List partitionTypes; + + public EntriesTable(TypeManager typeManager, SchemaTableName tableName, Table icebergTable, MetadataTableType metadataTableType, ExecutorService executor) + { + super( + requireNonNull(icebergTable, "icebergTable is null"), + new ConnectorTableMetadata( + requireNonNull(tableName, "tableName is null"), + columns(requireNonNull(typeManager, "typeManager is null"), icebergTable)), + metadataTableType, + executor); + checkArgument(metadataTableType == ALL_ENTRIES || metadataTableType == ENTRIES, "Unexpected metadata table type: %s", metadataTableType); + idToTypeMapping = getIcebergIdToTypeMapping(icebergTable.schema()); + primitiveFields = IcebergUtil.primitiveFields(icebergTable.schema()).stream() + .sorted(Comparator.comparing(Types.NestedField::name)) + .collect(toImmutableList()); + List partitionFields = getAllPartitionFields(icebergTable); + partitionColumn = getPartitionColumnType(partitionFields, icebergTable.schema(), typeManager); + partitionTypes = partitionTypes(partitionFields, primitiveFieldTypes(icebergTable.schema())); + } + + private static List columns(TypeManager typeManager, Table icebergTable) + { + return ImmutableList.builder() + .add(new ColumnMetadata("status", INTEGER)) + .add(new ColumnMetadata("snapshot_id", BIGINT)) + .add(new ColumnMetadata("sequence_number", BIGINT)) + .add(new ColumnMetadata("file_sequence_number", BIGINT)) + .add(new ColumnMetadata("data_file", RowType.from(dataFileFieldMetadata(typeManager, icebergTable)))) + .add(new ColumnMetadata("readable_metrics", typeManager.getType(new TypeSignature(JSON)))) + .build(); + } + + private static List dataFileFieldMetadata(TypeManager typeManager, Table icebergTable) + { + List partitionFields = getAllPartitionFields(icebergTable); + Optional partitionColumnType = getPartitionColumnType(partitionFields, icebergTable.schema(), typeManager); + + ImmutableList.Builder fields = ImmutableList.builder(); + fields.add(new RowType.Field(Optional.of("content"), INTEGER)); + fields.add(new RowType.Field(Optional.of("file_path"), VARCHAR)); + fields.add(new RowType.Field(Optional.of("file_format"), VARCHAR)); + fields.add(new RowType.Field(Optional.of("spec_id"), INTEGER)); + partitionColumnType.ifPresent(type -> fields.add(new RowType.Field(Optional.of("partition"), type.rowType()))); + fields.add(new RowType.Field(Optional.of("record_count"), BIGINT)); + fields.add(new RowType.Field(Optional.of("file_size_in_bytes"), BIGINT)); + fields.add(new RowType.Field(Optional.of("column_sizes"), typeManager.getType(mapType(INTEGER.getTypeSignature(), BIGINT.getTypeSignature())))); + fields.add(new RowType.Field(Optional.of("value_counts"), typeManager.getType(mapType(INTEGER.getTypeSignature(), BIGINT.getTypeSignature())))); + fields.add(new RowType.Field(Optional.of("null_value_counts"), typeManager.getType(mapType(INTEGER.getTypeSignature(), BIGINT.getTypeSignature())))); + fields.add(new RowType.Field(Optional.of("nan_value_counts"), typeManager.getType(mapType(INTEGER.getTypeSignature(), BIGINT.getTypeSignature())))); + fields.add(new RowType.Field(Optional.of("lower_bounds"), typeManager.getType(mapType(INTEGER.getTypeSignature(), VARCHAR.getTypeSignature())))); + fields.add(new RowType.Field(Optional.of("upper_bounds"), typeManager.getType(mapType(INTEGER.getTypeSignature(), VARCHAR.getTypeSignature())))); + fields.add(new RowType.Field(Optional.of("key_metadata"), VARBINARY)); + fields.add(new RowType.Field(Optional.of("split_offsets"), new ArrayType(BIGINT))); + fields.add(new RowType.Field(Optional.of("equality_ids"), new ArrayType(INTEGER))); + fields.add(new RowType.Field(Optional.of("sort_order_id"), INTEGER)); + return fields.build(); + } + + @Override + protected void addRow(PageListBuilder pagesBuilder, Row row, TimeZoneKey timeZoneKey) + { + pagesBuilder.beginRow(); + pagesBuilder.appendInteger(row.get("status", Integer.class)); + pagesBuilder.appendBigint(row.get("snapshot_id", Long.class)); + pagesBuilder.appendBigint(row.get("sequence_number", Long.class)); + pagesBuilder.appendBigint(row.get("file_sequence_number", Long.class)); + StructProjection dataFile = row.get("data_file", StructProjection.class); + appendDataFile((RowBlockBuilder) pagesBuilder.nextColumn(), dataFile); + ReadableMetricsStruct readableMetrics = row.get("readable_metrics", ReadableMetricsStruct.class); + String readableMetricsJson = FilesTable.toJson(readableMetrics, primitiveFields); + pagesBuilder.appendVarchar(readableMetricsJson); + pagesBuilder.endRow(); + } + + private void appendDataFile(RowBlockBuilder blockBuilder, StructProjection dataFile) + { + blockBuilder.buildEntry(fieldBuilders -> { + Integer content = dataFile.get(0, Integer.class); + INTEGER.writeLong(fieldBuilders.get(0), content); + + String filePath = dataFile.get(1, String.class); + VARCHAR.writeString(fieldBuilders.get(1), filePath); + + String fileFormat = dataFile.get(2, String.class); + VARCHAR.writeString(fieldBuilders.get(2), fileFormat); + + Integer specId = dataFile.get(3, Integer.class); + INTEGER.writeLong(fieldBuilders.get(3), Long.valueOf(specId)); + + partitionColumn.ifPresent(type -> { + StructProjection partition = dataFile.get(4, StructProjection.class); + RowBlockBuilder partitionBlockBuilder = (RowBlockBuilder) fieldBuilders.get(4); + partitionBlockBuilder.buildEntry(partitionBuilder -> { + for (int i = 0; i < type.rowType().getFields().size(); i++) { + Type icebergType = partitionTypes.get(i); + io.trino.spi.type.Type trinoType = type.rowType().getFields().get(i).getType(); + Object value = null; + Integer fieldId = type.fieldIds().get(i); + if (fieldId != null) { + value = convertIcebergValueToTrino(icebergType, partition.get(i, icebergType.typeId().javaClass())); + } + writeNativeValue(trinoType, partitionBuilder.get(i), value); + } + }); + }); + + int position = partitionColumn.isEmpty() ? 4 : 5; + Long recordCount = dataFile.get(position, Long.class); + BIGINT.writeLong(fieldBuilders.get(position), recordCount); + + Long fileSizeInBytes = dataFile.get(++position, Long.class); + BIGINT.writeLong(fieldBuilders.get(position), fileSizeInBytes); + + //noinspection unchecked + Map columnSizes = dataFile.get(++position, Map.class); + appendIntegerBigintMap((MapBlockBuilder) fieldBuilders.get(position), columnSizes); + + //noinspection unchecked + Map valueCounts = dataFile.get(++position, Map.class); + appendIntegerBigintMap((MapBlockBuilder) fieldBuilders.get(position), valueCounts); + + //noinspection unchecked + Map nullValueCounts = dataFile.get(++position, Map.class); + appendIntegerBigintMap((MapBlockBuilder) fieldBuilders.get(position), nullValueCounts); + + //noinspection unchecked + Map nanValueCounts = dataFile.get(++position, Map.class); + appendIntegerBigintMap((MapBlockBuilder) fieldBuilders.get(position), nanValueCounts); + + switch (ContentType.of(content)) { + case DATA, EQUALITY_DELETE -> { + //noinspection unchecked + Map lowerBounds = dataFile.get(++position, Map.class); + appendIntegerVarcharMap((MapBlockBuilder) fieldBuilders.get(position), lowerBounds); + + //noinspection unchecked + Map upperBounds = dataFile.get(++position, Map.class); + appendIntegerVarcharMap((MapBlockBuilder) fieldBuilders.get(position), upperBounds); + } + case POSITION_DELETE -> { + //noinspection unchecked + Map lowerBounds = dataFile.get(++position, Map.class); + appendBoundsForPositionDelete((MapBlockBuilder) fieldBuilders.get(position), lowerBounds); + + //noinspection unchecked + Map upperBounds = dataFile.get(++position, Map.class); + appendBoundsForPositionDelete((MapBlockBuilder) fieldBuilders.get(position), upperBounds); + } + } + + ByteBuffer keyMetadata = dataFile.get(++position, ByteBuffer.class); + if (keyMetadata == null) { + fieldBuilders.get(position).appendNull(); + } + else { + VARBINARY.writeSlice(fieldBuilders.get(position), wrappedBuffer(keyMetadata)); + } + + //noinspection unchecked + List splitOffsets = dataFile.get(++position, List.class); + appendBigintArray((ArrayBlockBuilder) fieldBuilders.get(position), splitOffsets); + + switch (ContentType.of(content)) { + case DATA -> { + // data files don't have equality ids + fieldBuilders.get(++position).appendNull(); + + Integer sortOrderId = dataFile.get(++position, Integer.class); + INTEGER.writeLong(fieldBuilders.get(position), Long.valueOf(sortOrderId)); + } + case POSITION_DELETE -> { + // position delete files don't have equality ids + fieldBuilders.get(++position).appendNull(); + + // position delete files don't have sort order id + fieldBuilders.get(++position).appendNull(); + } + case EQUALITY_DELETE -> { + //noinspection unchecked + List equalityIds = dataFile.get(++position, List.class); + appendIntegerArray((ArrayBlockBuilder) fieldBuilders.get(position), equalityIds); + + Integer sortOrderId = dataFile.get(++position, Integer.class); + INTEGER.writeLong(fieldBuilders.get(position), Long.valueOf(sortOrderId)); + } + } + }); + } + + public static void appendBigintArray(ArrayBlockBuilder blockBuilder, @Nullable List values) + { + if (values == null) { + blockBuilder.appendNull(); + return; + } + blockBuilder.buildEntry(elementBuilder -> { + for (Long value : values) { + BIGINT.writeLong(elementBuilder, value); + } + }); + } + + public static void appendIntegerArray(ArrayBlockBuilder blockBuilder, @Nullable List values) + { + if (values == null) { + blockBuilder.appendNull(); + return; + } + blockBuilder.buildEntry(elementBuilder -> { + for (Integer value : values) { + INTEGER.writeLong(elementBuilder, value); + } + }); + } + + private static void appendIntegerBigintMap(MapBlockBuilder blockBuilder, @Nullable Map values) + { + if (values == null) { + blockBuilder.appendNull(); + return; + } + blockBuilder.buildEntry((keyBuilder, valueBuilder) -> values.forEach((key, value) -> { + INTEGER.writeLong(keyBuilder, key); + BIGINT.writeLong(valueBuilder, value); + })); + } + + private void appendIntegerVarcharMap(MapBlockBuilder blockBuilder, @Nullable Map values) + { + if (values == null) { + blockBuilder.appendNull(); + return; + } + blockBuilder.buildEntry((keyBuilder, valueBuilder) -> values.forEach((key, value) -> { + Type type = idToTypeMapping.get(key); + INTEGER.writeLong(keyBuilder, key); + VARCHAR.writeString(valueBuilder, Transforms.identity().toHumanString(type, Conversions.fromByteBuffer(type, value))); + })); + } + + private static void appendBoundsForPositionDelete(MapBlockBuilder blockBuilder, @Nullable Map values) + { + if (values == null) { + blockBuilder.appendNull(); + return; + } + + blockBuilder.buildEntry((keyBuilder, valueBuilder) -> { + INTEGER.writeLong(keyBuilder, DELETE_FILE_POS.fieldId()); + ByteBuffer pos = values.get(DELETE_FILE_POS.fieldId()); + checkArgument(pos != null, "delete file pos is null"); + VARCHAR.writeString(valueBuilder, Transforms.identity().toHumanString(Types.LongType.get(), Conversions.fromByteBuffer(Types.LongType.get(), pos))); + + INTEGER.writeLong(keyBuilder, DELETE_FILE_PATH.fieldId()); + ByteBuffer path = values.get(DELETE_FILE_PATH.fieldId()); + checkArgument(path != null, "delete file path is null"); + VARCHAR.writeString(valueBuilder, Transforms.identity().toHumanString(Types.StringType.get(), Conversions.fromByteBuffer(Types.StringType.get(), path))); + }); + } + + private enum ContentType + { + DATA, + POSITION_DELETE, + EQUALITY_DELETE; + + static ContentType of(int content) + { + checkArgument(content >= 0 && content <= 2, "Unexpected content type: %s", content); + if (content == 0) { + return DATA; + } + if (content == 1) { + return POSITION_DELETE; + } + return EQUALITY_DELETE; + } + } +} diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/ExpressionConverter.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/ExpressionConverter.java index ceeeeb7b8e03..40819695b520 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/ExpressionConverter.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/ExpressionConverter.java @@ -32,8 +32,10 @@ import java.util.function.BiFunction; import static com.google.common.base.Preconditions.checkArgument; +import static io.trino.plugin.hive.util.HiveUtil.isStructuralType; import static io.trino.plugin.iceberg.IcebergMetadataColumn.isMetadataColumnId; import static io.trino.plugin.iceberg.IcebergTypes.convertTrinoValueToIceberg; +import static io.trino.spi.type.UuidType.UUID; import static java.lang.String.format; import static org.apache.iceberg.expressions.Expressions.alwaysFalse; import static org.apache.iceberg.expressions.Expressions.alwaysTrue; @@ -50,6 +52,21 @@ public final class ExpressionConverter { private ExpressionConverter() {} + public static boolean isConvertibleToIcebergExpression(Domain domain) + { + if (isStructuralType(domain.getType())) { + // structural types cannot be used to filter a table scan in Iceberg library. + return false; + } + + if (domain.getType() == UUID) { + // Iceberg orders UUID values differently than Trino (perhaps due to https://bugs.openjdk.org/browse/JDK-7025832), so allow only IS NULL / IS NOT NULL checks + return domain.isOnlyNull() || domain.getValues().isAll(); + } + + return true; + } + public static Expression toIcebergExpression(TupleDomain tupleDomain) { if (tupleDomain.isAll()) { diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/FilesTable.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/FilesTable.java index c70f79f74840..abfe2acfd3a2 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/FilesTable.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/FilesTable.java @@ -13,10 +13,13 @@ */ package io.trino.plugin.iceberg; +import com.fasterxml.jackson.core.JsonFactory; +import com.fasterxml.jackson.core.JsonGenerator; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import io.airlift.slice.Slice; import io.airlift.slice.Slices; +import io.trino.plugin.base.util.JsonUtils; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; import io.trino.spi.connector.ColumnMetadata; @@ -30,11 +33,18 @@ import io.trino.spi.predicate.TupleDomain; import io.trino.spi.type.ArrayType; import io.trino.spi.type.MapType; +import io.trino.spi.type.RowType; import io.trino.spi.type.TypeManager; +import io.trino.spi.type.TypeSignature; import jakarta.annotation.Nullable; -import org.apache.iceberg.DataFile; +import org.apache.iceberg.DataTask; import org.apache.iceberg.FileScanTask; +import org.apache.iceberg.MetricsUtil.ReadableColMetricsStruct; +import org.apache.iceberg.MetricsUtil.ReadableMetricsStruct; +import org.apache.iceberg.PartitionField; import org.apache.iceberg.Schema; +import org.apache.iceberg.SingleValueParser; +import org.apache.iceberg.StructLike; import org.apache.iceberg.Table; import org.apache.iceberg.TableScan; import org.apache.iceberg.io.CloseableGroup; @@ -46,55 +56,104 @@ import org.apache.iceberg.types.Types; import java.io.IOException; +import java.io.StringWriter; import java.io.UncheckedIOException; import java.nio.ByteBuffer; import java.util.ArrayList; +import java.util.Comparator; import java.util.List; import java.util.Map; import java.util.Optional; +import java.util.concurrent.ExecutorService; import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.collect.ImmutableList.toImmutableList; import static com.google.common.collect.ImmutableMap.toImmutableMap; +import static com.google.common.collect.Maps.immutableEntry; +import static com.google.common.collect.Streams.mapWithIndex; +import static io.trino.plugin.iceberg.IcebergTypes.convertIcebergValueToTrino; +import static io.trino.plugin.iceberg.IcebergUtil.getPartitionColumnType; +import static io.trino.plugin.iceberg.IcebergUtil.partitionTypes; import static io.trino.spi.block.MapValueBuilder.buildMapValue; +import static io.trino.spi.block.RowValueBuilder.buildRowValue; import static io.trino.spi.type.BigintType.BIGINT; import static io.trino.spi.type.IntegerType.INTEGER; +import static io.trino.spi.type.StandardTypes.JSON; import static io.trino.spi.type.TypeSignature.mapType; +import static io.trino.spi.type.TypeUtils.writeNativeValue; import static io.trino.spi.type.VarbinaryType.VARBINARY; import static io.trino.spi.type.VarcharType.VARCHAR; import static java.util.Objects.requireNonNull; +import static org.apache.iceberg.MetadataTableType.FILES; +import static org.apache.iceberg.MetadataTableUtils.createMetadataTableInstance; public class FilesTable implements SystemTable { + private static final JsonFactory JSON_FACTORY = JsonUtils.jsonFactoryBuilder().build(); + + private static final String CONTENT_COLUMN_NAME = "content"; + private static final String FILE_PATH_COLUMN_NAME = "file_path"; + private static final String FILE_FORMAT_COLUMN_NAME = "file_format"; + private static final String SPEC_ID_COLUMN_NAME = "spec_id"; + private static final String PARTITION_COLUMN_NAME = "partition"; + private static final String RECORD_COUNT_COLUMN_NAME = "record_count"; + private static final String FILE_SIZE_IN_BYTES_COLUMN_NAME = "file_size_in_bytes"; + private static final String COLUMN_SIZES_COLUMN_NAME = "column_sizes"; + private static final String VALUE_COUNTS_COLUMN_NAME = "value_counts"; + private static final String NULL_VALUE_COUNTS_COLUMN_NAME = "null_value_counts"; + private static final String NAN_VALUE_COUNTS_COLUMN_NAME = "nan_value_counts"; + private static final String LOWER_BOUNDS_COLUMN_NAME = "lower_bounds"; + private static final String UPPER_BOUNDS_COLUMN_NAME = "upper_bounds"; + private static final String KEY_METADATA_COLUMN_NAME = "key_metadata"; + private static final String SPLIT_OFFSETS_COLUMN_NAME = "split_offsets"; + private static final String EQUALITY_IDS_COLUMN_NAME = "equality_ids"; + private static final String SORT_ORDER_ID_COLUMN_NAME = "sort_order_id"; + private static final String READABLE_METRICS_COLUMN_NAME = "readable_metrics"; private final ConnectorTableMetadata tableMetadata; private final TypeManager typeManager; private final Table icebergTable; private final Optional snapshotId; + private final Optional partitionColumnType; + private final Map idToPrimitiveTypeMapping; + private final List primitiveFields; + private final ExecutorService executor; - public FilesTable(SchemaTableName tableName, TypeManager typeManager, Table icebergTable, Optional snapshotId) + public FilesTable(SchemaTableName tableName, TypeManager typeManager, Table icebergTable, Optional snapshotId, ExecutorService executor) { this.icebergTable = requireNonNull(icebergTable, "icebergTable is null"); this.typeManager = requireNonNull(typeManager, "typeManager is null"); - tableMetadata = new ConnectorTableMetadata(requireNonNull(tableName, "tableName is null"), - ImmutableList.builder() - .add(new ColumnMetadata("content", INTEGER)) - .add(new ColumnMetadata("file_path", VARCHAR)) - .add(new ColumnMetadata("file_format", VARCHAR)) - .add(new ColumnMetadata("record_count", BIGINT)) - .add(new ColumnMetadata("file_size_in_bytes", BIGINT)) - .add(new ColumnMetadata("column_sizes", typeManager.getType(mapType(INTEGER.getTypeSignature(), BIGINT.getTypeSignature())))) - .add(new ColumnMetadata("value_counts", typeManager.getType(mapType(INTEGER.getTypeSignature(), BIGINT.getTypeSignature())))) - .add(new ColumnMetadata("null_value_counts", typeManager.getType(mapType(INTEGER.getTypeSignature(), BIGINT.getTypeSignature())))) - .add(new ColumnMetadata("nan_value_counts", typeManager.getType(mapType(INTEGER.getTypeSignature(), BIGINT.getTypeSignature())))) - .add(new ColumnMetadata("lower_bounds", typeManager.getType(mapType(INTEGER.getTypeSignature(), VARCHAR.getTypeSignature())))) - .add(new ColumnMetadata("upper_bounds", typeManager.getType(mapType(INTEGER.getTypeSignature(), VARCHAR.getTypeSignature())))) - .add(new ColumnMetadata("key_metadata", VARBINARY)) - .add(new ColumnMetadata("split_offsets", new ArrayType(BIGINT))) - .add(new ColumnMetadata("equality_ids", new ArrayType(INTEGER))) - .build()); + List partitionFields = PartitionsTable.getAllPartitionFields(icebergTable); + partitionColumnType = getPartitionColumnType(partitionFields, icebergTable.schema(), typeManager); + idToPrimitiveTypeMapping = IcebergUtil.primitiveFieldTypes(icebergTable.schema()); + primitiveFields = IcebergUtil.primitiveFields(icebergTable.schema()).stream() + .sorted(Comparator.comparing(Types.NestedField::name)) + .collect(toImmutableList()); + + ImmutableList.Builder columns = ImmutableList.builder(); + columns.add(new ColumnMetadata(CONTENT_COLUMN_NAME, INTEGER)); + columns.add(new ColumnMetadata(FILE_PATH_COLUMN_NAME, VARCHAR)); + columns.add(new ColumnMetadata(FILE_FORMAT_COLUMN_NAME, VARCHAR)); + columns.add(new ColumnMetadata(SPEC_ID_COLUMN_NAME, INTEGER)); + partitionColumnType.ifPresent(type -> columns.add(new ColumnMetadata(PARTITION_COLUMN_NAME, type.rowType()))); + columns.add(new ColumnMetadata(RECORD_COUNT_COLUMN_NAME, BIGINT)); + columns.add(new ColumnMetadata(FILE_SIZE_IN_BYTES_COLUMN_NAME, BIGINT)); + columns.add(new ColumnMetadata(COLUMN_SIZES_COLUMN_NAME, typeManager.getType(mapType(INTEGER.getTypeSignature(), BIGINT.getTypeSignature())))); + columns.add(new ColumnMetadata(VALUE_COUNTS_COLUMN_NAME, typeManager.getType(mapType(INTEGER.getTypeSignature(), BIGINT.getTypeSignature())))); + columns.add(new ColumnMetadata(NULL_VALUE_COUNTS_COLUMN_NAME, typeManager.getType(mapType(INTEGER.getTypeSignature(), BIGINT.getTypeSignature())))); + columns.add(new ColumnMetadata(NAN_VALUE_COUNTS_COLUMN_NAME, typeManager.getType(mapType(INTEGER.getTypeSignature(), BIGINT.getTypeSignature())))); + columns.add(new ColumnMetadata(LOWER_BOUNDS_COLUMN_NAME, typeManager.getType(mapType(INTEGER.getTypeSignature(), VARCHAR.getTypeSignature())))); + columns.add(new ColumnMetadata(UPPER_BOUNDS_COLUMN_NAME, typeManager.getType(mapType(INTEGER.getTypeSignature(), VARCHAR.getTypeSignature())))); + columns.add(new ColumnMetadata(KEY_METADATA_COLUMN_NAME, VARBINARY)); + columns.add(new ColumnMetadata(SPLIT_OFFSETS_COLUMN_NAME, new ArrayType(BIGINT))); + columns.add(new ColumnMetadata(EQUALITY_IDS_COLUMN_NAME, new ArrayType(INTEGER))); + columns.add(new ColumnMetadata(SORT_ORDER_ID_COLUMN_NAME, INTEGER)); + columns.add(new ColumnMetadata(READABLE_METRICS_COLUMN_NAME, typeManager.getType(new TypeSignature(JSON)))); + + tableMetadata = new ConnectorTableMetadata(requireNonNull(tableName, "tableName is null"), columns.build()); this.snapshotId = requireNonNull(snapshotId, "snapshotId is null"); + this.executor = requireNonNull(executor, "executor is null"); } @Override @@ -120,11 +179,26 @@ public RecordCursor cursor(ConnectorTransactionHandle transactionHandle, Connect } Map idToTypeMapping = getIcebergIdToTypeMapping(icebergTable.schema()); - TableScan tableScan = icebergTable.newScan() + TableScan tableScan = createMetadataTableInstance(icebergTable, FILES) + .newScan() .useSnapshot(snapshotId.get()) - .includeColumnStats(); + .includeColumnStats() + .planWith(executor); + + Map columnNameToPosition = mapWithIndex(tableScan.schema().columns().stream(), + (column, position) -> immutableEntry(column.name(), Long.valueOf(position).intValue())) + .collect(toImmutableMap(Map.Entry::getKey, Map.Entry::getValue)); - PlanFilesIterable planFilesIterable = new PlanFilesIterable(tableScan.planFiles(), idToTypeMapping, types, typeManager); + PlanFilesIterable planFilesIterable = new PlanFilesIterable( + tableScan.planFiles(), + primitiveFields, + idToTypeMapping, + types, + columnNameToPosition, + typeManager, + partitionColumnType, + PartitionsTable.getAllPartitionFields(icebergTable), + idToPrimitiveTypeMapping); return planFilesIterable.cursor(); } @@ -133,26 +207,46 @@ private static class PlanFilesIterable implements Iterable> { private final CloseableIterable planFiles; + private final List primitiveFields; private final Map idToTypeMapping; private final List types; + private final Map columnNameToPosition; private boolean closed; private final MapType integerToBigintMapType; private final MapType integerToVarcharMapType; + private final Optional partitionColumnType; + private final List partitionFields; + private final Map idToPrimitiveTypeMapping; - public PlanFilesIterable(CloseableIterable planFiles, Map idToTypeMapping, List types, TypeManager typeManager) + public PlanFilesIterable( + CloseableIterable planFiles, + List primitiveFields, + Map idToTypeMapping, + List types, + Map columnNameToPosition, + TypeManager typeManager, + Optional partitionColumnType, + List partitionFields, + Map idToPrimitiveTypeMapping) { this.planFiles = requireNonNull(planFiles, "planFiles is null"); + this.primitiveFields = ImmutableList.copyOf(requireNonNull(primitiveFields, "primitiveFields is null")); this.idToTypeMapping = ImmutableMap.copyOf(requireNonNull(idToTypeMapping, "idToTypeMapping is null")); this.types = ImmutableList.copyOf(requireNonNull(types, "types is null")); + this.columnNameToPosition = ImmutableMap.copyOf(requireNonNull(columnNameToPosition, "columnNameToPosition is null")); this.integerToBigintMapType = new MapType(INTEGER, BIGINT, typeManager.getTypeOperators()); this.integerToVarcharMapType = new MapType(INTEGER, VARCHAR, typeManager.getTypeOperators()); + this.partitionColumnType = requireNonNull(partitionColumnType, "partitionColumnType is null"); + this.partitionFields = ImmutableList.copyOf(requireNonNull(partitionFields, "partitionFields is null")); + this.idToPrimitiveTypeMapping = ImmutableMap.copyOf(requireNonNull(idToPrimitiveTypeMapping, "idToPrimitiveTypeMapping is null")); addCloseable(planFiles); } public RecordCursor cursor() { CloseableIterator> iterator = this.iterator(); - return new InMemoryRecordSet.InMemoryRecordCursor(types, iterator) { + return new InMemoryRecordSet.InMemoryRecordCursor(types, iterator) + { @Override public void close() { @@ -172,64 +266,115 @@ public CloseableIterator> iterator() final CloseableIterator planFilesIterator = planFiles.iterator(); addCloseable(planFilesIterator); - return new CloseableIterator<>() { + return new CloseableIterator<>() + { + private CloseableIterator currentIterator = CloseableIterator.empty(); + @Override public boolean hasNext() { - return !closed && planFilesIterator.hasNext(); + updateCurrentIterator(); + return !closed && currentIterator.hasNext(); } @Override public List next() { - return getRecord(planFilesIterator.next().file()); + updateCurrentIterator(); + return getRecord(currentIterator.next()); + } + + private void updateCurrentIterator() + { + try { + while (!closed && !currentIterator.hasNext() && planFilesIterator.hasNext()) { + currentIterator.close(); + DataTask dataTask = (DataTask) planFilesIterator.next(); + currentIterator = dataTask.rows().iterator(); + } + } + catch (IOException e) { + throw new UncheckedIOException(e); + } } @Override public void close() throws IOException { - PlanFilesIterable.super.close(); + currentIterator.close(); + FilesTable.PlanFilesIterable.super.close(); closed = true; } }; } - private List getRecord(DataFile dataFile) + private List getRecord(StructLike structLike) { List columns = new ArrayList<>(); - columns.add(dataFile.content().id()); - columns.add(dataFile.path().toString()); - columns.add(dataFile.format().name()); - columns.add(dataFile.recordCount()); - columns.add(dataFile.fileSizeInBytes()); - columns.add(getIntegerBigintMapBlock(dataFile.columnSizes())); - columns.add(getIntegerBigintMapBlock(dataFile.valueCounts())); - columns.add(getIntegerBigintMapBlock(dataFile.nullValueCounts())); - columns.add(getIntegerBigintMapBlock(dataFile.nanValueCounts())); - columns.add(getIntegerVarcharMapBlock(dataFile.lowerBounds())); - columns.add(getIntegerVarcharMapBlock(dataFile.upperBounds())); - columns.add(toVarbinarySlice(dataFile.keyMetadata())); - columns.add(toBigintArrayBlock(dataFile.splitOffsets())); - columns.add(toIntegerArrayBlock(dataFile.equalityFieldIds())); + columns.add(structLike.get(columnNameToPosition.get(CONTENT_COLUMN_NAME), Integer.class)); + columns.add(structLike.get(columnNameToPosition.get(FILE_PATH_COLUMN_NAME), String.class)); + columns.add(structLike.get(columnNameToPosition.get(FILE_FORMAT_COLUMN_NAME), String.class)); + columns.add(structLike.get(columnNameToPosition.get(SPEC_ID_COLUMN_NAME), Integer.class)); + + if (columnNameToPosition.containsKey(PARTITION_COLUMN_NAME)) { + List partitionTypes = partitionTypes(partitionFields, idToPrimitiveTypeMapping); + StructLike partitionStruct = structLike.get(columnNameToPosition.get(PARTITION_COLUMN_NAME), org.apache.iceberg.PartitionData.class); + List partitionColumnTypes = partitionColumnType.orElseThrow().rowType().getFields().stream() + .map(RowType.Field::getType) + .collect(toImmutableList()); + if (partitionStruct != null) { + columns.add(buildRowValue( + partitionColumnType.get().rowType(), + fields -> { + for (int i = 0; i < partitionColumnTypes.size(); i++) { + Type type = partitionTypes.get(i); + io.trino.spi.type.Type trinoType = partitionColumnType.get().rowType().getFields().get(i).getType(); + Object value = null; + Integer fieldId = partitionColumnType.get().fieldIds().get(i); + if (fieldId != null) { + value = convertIcebergValueToTrino(type, partitionStruct.get(i, type.typeId().javaClass())); + } + writeNativeValue(trinoType, fields.get(i), value); + } + })); + } + } + + columns.add(structLike.get(columnNameToPosition.get(RECORD_COUNT_COLUMN_NAME), Long.class)); + columns.add(structLike.get(columnNameToPosition.get(FILE_SIZE_IN_BYTES_COLUMN_NAME), Long.class)); + columns.add(getIntegerBigintSqlMap(structLike.get(columnNameToPosition.get(COLUMN_SIZES_COLUMN_NAME), Map.class))); + columns.add(getIntegerBigintSqlMap(structLike.get(columnNameToPosition.get(VALUE_COUNTS_COLUMN_NAME), Map.class))); + columns.add(getIntegerBigintSqlMap(structLike.get(columnNameToPosition.get(NULL_VALUE_COUNTS_COLUMN_NAME), Map.class))); + columns.add(getIntegerBigintSqlMap(structLike.get(columnNameToPosition.get(NAN_VALUE_COUNTS_COLUMN_NAME), Map.class))); + columns.add(getIntegerVarcharSqlMap(structLike.get(columnNameToPosition.get(LOWER_BOUNDS_COLUMN_NAME), Map.class))); + columns.add(getIntegerVarcharSqlMap(structLike.get(columnNameToPosition.get(UPPER_BOUNDS_COLUMN_NAME), Map.class))); + columns.add(toVarbinarySlice(structLike.get(columnNameToPosition.get(KEY_METADATA_COLUMN_NAME), ByteBuffer.class))); + columns.add(toBigintArrayBlock(structLike.get(columnNameToPosition.get(SPLIT_OFFSETS_COLUMN_NAME), List.class))); + columns.add(toIntegerArrayBlock(structLike.get(columnNameToPosition.get(EQUALITY_IDS_COLUMN_NAME), List.class))); + columns.add(structLike.get(columnNameToPosition.get(SORT_ORDER_ID_COLUMN_NAME), Integer.class)); + + ReadableMetricsStruct readableMetrics = structLike.get(columnNameToPosition.get(READABLE_METRICS_COLUMN_NAME), ReadableMetricsStruct.class); + columns.add(toJson(readableMetrics, primitiveFields)); + checkArgument(columns.size() == types.size(), "Expected %s types in row, but got %s values", types.size(), columns.size()); return columns; } - private Object getIntegerBigintMapBlock(Map value) + private Block getIntegerBigintSqlMap(Map value) { if (value == null) { return null; } - return toIntegerBigintMapBlock(value); + return toIntegerBigintSqlMap(value); } - private Object getIntegerVarcharMapBlock(Map value) + private Block getIntegerVarcharSqlMap(Map value) { if (value == null) { return null; } - return toIntegerVarcharMapBlock( + return toIntegerVarcharSqlMap( value.entrySet().stream() .filter(entry -> idToTypeMapping.containsKey(entry.getKey())) .collect(toImmutableMap( @@ -238,7 +383,7 @@ private Object getIntegerVarcharMapBlock(Map value) idToTypeMapping.get(entry.getKey()), Conversions.fromByteBuffer(idToTypeMapping.get(entry.getKey()), entry.getValue()))))); } - private Object toIntegerBigintMapBlock(Map values) + private Block toIntegerBigintSqlMap(Map values) { return buildMapValue( integerToBigintMapType, @@ -249,7 +394,7 @@ private Object toIntegerBigintMapBlock(Map values) })); } - private Object toIntegerVarcharMapBlock(Map values) + private Block toIntegerVarcharSqlMap(Map values) { return buildMapValue( integerToVarcharMapType, @@ -266,7 +411,7 @@ private static Block toIntegerArrayBlock(List values) if (values == null) { return null; } - BlockBuilder builder = INTEGER.createBlockBuilder(null, values.size()); + BlockBuilder builder = INTEGER.createFixedSizeBlockBuilder(values.size()); values.forEach(value -> INTEGER.writeLong(builder, value)); return builder.build(); } @@ -277,7 +422,7 @@ private static Block toBigintArrayBlock(List values) if (values == null) { return null; } - BlockBuilder builder = BIGINT.createBlockBuilder(null, values.size()); + BlockBuilder builder = BIGINT.createFixedSizeBlockBuilder(values.size()); values.forEach(value -> BIGINT.writeLong(builder, value)); return builder.build(); } @@ -292,7 +437,75 @@ private static Slice toVarbinarySlice(ByteBuffer value) } } - private static Map getIcebergIdToTypeMapping(Schema schema) + static String toJson(ReadableMetricsStruct readableMetrics, List primitiveFields) + { + StringWriter writer = new StringWriter(); + try { + JsonGenerator generator = JSON_FACTORY.createGenerator(writer); + generator.writeStartObject(); + + for (int i = 0; i < readableMetrics.size(); i++) { + Types.NestedField field = primitiveFields.get(i); + generator.writeFieldName(field.name()); + + generator.writeStartObject(); + ReadableColMetricsStruct columnMetrics = readableMetrics.get(i, ReadableColMetricsStruct.class); + + generator.writeFieldName("column_size"); + Long columnSize = columnMetrics.get(0, Long.class); + if (columnSize == null) { + generator.writeNull(); + } + else { + generator.writeNumber(columnSize); + } + + generator.writeFieldName("value_count"); + Long valueCount = columnMetrics.get(1, Long.class); + if (valueCount == null) { + generator.writeNull(); + } + else { + generator.writeNumber(valueCount); + } + + generator.writeFieldName("null_value_count"); + Long nullValueCount = columnMetrics.get(2, Long.class); + if (nullValueCount == null) { + generator.writeNull(); + } + else { + generator.writeNumber(nullValueCount); + } + + generator.writeFieldName("nan_value_count"); + Long nanValueCount = columnMetrics.get(3, Long.class); + if (nanValueCount == null) { + generator.writeNull(); + } + else { + generator.writeNumber(nanValueCount); + } + + generator.writeFieldName("lower_bound"); + SingleValueParser.toJson(field.type(), columnMetrics.get(4, Object.class), generator); + + generator.writeFieldName("upper_bound"); + SingleValueParser.toJson(field.type(), columnMetrics.get(5, Object.class), generator); + + generator.writeEndObject(); + } + + generator.writeEndObject(); + generator.flush(); + return writer.toString(); + } + catch (IOException e) { + throw new UncheckedIOException("JSON conversion failed for: " + readableMetrics, e); + } + } + + static Map getIcebergIdToTypeMapping(Schema schema) { ImmutableMap.Builder icebergIdToTypeMapping = ImmutableMap.builder(); for (Types.NestedField field : schema.columns()) { diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/AsyncIcebergSplitProducer.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/ForIcebergMetadata.java similarity index 82% rename from plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/AsyncIcebergSplitProducer.java rename to plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/ForIcebergMetadata.java index 678ad872d8fd..6a4e26046389 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/AsyncIcebergSplitProducer.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/ForIcebergMetadata.java @@ -11,6 +11,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ + package io.trino.plugin.iceberg; import com.google.inject.BindingAnnotation; @@ -18,10 +19,12 @@ import java.lang.annotation.Retention; import java.lang.annotation.Target; +import static java.lang.annotation.ElementType.FIELD; +import static java.lang.annotation.ElementType.METHOD; import static java.lang.annotation.ElementType.PARAMETER; import static java.lang.annotation.RetentionPolicy.RUNTIME; @Retention(RUNTIME) -@Target(PARAMETER) +@Target({FIELD, PARAMETER, METHOD}) @BindingAnnotation -public @interface AsyncIcebergSplitProducer {} +public @interface ForIcebergMetadata {} diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/ForIcebergScanPlanning.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/ForIcebergScanPlanning.java new file mode 100644 index 000000000000..02df2e2a9b7d --- /dev/null +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/ForIcebergScanPlanning.java @@ -0,0 +1,29 @@ +/* + * 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.iceberg; + +import com.google.inject.BindingAnnotation; + +import java.lang.annotation.Retention; +import java.lang.annotation.Target; + +import static java.lang.annotation.ElementType.FIELD; +import static java.lang.annotation.ElementType.METHOD; +import static java.lang.annotation.ElementType.PARAMETER; +import static java.lang.annotation.RetentionPolicy.RUNTIME; + +@Retention(RUNTIME) +@Target({FIELD, PARAMETER, METHOD}) +@BindingAnnotation +public @interface ForIcebergScanPlanning {} diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/ForIcebergSplitManager.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/ForIcebergSplitManager.java new file mode 100644 index 000000000000..f11f112791c9 --- /dev/null +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/ForIcebergSplitManager.java @@ -0,0 +1,29 @@ +/* + * 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.iceberg; + +import com.google.inject.BindingAnnotation; + +import java.lang.annotation.Retention; +import java.lang.annotation.Target; + +import static java.lang.annotation.ElementType.FIELD; +import static java.lang.annotation.ElementType.METHOD; +import static java.lang.annotation.ElementType.PARAMETER; +import static java.lang.annotation.RetentionPolicy.RUNTIME; + +@Retention(RUNTIME) +@Target({FIELD, PARAMETER, METHOD}) +@BindingAnnotation +public @interface ForIcebergSplitManager {} diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergAvroFileWriter.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergAvroFileWriter.java index f0d54a25fd0f..d42b11289e42 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergAvroFileWriter.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergAvroFileWriter.java @@ -18,7 +18,6 @@ import io.trino.spi.Page; import io.trino.spi.TrinoException; import io.trino.spi.type.Type; -import org.apache.iceberg.Metrics; import org.apache.iceberg.Schema; import org.apache.iceberg.avro.Avro; import org.apache.iceberg.data.Record; @@ -29,6 +28,7 @@ import java.io.Closeable; import java.io.IOException; import java.util.List; +import java.util.Optional; import static io.airlift.slice.SizeOf.instanceSize; import static io.trino.plugin.iceberg.IcebergAvroDataConversion.toIcebergRecords; @@ -151,8 +151,8 @@ private static String toIcebergAvroCompressionName(HiveCompressionCodec hiveComp } @Override - public Metrics getMetrics() + public FileMetrics getFileMetrics() { - return avroWriter.metrics(); + return new FileMetrics(avroWriter.metrics(), Optional.empty()); } } diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergAvroPageSource.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergAvroPageSource.java index 65fbad86c794..bf3d99f2d9a4 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergAvroPageSource.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergAvroPageSource.java @@ -49,11 +49,7 @@ public class IcebergAvroPageSource private final List columnNames; private final List columnTypes; private final Map icebergTypes; - /** - * Indicates whether the column at each index should be populated with the - * indices of its rows - */ - private final List rowIndexLocations; + private final boolean appendRowNumberColumn; private final PageBuilder pageBuilder; private final AggregatedMemoryContext memoryUsage; @@ -69,16 +65,16 @@ public IcebergAvroPageSource( Optional nameMapping, List columnNames, List columnTypes, - List rowIndexLocations, + boolean appendRowNumberColumn, AggregatedMemoryContext memoryUsage) { this.columnNames = ImmutableList.copyOf(requireNonNull(columnNames, "columnNames is null")); this.columnTypes = ImmutableList.copyOf(requireNonNull(columnTypes, "columnTypes is null")); - this.rowIndexLocations = ImmutableList.copyOf(requireNonNull(rowIndexLocations, "rowIndexLocations is null")); + this.appendRowNumberColumn = appendRowNumberColumn; this.memoryUsage = requireNonNull(memoryUsage, "memoryUsage is null"); checkArgument( - columnNames.size() == rowIndexLocations.size() && columnNames.size() == columnTypes.size(), - "names, rowIndexLocations, and types must correspond one-to-one-to-one"); + columnNames.size() == columnTypes.size(), + "names and types must correspond one-to-one-to-one"); // The column orders in the generated schema might be different from the original order Schema readSchema = fileSchema.select(columnNames); @@ -94,11 +90,6 @@ public IcebergAvroPageSource( recordIterator = avroReader.iterator(); } - private boolean isIndexColumn(int column) - { - return rowIndexLocations.get(column); - } - @Override public long getCompletedBytes() { @@ -131,13 +122,11 @@ public Page getNextPage() pageBuilder.declarePosition(); Record record = recordIterator.next(); for (int channel = 0; channel < columnTypes.size(); channel++) { - if (isIndexColumn(channel)) { - BIGINT.writeLong(pageBuilder.getBlockBuilder(channel), rowId); - } - else { - String name = columnNames.get(channel); - serializeToTrinoBlock(columnTypes.get(channel), icebergTypes.get(name), pageBuilder.getBlockBuilder(channel), record.getField(name)); - } + String name = columnNames.get(channel); + serializeToTrinoBlock(columnTypes.get(channel), icebergTypes.get(name), pageBuilder.getBlockBuilder(channel), record.getField(name)); + } + if (appendRowNumberColumn) { + BIGINT.writeLong(pageBuilder.getBlockBuilder(columnTypes.size()), rowId); } rowId++; } diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergBucketFunction.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergBucketFunction.java index 433466c966a0..cb7883376380 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergBucketFunction.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergBucketFunction.java @@ -13,133 +13,140 @@ */ package io.trino.plugin.iceberg; -import io.trino.plugin.iceberg.PartitionTransforms.ColumnTransform; import io.trino.plugin.iceberg.PartitionTransforms.ValueTransform; import io.trino.spi.Page; import io.trino.spi.block.Block; +import io.trino.spi.block.RowBlock; import io.trino.spi.connector.BucketFunction; -import io.trino.spi.type.Type; +import io.trino.spi.connector.ConnectorSplit; import io.trino.spi.type.TypeOperators; -import org.apache.iceberg.PartitionSpec; import java.lang.invoke.MethodHandle; -import java.util.HashMap; import java.util.List; -import java.util.Map; +import java.util.function.ToIntFunction; import static com.google.common.base.Preconditions.checkArgument; -import static com.google.common.base.Preconditions.checkState; import static com.google.common.collect.ImmutableList.toImmutableList; -import static io.trino.plugin.iceberg.PartitionTransforms.getColumnTransform; +import static io.trino.plugin.iceberg.IcebergPartitionFunction.Transform.BUCKET; import static io.trino.spi.function.InvocationConvention.InvocationArgumentConvention.NEVER_NULL; import static io.trino.spi.function.InvocationConvention.InvocationReturnConvention.FAIL_ON_NULL; import static io.trino.spi.function.InvocationConvention.simpleConvention; import static io.trino.spi.type.TypeUtils.NULL_HASH_CODE; import static java.util.Objects.requireNonNull; +import static java.util.Objects.requireNonNullElse; public class IcebergBucketFunction - implements BucketFunction + implements BucketFunction, ToIntFunction { private final int bucketCount; + private final List functions; - private final List partitionColumns; - private final List hashCodeInvokers; + private final boolean singleBucketFunction; - public IcebergBucketFunction( - TypeOperators typeOperators, - PartitionSpec partitionSpec, - List partitioningColumns, - int bucketCount) + public IcebergBucketFunction(IcebergPartitioningHandle partitioningHandle, TypeOperators typeOperators, int bucketCount) { - requireNonNull(partitionSpec, "partitionSpec is null"); - checkArgument(!partitionSpec.isUnpartitioned(), "empty partitionSpec"); - requireNonNull(partitioningColumns, "partitioningColumns is null"); + requireNonNull(partitioningHandle, "partitioningHandle is null"); requireNonNull(typeOperators, "typeOperators is null"); checkArgument(bucketCount > 0, "Invalid bucketCount: %s", bucketCount); this.bucketCount = bucketCount; - - Map fieldIdToInputChannel = new HashMap<>(); - for (int i = 0; i < partitioningColumns.size(); i++) { - Integer previous = fieldIdToInputChannel.put(partitioningColumns.get(i).getId(), i); - checkState(previous == null, "Duplicate id %s in %s at %s and %s", partitioningColumns.get(i).getId(), partitioningColumns, i, previous); - } - partitionColumns = partitionSpec.fields().stream() - .map(field -> { - Integer channel = fieldIdToInputChannel.get(field.sourceId()); - checkArgument(channel != null, "partition field not found: %s", field); - Type inputType = partitioningColumns.get(channel).getType(); - ColumnTransform transform = getColumnTransform(field, inputType); - return new PartitionColumn(channel, transform.getValueTransform(), transform.getType()); - }) - .collect(toImmutableList()); - hashCodeInvokers = partitionColumns.stream() - .map(PartitionColumn::getResultType) - .map(type -> typeOperators.getHashCodeOperator(type, simpleConvention(FAIL_ON_NULL, NEVER_NULL))) + List partitionFunctions = partitioningHandle.partitionFunctions(); + this.functions = partitionFunctions.stream() + .map(partitionFunction -> HashFunction.create(partitionFunction, typeOperators)) .collect(toImmutableList()); + + this.singleBucketFunction = partitionFunctions.size() == 1 && + partitionFunctions.get(0).transform() == BUCKET && + partitionFunctions.get(0).size().orElseThrow() == bucketCount; } @Override public int getBucket(Page page, int position) { - long hash = 0; + if (singleBucketFunction) { + long bucket = (long) requireNonNullElse(functions.get(0).getValue(page, position), 0L); + checkArgument(0 <= bucket && bucket < bucketCount, "Bucket value out of range: %s (bucketCount: %s)", bucket, bucketCount); + return (int) bucket; + } - for (int i = 0; i < partitionColumns.size(); i++) { - PartitionColumn partitionColumn = partitionColumns.get(i); - Block block = page.getBlock(partitionColumn.getSourceChannel()); - Object value = partitionColumn.getValueTransform().apply(block, position); - long valueHash = hashValue(hashCodeInvokers.get(i), value); + long hash = 0; + for (HashFunction function : functions) { + long valueHash = function.computeHash(page, position); hash = (31 * hash) + valueHash; } return (int) ((hash & Long.MAX_VALUE) % bucketCount); } - private static long hashValue(MethodHandle method, Object value) + @Override + public int applyAsInt(ConnectorSplit split) { - if (value == null) { - return NULL_HASH_CODE; - } - try { - return (long) method.invoke(value); + List partitionValues = ((IcebergSplit) split).getPartitionValues() + .orElseThrow(() -> new IllegalArgumentException("Split does not contain partition values")); + + if (singleBucketFunction) { + long bucket = (long) requireNonNullElse(partitionValues.get(0), 0); + checkArgument(0 <= bucket && bucket < bucketCount, "Bucket value out of range: %s (bucketCount: %s)", bucket, bucketCount); + return (int) bucket; } - catch (Throwable throwable) { - if (throwable instanceof Error) { - throw (Error) throwable; - } - if (throwable instanceof RuntimeException) { - throw (RuntimeException) throwable; - } - throw new RuntimeException(throwable); + + long hash = 0; + for (int i = 0; i < functions.size(); i++) { + long valueHash = functions.get(i).computeHash(partitionValues.get(i)); + hash = (31 * hash) + valueHash; } + + return (int) ((hash & Long.MAX_VALUE) % bucketCount); } - private static class PartitionColumn + private record HashFunction(List dataPath, ValueTransform valueTransform, MethodHandle hashCodeOperator) { - private final int sourceChannel; - private final ValueTransform valueTransform; - private final Type resultType; + private static HashFunction create(IcebergPartitionFunction partitionFunction, TypeOperators typeOperators) + { + PartitionTransforms.ColumnTransform columnTransform = PartitionTransforms.getColumnTransform(partitionFunction); + return new HashFunction( + partitionFunction.dataPath(), + columnTransform.getValueTransform(), + typeOperators.getHashCodeOperator(columnTransform.getType(), simpleConvention(FAIL_ON_NULL, NEVER_NULL))); + } - public PartitionColumn(int sourceChannel, ValueTransform valueTransform, Type resultType) + private HashFunction { - this.sourceChannel = sourceChannel; - this.valueTransform = requireNonNull(valueTransform, "valueTransform is null"); - this.resultType = requireNonNull(resultType, "resultType is null"); + requireNonNull(valueTransform, "valueTransform is null"); + requireNonNull(hashCodeOperator, "hashCodeOperator is null"); } - public int getSourceChannel() + public Object getValue(Page page, int position) { - return sourceChannel; + Block block = page.getBlock(dataPath.get(0)); + for (int i = 1; i < dataPath.size(); i++) { + block = ((RowBlock) block).getFieldBlock(dataPath.get(i)); + } + return valueTransform.apply(block, position); } - public Type getResultType() + public long computeHash(Page page, int position) { - return resultType; + return computeHash(getValue(page, position)); } - public ValueTransform getValueTransform() + private long computeHash(Object value) { - return valueTransform; + if (value == null) { + return NULL_HASH_CODE; + } + try { + return (long) hashCodeOperator.invoke(value); + } + catch (Throwable throwable) { + if (throwable instanceof Error) { + throw (Error) throwable; + } + if (throwable instanceof RuntimeException) { + throw (RuntimeException) throwable; + } + throw new RuntimeException(throwable); + } } } } diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergColumnHandle.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergColumnHandle.java index ebb49d12a00f..f7c875b4b7ba 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergColumnHandle.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergColumnHandle.java @@ -18,6 +18,7 @@ import com.fasterxml.jackson.annotation.JsonProperty; import com.google.common.collect.ImmutableList; import com.google.common.collect.Iterables; +import io.airlift.slice.SizeOf; import io.trino.spi.connector.ColumnHandle; import io.trino.spi.connector.ColumnMetadata; import io.trino.spi.type.Type; @@ -26,8 +27,12 @@ import java.util.Objects; import java.util.Optional; +import static io.airlift.slice.SizeOf.estimatedSizeOf; +import static io.airlift.slice.SizeOf.instanceSize; +import static io.airlift.slice.SizeOf.sizeOf; import static io.trino.plugin.iceberg.IcebergMetadataColumn.FILE_MODIFIED_TIME; import static io.trino.plugin.iceberg.IcebergMetadataColumn.FILE_PATH; +import static io.trino.plugin.iceberg.IcebergMetadataColumn.PARTITION; import static java.util.Objects.requireNonNull; import static org.apache.iceberg.MetadataColumns.IS_DELETED; import static org.apache.iceberg.MetadataColumns.ROW_POSITION; @@ -35,19 +40,30 @@ public class IcebergColumnHandle implements ColumnHandle { + private static final int INSTANCE_SIZE = instanceSize(IcebergColumnHandle.class); + // Iceberg reserved row ids begin at INTEGER.MAX_VALUE and count down. Starting with MIN_VALUE here to avoid conflicts. - public static final int TRINO_UPDATE_ROW_ID = Integer.MIN_VALUE; - public static final int TRINO_MERGE_ROW_ID = Integer.MIN_VALUE + 1; + public static final int TRINO_MERGE_ROW_ID = Integer.MIN_VALUE; public static final String TRINO_ROW_ID_NAME = "$row_id"; - public static final int TRINO_MERGE_PARTITION_SPEC_ID = Integer.MIN_VALUE + 2; - public static final int TRINO_MERGE_PARTITION_DATA = Integer.MIN_VALUE + 3; + public static final int TRINO_MERGE_PARTITION_SPEC_ID = Integer.MIN_VALUE + 1; + public static final int TRINO_MERGE_PARTITION_DATA = Integer.MIN_VALUE + 2; + + public static final String DATA_CHANGE_TYPE_NAME = "_change_type"; + public static final int DATA_CHANGE_TYPE_ID = Integer.MIN_VALUE + 3; + public static final String DATA_CHANGE_VERSION_NAME = "_change_version_id"; + public static final int DATA_CHANGE_VERSION_ID = Integer.MIN_VALUE + 4; + public static final String DATA_CHANGE_TIMESTAMP_NAME = "_change_timestamp"; + public static final int DATA_CHANGE_TIMESTAMP_ID = Integer.MIN_VALUE + 5; + public static final String DATA_CHANGE_ORDINAL_NAME = "_change_ordinal"; + public static final int DATA_CHANGE_ORDINAL_ID = Integer.MIN_VALUE + 6; private final ColumnIdentity baseColumnIdentity; private final Type baseType; // The list of field ids to indicate the projected part of the top-level column represented by baseColumnIdentity private final List path; private final Type type; + private final boolean nullable; private final Optional comment; // Cache of ColumnIdentity#getId to ensure quick access, even with dereferences private final int id; @@ -58,12 +74,14 @@ public IcebergColumnHandle( @JsonProperty("baseType") Type baseType, @JsonProperty("path") List path, @JsonProperty("type") Type type, + @JsonProperty("nullable") boolean nullable, @JsonProperty("comment") Optional comment) { this.baseColumnIdentity = requireNonNull(baseColumnIdentity, "baseColumnIdentity is null"); this.baseType = requireNonNull(baseType, "baseType is null"); this.path = ImmutableList.copyOf(requireNonNull(path, "path is null")); this.type = requireNonNull(type, "type is null"); + this.nullable = nullable; this.comment = requireNonNull(comment, "comment is null"); this.id = path.isEmpty() ? baseColumnIdentity.getId() : Iterables.getLast(path); } @@ -99,7 +117,13 @@ public Type getBaseType() @JsonIgnore public IcebergColumnHandle getBaseColumn() { - return new IcebergColumnHandle(getBaseColumnIdentity(), getBaseType(), ImmutableList.of(), getBaseType(), Optional.empty()); + return new IcebergColumnHandle(getBaseColumnIdentity(), getBaseType(), ImmutableList.of(), getBaseType(), isNullable(), Optional.empty()); + } + + @JsonProperty + public boolean isNullable() + { + return nullable; } @JsonProperty @@ -158,12 +182,6 @@ public boolean isRowPositionColumn() return id == ROW_POSITION.fieldId(); } - @JsonIgnore - public boolean isUpdateRowIdColumn() - { - return id == TRINO_UPDATE_ROW_ID; - } - @JsonIgnore public boolean isMergeRowIdColumn() { @@ -179,6 +197,12 @@ public boolean isIsDeletedColumn() return id == IS_DELETED.fieldId(); } + @JsonIgnore + public boolean isPartitionColumn() + { + return id == PARTITION.getId(); + } + @JsonIgnore public boolean isFileModifiedTimeColumn() { @@ -188,7 +212,7 @@ public boolean isFileModifiedTimeColumn() @Override public int hashCode() { - return Objects.hash(baseColumnIdentity, baseType, path, type, comment); + return Objects.hash(baseColumnIdentity, baseType, path, type, nullable, comment); } @Override @@ -205,6 +229,7 @@ public boolean equals(Object obj) Objects.equals(this.baseType, other.baseType) && Objects.equals(this.path, other.path) && Objects.equals(this.type, other.type) && + Objects.equals(this.nullable, other.nullable) && Objects.equals(this.comment, other.comment); } @@ -214,6 +239,37 @@ public String toString() return getId() + ":" + getName() + ":" + type.getDisplayName(); } + public long getRetainedSizeInBytes() + { + // type is not accounted for as the instances are cached (by TypeRegistry) and shared + return INSTANCE_SIZE + + baseColumnIdentity.getRetainedSizeInBytes() + + estimatedSizeOf(path, SizeOf::sizeOf) + + sizeOf(nullable) + + sizeOf(comment, SizeOf::estimatedSizeOf) + + sizeOf(id); + } + + public static IcebergColumnHandle partitionColumnHandle() + { + return new IcebergColumnHandle( + columnIdentity(PARTITION), + PARTITION.getType(), + ImmutableList.of(), + PARTITION.getType(), + false, + Optional.empty()); + } + + public static ColumnMetadata partitionColumnMetadata() + { + return ColumnMetadata.builder() + .setName(PARTITION.getColumnName()) + .setType(PARTITION.getType()) + .setHidden(true) + .build(); + } + public static IcebergColumnHandle pathColumnHandle() { return new IcebergColumnHandle( @@ -221,6 +277,7 @@ public static IcebergColumnHandle pathColumnHandle() FILE_PATH.getType(), ImmutableList.of(), FILE_PATH.getType(), + false, Optional.empty()); } @@ -240,6 +297,7 @@ public static IcebergColumnHandle fileModifiedTimeColumnHandle() FILE_MODIFIED_TIME.getType(), ImmutableList.of(), FILE_MODIFIED_TIME.getType(), + false, Optional.empty()); } diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergConfig.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergConfig.java index f59284388194..45ead9a86afa 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergConfig.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergConfig.java @@ -13,27 +13,36 @@ */ package io.trino.plugin.iceberg; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableSet; import io.airlift.configuration.Config; import io.airlift.configuration.ConfigDescription; import io.airlift.configuration.DefunctConfig; import io.airlift.configuration.LegacyConfig; import io.airlift.units.DataSize; import io.airlift.units.Duration; -import io.trino.plugin.hive.HiveCompressionCodec; +import io.trino.plugin.hive.HiveCompressionOption; +import jakarta.validation.constraints.AssertFalse; import jakarta.validation.constraints.DecimalMax; import jakarta.validation.constraints.DecimalMin; import jakarta.validation.constraints.Max; import jakarta.validation.constraints.Min; import jakarta.validation.constraints.NotNull; +import java.util.List; import java.util.Optional; +import java.util.Set; +import static com.google.common.base.Preconditions.checkArgument; +import static com.google.common.collect.ImmutableSet.toImmutableSet; import static io.airlift.units.DataSize.Unit.GIGABYTE; -import static io.trino.plugin.hive.HiveCompressionCodec.ZSTD; +import static io.airlift.units.DataSize.Unit.MEGABYTE; import static io.trino.plugin.iceberg.CatalogType.HIVE_METASTORE; import static io.trino.plugin.iceberg.IcebergFileFormat.PARQUET; +import static java.util.Locale.ENGLISH; import static java.util.concurrent.TimeUnit.DAYS; import static java.util.concurrent.TimeUnit.SECONDS; +import static org.apache.iceberg.TableProperties.COMMIT_NUM_RETRIES_DEFAULT; @DefunctConfig({ "iceberg.allow-legacy-snapshot-syntax", @@ -46,33 +55,47 @@ public class IcebergConfig public static final String EXTENDED_STATISTICS_CONFIG = "iceberg.extended-statistics.enabled"; public static final String EXTENDED_STATISTICS_DESCRIPTION = "Enable collection (ANALYZE) and use of extended statistics."; public static final String COLLECT_EXTENDED_STATISTICS_ON_WRITE_DESCRIPTION = "Collect extended statistics during writes"; - public static final String EXPIRE_SNAPSHOTS_MIN_RETENTION = "iceberg.expire_snapshots.min-retention"; - public static final String REMOVE_ORPHAN_FILES_MIN_RETENTION = "iceberg.remove_orphan_files.min-retention"; + public static final String EXPIRE_SNAPSHOTS_MIN_RETENTION = "iceberg.expire-snapshots.min-retention"; + public static final String REMOVE_ORPHAN_FILES_MIN_RETENTION = "iceberg.remove-orphan-files.min-retention"; private IcebergFileFormat fileFormat = PARQUET; - private HiveCompressionCodec compressionCodec = ZSTD; + private HiveCompressionOption compressionCodec = HiveCompressionOption.ZSTD; + private int maxCommitRetry = COMMIT_NUM_RETRIES_DEFAULT; private boolean useFileSizeFromMetadata = true; private int maxPartitionsPerWriter = 100; private boolean uniqueTableLocation = true; private CatalogType catalogType = HIVE_METASTORE; - private Duration dynamicFilteringWaitTimeout = new Duration(0, SECONDS); + private Duration dynamicFilteringWaitTimeout = new Duration(1, SECONDS); private boolean tableStatisticsEnabled = true; private boolean extendedStatisticsEnabled = true; private boolean collectExtendedStatisticsOnWrite = true; private boolean projectionPushdownEnabled = true; private boolean registerTableProcedureEnabled; + private boolean addFilesProcedureEnabled; private Optional hiveCatalogName = Optional.empty(); private int formatVersion = FORMAT_VERSION_SUPPORT_MAX; private Duration expireSnapshotsMinRetention = new Duration(7, DAYS); private Duration removeOrphanFilesMinRetention = new Duration(7, DAYS); private DataSize targetMaxFileSize = DataSize.of(1, GIGABYTE); + private DataSize idleWriterMinFileSize = DataSize.of(16, MEGABYTE); // This is meant to protect users who are misusing schema locations (by // putting schemas in locations with extraneous files), so default to false // to avoid deleting those files if Trino is unable to check. private boolean deleteSchemaLocationsFallback; private double minimumAssignedSplitWeight = 0.05; + private boolean hideMaterializedViewStorageTable = true; private Optional materializedViewsStorageSchema = Optional.empty(); private boolean sortedWritingEnabled = true; + private boolean queryPartitionFilterRequired; + private int splitManagerThreads = Runtime.getRuntime().availableProcessors() * 2; + private Set queryPartitionFilterRequiredSchemas = ImmutableSet.of(); + private List allowedExtraProperties = ImmutableList.of(); + private boolean incrementalRefreshEnabled = true; + private boolean metadataCacheEnabled = true; + private boolean objectStoreLayoutEnabled; + private int metadataParallelism = 8; + private boolean bucketExecutionEnabled = true; + private boolean fileBasedConflictDetectionEnabled = true; public CatalogType getCatalogType() { @@ -100,18 +123,32 @@ public IcebergConfig setFileFormat(IcebergFileFormat fileFormat) } @NotNull - public HiveCompressionCodec getCompressionCodec() + public HiveCompressionOption getCompressionCodec() { return compressionCodec; } @Config("iceberg.compression-codec") - public IcebergConfig setCompressionCodec(HiveCompressionCodec compressionCodec) + public IcebergConfig setCompressionCodec(HiveCompressionOption compressionCodec) { this.compressionCodec = compressionCodec; return this; } + @Min(0) + public int getMaxCommitRetry() + { + return maxCommitRetry; + } + + @Config("iceberg.max-commit-retry") + @ConfigDescription("Number of times to retry a commit before failing") + public IcebergConfig setMaxCommitRetry(int maxCommitRetry) + { + this.maxCommitRetry = maxCommitRetry; + return this; + } + @Deprecated public boolean isUseFileSizeFromMetadata() { @@ -242,6 +279,20 @@ public IcebergConfig setRegisterTableProcedureEnabled(boolean registerTableProce return this; } + public boolean isAddFilesProcedureEnabled() + { + return addFilesProcedureEnabled; + } + + @Config("iceberg.add-files-procedure.enabled") + @LegacyConfig("iceberg.add_files-procedure.enabled") + @ConfigDescription("Allow users to call the add_files procedure") + public IcebergConfig setAddFilesProcedureEnabled(boolean addFilesProcedureEnabled) + { + this.addFilesProcedureEnabled = addFilesProcedureEnabled; + return this; + } + public Optional getHiveCatalogName() { return hiveCatalogName; @@ -312,6 +363,20 @@ public IcebergConfig setTargetMaxFileSize(DataSize targetMaxFileSize) return this; } + @NotNull + public DataSize getIdleWriterMinFileSize() + { + return idleWriterMinFileSize; + } + + @Config("iceberg.idle-writer-min-file-size") + @ConfigDescription("Minimum data written by a single partition writer before it can be consider as 'idle' and could be closed by the engine") + public IcebergConfig setIdleWriterMinFileSize(DataSize idleWriterMinFileSize) + { + this.idleWriterMinFileSize = idleWriterMinFileSize; + return this; + } + public boolean isDeleteSchemaLocationsFallback() { return this.deleteSchemaLocationsFallback; @@ -334,6 +399,19 @@ public IcebergConfig setMinimumAssignedSplitWeight(double minimumAssignedSplitWe return this; } + public boolean isHideMaterializedViewStorageTable() + { + return hideMaterializedViewStorageTable; + } + + @Config("iceberg.materialized-views.hide-storage-table") + @ConfigDescription("Hide materialized view storage tables in metastore") + public IcebergConfig setHideMaterializedViewStorageTable(boolean hideMaterializedViewStorageTable) + { + this.hideMaterializedViewStorageTable = hideMaterializedViewStorageTable; + return this; + } + @DecimalMax("1") @DecimalMin(value = "0", inclusive = false) public double getMinimumAssignedSplitWeight() @@ -367,4 +445,146 @@ public IcebergConfig setSortedWritingEnabled(boolean sortedWritingEnabled) this.sortedWritingEnabled = sortedWritingEnabled; return this; } + + @Config("iceberg.query-partition-filter-required") + @ConfigDescription("Require a filter on at least one partition column") + public IcebergConfig setQueryPartitionFilterRequired(boolean queryPartitionFilterRequired) + { + this.queryPartitionFilterRequired = queryPartitionFilterRequired; + return this; + } + + public boolean isQueryPartitionFilterRequired() + { + return queryPartitionFilterRequired; + } + + public Set getQueryPartitionFilterRequiredSchemas() + { + return queryPartitionFilterRequiredSchemas; + } + + @Config("iceberg.query-partition-filter-required-schemas") + @ConfigDescription("List of schemas for which filter on partition column is enforced") + public IcebergConfig setQueryPartitionFilterRequiredSchemas(List queryPartitionFilterRequiredSchemas) + { + this.queryPartitionFilterRequiredSchemas = queryPartitionFilterRequiredSchemas.stream() + .map(value -> value.toLowerCase(ENGLISH)) + .collect(toImmutableSet()); + return this; + } + + @Min(0) + public int getSplitManagerThreads() + { + return splitManagerThreads; + } + + public List getAllowedExtraProperties() + { + return allowedExtraProperties; + } + + @Config("iceberg.allowed-extra-properties") + @ConfigDescription("List of extra properties that are allowed to be set on Iceberg tables") + public IcebergConfig setAllowedExtraProperties(List allowedExtraProperties) + { + this.allowedExtraProperties = ImmutableList.copyOf(allowedExtraProperties); + checkArgument(!allowedExtraProperties.contains("*") || allowedExtraProperties.size() == 1, + "Wildcard * should be the only element in the list"); + return this; + } + + public boolean isIncrementalRefreshEnabled() + { + return incrementalRefreshEnabled; + } + + @Config("iceberg.incremental-refresh-enabled") + @ConfigDescription("Enable Incremental refresh for MVs backed by Iceberg tables, when possible") + public IcebergConfig setIncrementalRefreshEnabled(boolean incrementalRefreshEnabled) + { + this.incrementalRefreshEnabled = incrementalRefreshEnabled; + return this; + } + + @Config("iceberg.split-manager-threads") + @ConfigDescription("Number of threads to use for generating splits") + public IcebergConfig setSplitManagerThreads(int splitManagerThreads) + { + this.splitManagerThreads = splitManagerThreads; + return this; + } + + @AssertFalse(message = "iceberg.materialized-views.storage-schema may only be set when iceberg.materialized-views.hide-storage-table is set to false") + public boolean isStorageSchemaSetWhenHidingIsEnabled() + { + return hideMaterializedViewStorageTable && materializedViewsStorageSchema.isPresent(); + } + + public boolean isMetadataCacheEnabled() + { + return metadataCacheEnabled; + } + + @Config("iceberg.metadata-cache.enabled") + @ConfigDescription("Enables in-memory caching of metadata files on coordinator if fs.cache.enabled is not set to true") + public IcebergConfig setMetadataCacheEnabled(boolean metadataCacheEnabled) + { + this.metadataCacheEnabled = metadataCacheEnabled; + return this; + } + + public boolean isObjectStoreLayoutEnabled() + { + return objectStoreLayoutEnabled; + } + + @Config("iceberg.object-store-layout.enabled") + @ConfigDescription("Enable the Iceberg object store file layout") + public IcebergConfig setObjectStoreLayoutEnabled(boolean objectStoreLayoutEnabled) + { + this.objectStoreLayoutEnabled = objectStoreLayoutEnabled; + return this; + } + + @Min(1) + public int getMetadataParallelism() + { + return metadataParallelism; + } + + @ConfigDescription("Limits metadata enumeration calls parallelism") + @Config("iceberg.metadata.parallelism") + public IcebergConfig setMetadataParallelism(int metadataParallelism) + { + this.metadataParallelism = metadataParallelism; + return this; + } + + public boolean isBucketExecutionEnabled() + { + return bucketExecutionEnabled; + } + + @Config("iceberg.bucket-execution") + @ConfigDescription("Enable bucket-aware execution: use physical bucketing information to optimize queries") + public IcebergConfig setBucketExecutionEnabled(boolean bucketExecutionEnabled) + { + this.bucketExecutionEnabled = bucketExecutionEnabled; + return this; + } + + public boolean isFileBasedConflictDetectionEnabled() + { + return fileBasedConflictDetectionEnabled; + } + + @Config("iceberg.file-based-conflict-detection") + @ConfigDescription("Enable file-based conflict detection: take partition information from the actual written files as a source for the conflict detection system") + public IcebergConfig setFileBasedConflictDetectionEnabled(boolean fileBasedConflictDetectionEnabled) + { + this.fileBasedConflictDetectionEnabled = fileBasedConflictDetectionEnabled; + return this; + } } diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergConnector.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergConnector.java index 1b9418d1b44a..3ab74f5eebab 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergConnector.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergConnector.java @@ -16,6 +16,7 @@ import com.google.common.annotations.VisibleForTesting; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableSet; +import com.google.inject.Inject; import com.google.inject.Injector; import io.airlift.bootstrap.LifeCycleManager; import io.trino.plugin.base.classloader.ClassLoaderSafeConnectorMetadata; @@ -32,6 +33,8 @@ import io.trino.spi.connector.ConnectorSplitManager; import io.trino.spi.connector.ConnectorTransactionHandle; import io.trino.spi.connector.TableProcedureMetadata; +import io.trino.spi.function.FunctionProvider; +import io.trino.spi.function.table.ConnectorTableFunction; import io.trino.spi.procedure.Procedure; import io.trino.spi.session.PropertyMetadata; import io.trino.spi.transaction.IsolationLevel; @@ -66,7 +69,10 @@ public class IcebergConnector private final Optional accessControl; private final Set procedures; private final Set tableProcedures; + private final Set tableFunctions; + private final FunctionProvider functionProvider; + @Inject public IcebergConnector( Injector injector, LifeCycleManager lifeCycleManager, @@ -76,13 +82,15 @@ public IcebergConnector( ConnectorPageSinkProvider pageSinkProvider, ConnectorNodePartitioningProvider nodePartitioningProvider, Set sessionPropertiesProviders, - List> schemaProperties, - List> tableProperties, - List> materializedViewProperties, - List> analyzeProperties, + IcebergSchemaProperties schemaProperties, + IcebergTableProperties tableProperties, + IcebergMaterializedViewProperties materializedViewProperties, + IcebergAnalyzeProperties analyzeProperties, Optional accessControl, Set procedures, - Set tableProcedures) + Set tableProcedures, + Set tableFunctions, + FunctionProvider functionProvider) { this.injector = requireNonNull(injector, "injector is null"); this.lifeCycleManager = requireNonNull(lifeCycleManager, "lifeCycleManager is null"); @@ -94,13 +102,15 @@ public IcebergConnector( this.sessionProperties = sessionPropertiesProviders.stream() .flatMap(sessionPropertiesProvider -> sessionPropertiesProvider.getSessionProperties().stream()) .collect(toImmutableList()); - this.schemaProperties = ImmutableList.copyOf(requireNonNull(schemaProperties, "schemaProperties is null")); - this.tableProperties = ImmutableList.copyOf(requireNonNull(tableProperties, "tableProperties is null")); - this.materializedViewProperties = ImmutableList.copyOf(requireNonNull(materializedViewProperties, "materializedViewProperties is null")); - this.analyzeProperties = ImmutableList.copyOf(requireNonNull(analyzeProperties, "analyzeProperties is null")); + this.schemaProperties = ImmutableList.copyOf(requireNonNull(schemaProperties, "schemaProperties is null").getSchemaProperties()); + this.tableProperties = ImmutableList.copyOf(requireNonNull(tableProperties, "tableProperties is null").getTableProperties()); + this.materializedViewProperties = ImmutableList.copyOf(requireNonNull(materializedViewProperties, "materializedViewProperties is null").getMaterializedViewProperties()); + this.analyzeProperties = ImmutableList.copyOf(requireNonNull(analyzeProperties, "analyzeProperties is null").getAnalyzeProperties()); this.accessControl = requireNonNull(accessControl, "accessControl is null"); this.procedures = ImmutableSet.copyOf(requireNonNull(procedures, "procedures is null")); this.tableProcedures = ImmutableSet.copyOf(requireNonNull(tableProcedures, "tableProcedures is null")); + this.tableFunctions = ImmutableSet.copyOf(requireNonNull(tableFunctions, "tableFunctions is null")); + this.functionProvider = requireNonNull(functionProvider, "functionProvider is null"); } @Override @@ -154,6 +164,18 @@ public Set getTableProcedures() return tableProcedures; } + @Override + public Set getTableFunctions() + { + return tableFunctions; + } + + @Override + public Optional getFunctionProvider() + { + return Optional.of(functionProvider); + } + @Override public List> getSessionProperties() { diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergConnectorFactory.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergConnectorFactory.java index 39e2c9e52617..e8fc424105cc 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergConnectorFactory.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergConnectorFactory.java @@ -14,34 +14,43 @@ package io.trino.plugin.iceberg; import com.google.inject.Binder; +import com.google.inject.Injector; +import com.google.inject.Key; import com.google.inject.Module; +import io.airlift.bootstrap.Bootstrap; +import io.airlift.configuration.AbstractConfigurationAwareModule; +import io.airlift.json.JsonModule; +import io.opentelemetry.api.OpenTelemetry; +import io.opentelemetry.api.trace.Tracer; +import io.trino.filesystem.manager.FileSystemModule; +import io.trino.plugin.base.CatalogName; +import io.trino.plugin.base.jmx.ConnectorObjectNameGeneratorModule; +import io.trino.plugin.base.jmx.MBeanServerModule; +import io.trino.plugin.hive.HiveConfig; +import io.trino.plugin.hive.NodeVersion; +import io.trino.plugin.iceberg.catalog.IcebergCatalogModule; +import io.trino.spi.NodeManager; +import io.trino.spi.PageIndexerFactory; +import io.trino.spi.PageSorter; +import io.trino.spi.classloader.ThreadContextClassLoader; +import io.trino.spi.connector.CatalogHandle; import io.trino.spi.connector.Connector; import io.trino.spi.connector.ConnectorContext; import io.trino.spi.connector.ConnectorFactory; +import io.trino.spi.type.TypeManager; +import org.weakref.jmx.guice.MBeanModule; -import java.lang.reflect.InvocationTargetException; import java.util.Map; import java.util.Optional; -import static com.google.common.base.Throwables.throwIfUnchecked; +import static com.google.common.base.Verify.verify; +import static com.google.inject.util.Modules.EMPTY_MODULE; import static io.trino.plugin.base.Versions.checkStrictSpiVersionMatch; import static java.util.Objects.requireNonNull; public class IcebergConnectorFactory implements ConnectorFactory { - private final Class module; - - public IcebergConnectorFactory() - { - this(EmptyModule.class); - } - - public IcebergConnectorFactory(Class module) - { - this.module = requireNonNull(module, "module is null"); - } - @Override public String getName() { @@ -52,29 +61,71 @@ public String getName() public Connector create(String catalogName, Map config, ConnectorContext context) { checkStrictSpiVersionMatch(context, this); + return createConnector(catalogName, config, context, EMPTY_MODULE, Optional.empty()); + } - ClassLoader classLoader = context.duplicatePluginClassLoader(); - try { - Object moduleInstance = classLoader.loadClass(module.getName()).getConstructor().newInstance(); - Class moduleClass = classLoader.loadClass(Module.class.getName()); - return (Connector) classLoader.loadClass(InternalIcebergConnectorFactory.class.getName()) - .getMethod("createConnector", String.class, Map.class, ConnectorContext.class, moduleClass, Optional.class, Optional.class) - .invoke(null, catalogName, config, context, moduleInstance, Optional.empty(), Optional.empty()); - } - catch (InvocationTargetException e) { - Throwable targetException = e.getTargetException(); - throwIfUnchecked(targetException); - throw new RuntimeException(targetException); - } - catch (ReflectiveOperationException e) { - throw new RuntimeException(e); + public static Connector createConnector( + String catalogName, + Map config, + ConnectorContext context, + Module module, + Optional icebergCatalogModule) + { + ClassLoader classLoader = IcebergConnectorFactory.class.getClassLoader(); + try (ThreadContextClassLoader ignore = new ThreadContextClassLoader(classLoader)) { + Bootstrap app = new Bootstrap( + new MBeanModule(), + new ConnectorObjectNameGeneratorModule("io.trino.plugin.iceberg", "trino.plugin.iceberg"), + new JsonModule(), + new IcebergModule(), + new IcebergSecurityModule(), + icebergCatalogModule.orElse(new IcebergCatalogModule()), + new MBeanServerModule(), + new IcebergFileSystemModule(catalogName, context), + binder -> { + binder.bind(ClassLoader.class).toInstance(IcebergConnectorFactory.class.getClassLoader()); + binder.bind(OpenTelemetry.class).toInstance(context.getOpenTelemetry()); + binder.bind(Tracer.class).toInstance(context.getTracer()); + binder.bind(NodeVersion.class).toInstance(new NodeVersion(context.getNodeManager().getCurrentNode().getVersion())); + binder.bind(NodeManager.class).toInstance(context.getNodeManager()); + binder.bind(TypeManager.class).toInstance(context.getTypeManager()); + binder.bind(PageIndexerFactory.class).toInstance(context.getPageIndexerFactory()); + binder.bind(CatalogHandle.class).toInstance(context.getCatalogHandle()); + binder.bind(CatalogName.class).toInstance(new CatalogName(catalogName)); + binder.bind(PageSorter.class).toInstance(context.getPageSorter()); + }, + module); + + Injector injector = app + .doNotInitializeLogging() + .setRequiredConfigurationProperties(config) + .initialize(); + + verify(!injector.getBindings().containsKey(Key.get(HiveConfig.class)), "HiveConfig should not be bound"); + + return injector.getInstance(IcebergConnector.class); } } - public static class EmptyModule - implements Module + private static class IcebergFileSystemModule + extends AbstractConfigurationAwareModule { + private final String catalogName; + private final NodeManager nodeManager; + private final OpenTelemetry openTelemetry; + + public IcebergFileSystemModule(String catalogName, ConnectorContext context) + { + this.catalogName = requireNonNull(catalogName, "catalogName is null"); + this.nodeManager = context.getNodeManager(); + this.openTelemetry = context.getOpenTelemetry(); + } + @Override - public void configure(Binder binder) {} + protected void setup(Binder binder) + { + boolean metadataCacheEnabled = buildConfigObject(IcebergConfig.class).isMetadataCacheEnabled(); + install(new FileSystemModule(catalogName, nodeManager, openTelemetry, metadataCacheEnabled)); + } } } diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergEnvironmentContext.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergEnvironmentContext.java new file mode 100644 index 000000000000..50f34c228530 --- /dev/null +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergEnvironmentContext.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.iceberg; + +import com.google.inject.Inject; +import io.trino.spi.NodeManager; +import org.apache.iceberg.EnvironmentContext; + +import static java.util.Objects.requireNonNull; +import static org.apache.iceberg.EnvironmentContext.ENGINE_NAME; +import static org.apache.iceberg.EnvironmentContext.ENGINE_VERSION; + +public class IcebergEnvironmentContext +{ + @Inject + public IcebergEnvironmentContext(NodeManager nodeManager) + { + requireNonNull(nodeManager, "nodeManager is null"); + EnvironmentContext.put(ENGINE_NAME, "trino"); + EnvironmentContext.put(ENGINE_VERSION, nodeManager.getCurrentNode().getVersion()); + } +} diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergErrorCode.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergErrorCode.java index 0ce831abb7e7..edb866436aa5 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergErrorCode.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergErrorCode.java @@ -40,6 +40,8 @@ public enum IcebergErrorCode ICEBERG_CATALOG_ERROR(13, EXTERNAL), ICEBERG_WRITER_CLOSE_ERROR(14, EXTERNAL), ICEBERG_MISSING_METADATA(15, EXTERNAL), + ICEBERG_WRITER_DATA_ERROR(16, EXTERNAL), + ICEBERG_UNSUPPORTED_VIEW_DIALECT(17, EXTERNAL) /**/; private final ErrorCode errorCode; diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergExceptions.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergExceptions.java new file mode 100644 index 000000000000..0fbcfe003346 --- /dev/null +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergExceptions.java @@ -0,0 +1,56 @@ +/* + * 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.iceberg; + +import io.trino.spi.StandardErrorCode; +import io.trino.spi.TrinoException; +import org.apache.iceberg.exceptions.ValidationException; + +import java.io.FileNotFoundException; + +import static com.google.common.base.Throwables.getCausalChain; +import static io.trino.plugin.iceberg.IcebergErrorCode.ICEBERG_INVALID_METADATA; +import static io.trino.plugin.iceberg.IcebergErrorCode.ICEBERG_MISSING_METADATA; + +public final class IcebergExceptions +{ + private IcebergExceptions() {} + + private static boolean isNotFoundException(Throwable failure) + { + return getCausalChain(failure).stream().anyMatch(e -> + e instanceof org.apache.iceberg.exceptions.NotFoundException + || e instanceof FileNotFoundException); + } + + public static boolean isFatalException(Throwable failure) + { + return isNotFoundException(failure) || failure instanceof ValidationException; + } + + public static RuntimeException translateMetadataException(Throwable failure, String tableName) + { + if (failure instanceof TrinoException trinoException) { + return trinoException; + } + if (isNotFoundException(failure)) { + throw new TrinoException(ICEBERG_MISSING_METADATA, "Metadata not found in metadata location for table " + tableName, failure); + } + if (failure instanceof ValidationException) { + throw new TrinoException(ICEBERG_INVALID_METADATA, "Invalid metadata file for table " + tableName, failure); + } + + return new TrinoException(StandardErrorCode.GENERIC_INTERNAL_ERROR, "Error processing metadata for table " + tableName, failure); + } +} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/aws/athena/projection/ProjectionFactory.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergFileSystemFactory.java similarity index 68% rename from plugin/trino-hive/src/main/java/io/trino/plugin/hive/aws/athena/projection/ProjectionFactory.java rename to plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergFileSystemFactory.java index 88d640f4bc89..a8bd36cce40c 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/aws/athena/projection/ProjectionFactory.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergFileSystemFactory.java @@ -11,15 +11,14 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package io.trino.plugin.hive.aws.athena.projection; +package io.trino.plugin.iceberg; -import io.trino.spi.type.Type; +import io.trino.filesystem.TrinoFileSystem; +import io.trino.spi.security.ConnectorIdentity; import java.util.Map; -public interface ProjectionFactory +public interface IcebergFileSystemFactory { - boolean isSupportedColumnType(Type columnType); - - Projection create(String columnName, Type columnType, Map columnProperties); + TrinoFileSystem create(ConnectorIdentity identity, Map fileIoProperties); } diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergFileWriter.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergFileWriter.java index 0d04338e413e..f6616bab1c9e 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergFileWriter.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergFileWriter.java @@ -16,8 +16,13 @@ import io.trino.plugin.hive.FileWriter; import org.apache.iceberg.Metrics; +import java.util.List; +import java.util.Optional; + public interface IcebergFileWriter extends FileWriter { - Metrics getMetrics(); + FileMetrics getFileMetrics(); + + record FileMetrics(Metrics metrics, Optional> splitOffsets) {} } diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergFileWriterFactory.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergFileWriterFactory.java index 48f18e3e782f..57d8c716c897 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergFileWriterFactory.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergFileWriterFactory.java @@ -30,6 +30,7 @@ import io.trino.orc.OutputStreamOrcDataSink; import io.trino.parquet.writer.ParquetWriterOptions; import io.trino.plugin.hive.FileFormatDataSourceStats; +import io.trino.plugin.hive.HiveCompressionCodec; import io.trino.plugin.hive.NodeVersion; import io.trino.plugin.hive.orc.OrcWriterConfig; import io.trino.plugin.iceberg.fileio.ForwardingOutputFile; @@ -52,16 +53,16 @@ import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.collect.ImmutableList.toImmutableList; -import static io.trino.plugin.hive.HiveMetadata.PRESTO_QUERY_ID_NAME; -import static io.trino.plugin.hive.HiveMetadata.PRESTO_VERSION_NAME; +import static io.trino.plugin.hive.HiveCompressionCodecs.toCompressionCodec; +import static io.trino.plugin.hive.HiveMetadata.TRINO_QUERY_ID_NAME; +import static io.trino.plugin.hive.HiveMetadata.TRINO_VERSION_NAME; import static io.trino.plugin.iceberg.IcebergErrorCode.ICEBERG_INVALID_METADATA; import static io.trino.plugin.iceberg.IcebergErrorCode.ICEBERG_WRITER_OPEN_ERROR; import static io.trino.plugin.iceberg.IcebergErrorCode.ICEBERG_WRITE_VALIDATION_FAILED; -import static io.trino.plugin.iceberg.IcebergMetadata.ORC_BLOOM_FILTER_COLUMNS_KEY; -import static io.trino.plugin.iceberg.IcebergMetadata.ORC_BLOOM_FILTER_FPP_KEY; import static io.trino.plugin.iceberg.IcebergSessionProperties.getCompressionCodec; import static io.trino.plugin.iceberg.IcebergSessionProperties.getOrcStringStatisticsLimit; import static io.trino.plugin.iceberg.IcebergSessionProperties.getOrcWriterMaxDictionaryMemory; +import static io.trino.plugin.iceberg.IcebergSessionProperties.getOrcWriterMaxRowGroupRows; import static io.trino.plugin.iceberg.IcebergSessionProperties.getOrcWriterMaxStripeRows; import static io.trino.plugin.iceberg.IcebergSessionProperties.getOrcWriterMaxStripeSize; import static io.trino.plugin.iceberg.IcebergSessionProperties.getOrcWriterMinStripeSize; @@ -69,13 +70,16 @@ import static io.trino.plugin.iceberg.IcebergSessionProperties.getParquetWriterBatchSize; import static io.trino.plugin.iceberg.IcebergSessionProperties.getParquetWriterBlockSize; import static io.trino.plugin.iceberg.IcebergSessionProperties.getParquetWriterPageSize; +import static io.trino.plugin.iceberg.IcebergSessionProperties.getParquetWriterPageValueCount; import static io.trino.plugin.iceberg.IcebergSessionProperties.isOrcWriterValidate; -import static io.trino.plugin.iceberg.IcebergTableProperties.ORC_BLOOM_FILTER_FPP; +import static io.trino.plugin.iceberg.IcebergTableProperties.ORC_BLOOM_FILTER_FPP_PROPERTY; +import static io.trino.plugin.iceberg.IcebergUtil.getOrcBloomFilterColumns; +import static io.trino.plugin.iceberg.IcebergUtil.getOrcBloomFilterFpp; +import static io.trino.plugin.iceberg.IcebergUtil.getParquetBloomFilterColumns; import static io.trino.plugin.iceberg.TypeConverter.toOrcType; import static io.trino.plugin.iceberg.TypeConverter.toTrinoType; import static io.trino.plugin.iceberg.util.PrimitiveTypeMapBuilder.makeTypeMap; import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED; -import static java.lang.Double.parseDouble; import static java.lang.String.format; import static java.util.Objects.requireNonNull; import static org.apache.iceberg.TableProperties.DEFAULT_WRITE_METRICS_MODE; @@ -126,7 +130,7 @@ public IcebergFileWriter createDataFileWriter( switch (fileFormat) { case PARQUET: // TODO use metricsConfig https://github.com/trinodb/trino/issues/9791 - return createParquetWriter(MetricsConfig.getDefault(), fileSystem, outputPath, icebergSchema, session); + return createParquetWriter(MetricsConfig.getDefault(), fileSystem, outputPath, icebergSchema, session, storageProperties); case ORC: return createOrcWriter(metricsConfig, fileSystem, outputPath, icebergSchema, session, storageProperties, getOrcStringStatisticsLimit(session)); case AVRO: @@ -145,7 +149,7 @@ public IcebergFileWriter createPositionDeleteWriter( { switch (fileFormat) { case PARQUET: - return createParquetWriter(FULL_METRICS_CONFIG, fileSystem, outputPath, POSITION_DELETE_SCHEMA, session); + return createParquetWriter(FULL_METRICS_CONFIG, fileSystem, outputPath, POSITION_DELETE_SCHEMA, session, storageProperties); case ORC: return createOrcWriter(FULL_METRICS_CONFIG, fileSystem, outputPath, POSITION_DELETE_SCHEMA, session, storageProperties, DataSize.ofBytes(Integer.MAX_VALUE)); case AVRO: @@ -160,7 +164,8 @@ private IcebergFileWriter createParquetWriter( TrinoFileSystem fileSystem, Location outputPath, Schema icebergSchema, - ConnectorSession session) + ConnectorSession session, + Map storageProperties) { List fileColumnNames = icebergSchema.columns().stream() .map(Types.NestedField::name) @@ -176,10 +181,13 @@ private IcebergFileWriter createParquetWriter( ParquetWriterOptions parquetWriterOptions = ParquetWriterOptions.builder() .setMaxPageSize(getParquetWriterPageSize(session)) + .setMaxPageValueCount(getParquetWriterPageValueCount(session)) .setMaxBlockSize(getParquetWriterBlockSize(session)) .setBatchSize(getParquetWriterBatchSize(session)) + .setBloomFilterColumns(getParquetBloomFilterColumns(storageProperties)) .build(); + HiveCompressionCodec hiveCompressionCodec = toCompressionCodec(getCompressionCodec(session)); return new IcebergParquetFileWriter( metricsConfig, outputFile, @@ -190,9 +198,9 @@ private IcebergFileWriter createParquetWriter( makeTypeMap(fileColumnTypes, fileColumnNames), parquetWriterOptions, IntStream.range(0, fileColumnNames.size()).toArray(), - getCompressionCodec(session).getParquetCompressionCodec(), - nodeVersion.toString(), - fileSystem); + hiveCompressionCodec.getParquetCompressionCodec() + .orElseThrow(() -> new TrinoException(NOT_SUPPORTED, "Compression codec %s not supported for Parquet".formatted(hiveCompressionCodec))), + nodeVersion.toString()); } catch (IOException e) { throw new TrinoException(ICEBERG_WRITER_OPEN_ERROR, "Error creating Parquet file", e); @@ -243,17 +251,18 @@ private IcebergFileWriter createOrcWriter( fileColumnNames, fileColumnTypes, toOrcType(icebergSchema), - getCompressionCodec(session).getOrcCompressionKind(), + toCompressionCodec(getCompressionCodec(session)).getOrcCompressionKind(), withBloomFilterOptions(orcWriterOptions, storageProperties) .withStripeMinSize(getOrcWriterMinStripeSize(session)) .withStripeMaxSize(getOrcWriterMaxStripeSize(session)) .withStripeMaxRowCount(getOrcWriterMaxStripeRows(session)) + .withRowGroupMaxRowCount(getOrcWriterMaxRowGroupRows(session)) .withDictionaryMaxMemory(getOrcWriterMaxDictionaryMemory(session)) .withMaxStringStatisticsLimit(stringStatisticsLimit), IntStream.range(0, fileColumnNames.size()).toArray(), ImmutableMap.builder() - .put(PRESTO_VERSION_NAME, nodeVersion.toString()) - .put(PRESTO_QUERY_ID_NAME, session.getQueryId()) + .put(TRINO_VERSION_NAME, nodeVersion.toString()) + .put(TRINO_QUERY_ID_NAME, session.getQueryId()) .buildOrThrow(), validationInputFactory, getOrcWriterValidateMode(session), @@ -266,19 +275,18 @@ private IcebergFileWriter createOrcWriter( public static OrcWriterOptions withBloomFilterOptions(OrcWriterOptions orcWriterOptions, Map storageProperties) { - if (storageProperties.containsKey(ORC_BLOOM_FILTER_COLUMNS_KEY)) { - if (!storageProperties.containsKey(ORC_BLOOM_FILTER_FPP_KEY)) { - throw new TrinoException(ICEBERG_INVALID_METADATA, "FPP for Bloom filter is missing"); - } + Optional orcBloomFilterColumns = getOrcBloomFilterColumns(storageProperties); + Optional orcBloomFilterFpp = getOrcBloomFilterFpp(storageProperties); + if (orcBloomFilterColumns.isPresent()) { try { - double fpp = parseDouble(storageProperties.get(ORC_BLOOM_FILTER_FPP_KEY)); + double fpp = orcBloomFilterFpp.map(Double::parseDouble).orElseGet(orcWriterOptions::getBloomFilterFpp); return OrcWriterOptions.builderFrom(orcWriterOptions) - .setBloomFilterColumns(ImmutableSet.copyOf(COLUMN_NAMES_SPLITTER.splitToList(storageProperties.get(ORC_BLOOM_FILTER_COLUMNS_KEY)))) + .setBloomFilterColumns(ImmutableSet.copyOf(COLUMN_NAMES_SPLITTER.splitToList(orcBloomFilterColumns.get()))) .setBloomFilterFpp(fpp) .build(); } catch (NumberFormatException e) { - throw new TrinoException(ICEBERG_INVALID_METADATA, format("Invalid value for %s property: %s", ORC_BLOOM_FILTER_FPP, storageProperties.get(ORC_BLOOM_FILTER_FPP_KEY))); + throw new TrinoException(ICEBERG_INVALID_METADATA, format("Invalid value for %s property: %s", ORC_BLOOM_FILTER_FPP_PROPERTY, orcBloomFilterFpp.get())); } } return orcWriterOptions; @@ -297,10 +305,10 @@ private IcebergFileWriter createAvroWriter( .collect(toImmutableList()); return new IcebergAvroFileWriter( - new ForwardingOutputFile(fileSystem, outputPath.toString()), + new ForwardingOutputFile(fileSystem, outputPath), rollbackAction, icebergSchema, columnTypes, - getCompressionCodec(session)); + toCompressionCodec(getCompressionCodec(session))); } } diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergInputInfo.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergInputInfo.java index f6dac0b08dbb..d218360a6c8b 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergInputInfo.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergInputInfo.java @@ -16,6 +16,7 @@ import com.fasterxml.jackson.annotation.JsonCreator; import com.fasterxml.jackson.annotation.JsonProperty; +import java.util.List; import java.util.Objects; import java.util.Optional; @@ -24,18 +25,30 @@ public class IcebergInputInfo { private final Optional snapshotId; - private final Optional partitioned; + private final List partitionFields; private final String tableDefaultFileFormat; + private final Optional totalRecords; + private final Optional deletedRecords; + private final Optional totalDataFiles; + private final Optional totalDeleteFiles; @JsonCreator public IcebergInputInfo( @JsonProperty("snapshotId") Optional snapshotId, - @JsonProperty("partitioned") Optional partitioned, - @JsonProperty("fileFormat") String tableDefaultFileFormat) + @JsonProperty("partitionFields") List partitionFields, + @JsonProperty("fileFormat") String tableDefaultFileFormat, + @JsonProperty("totalRecords") Optional totalRecords, + @JsonProperty("deletedRecords") Optional deletedRecords, + @JsonProperty("totalDataFiles") Optional totalDataFiles, + @JsonProperty("totalDeleteFiles") Optional totalDeleteFiles) { this.snapshotId = requireNonNull(snapshotId, "snapshotId is null"); - this.partitioned = requireNonNull(partitioned, "partitioned is null"); + this.partitionFields = partitionFields; this.tableDefaultFileFormat = requireNonNull(tableDefaultFileFormat, "tableDefaultFileFormat is null"); + this.totalRecords = totalRecords; + this.deletedRecords = deletedRecords; + this.totalDataFiles = totalDataFiles; + this.totalDeleteFiles = totalDeleteFiles; } @JsonProperty @@ -45,9 +58,9 @@ public Optional getSnapshotId() } @JsonProperty - public Optional getPartitioned() + public List getPartitionFields() { - return partitioned; + return partitionFields; } @JsonProperty @@ -56,6 +69,30 @@ public String getTableDefaultFileFormat() return tableDefaultFileFormat; } + @JsonProperty + public Optional getDeletedRecords() + { + return deletedRecords; + } + + @JsonProperty + public Optional getTotalDataFiles() + { + return totalDataFiles; + } + + @JsonProperty + public Optional getTotalDeleteFiles() + { + return totalDeleteFiles; + } + + @JsonProperty + public Optional getTotalRecords() + { + return totalRecords; + } + @Override public boolean equals(Object o) { @@ -65,14 +102,18 @@ public boolean equals(Object o) if (!(o instanceof IcebergInputInfo that)) { return false; } - return partitioned.equals(that.partitioned) + return Objects.equals(partitionFields, that.partitionFields) && snapshotId.equals(that.snapshotId) + && totalRecords.equals(that.totalRecords) + && deletedRecords.equals(that.deletedRecords) + && totalDataFiles.equals(that.totalDataFiles) + && totalDeleteFiles.equals(that.totalDeleteFiles) && tableDefaultFileFormat.equals(that.tableDefaultFileFormat); } @Override public int hashCode() { - return Objects.hash(snapshotId, partitioned, tableDefaultFileFormat); + return Objects.hash(snapshotId, partitionFields, tableDefaultFileFormat, totalRecords, deletedRecords, totalDataFiles, totalDeleteFiles); } } diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergMaterializedViewAdditionalProperties.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergMaterializedViewProperties.java similarity index 65% rename from plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergMaterializedViewAdditionalProperties.java rename to plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergMaterializedViewProperties.java index 37bc14a1e1c3..036dd9d5a978 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergMaterializedViewAdditionalProperties.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergMaterializedViewProperties.java @@ -23,20 +23,24 @@ import static io.trino.spi.session.PropertyMetadata.stringProperty; -public class IcebergMaterializedViewAdditionalProperties +public class IcebergMaterializedViewProperties { public static final String STORAGE_SCHEMA = "storage_schema"; private final List> materializedViewProperties; @Inject - public IcebergMaterializedViewAdditionalProperties(IcebergConfig icebergConfig) + public IcebergMaterializedViewProperties(IcebergConfig icebergConfig, IcebergTableProperties tableProperties) { - materializedViewProperties = ImmutableList.of(stringProperty( - STORAGE_SCHEMA, - "Schema for creating materialized view storage table", - icebergConfig.getMaterializedViewsStorageSchema().orElse(null), - false)); + materializedViewProperties = ImmutableList.>builder() + .add(stringProperty( + STORAGE_SCHEMA, + "Schema for creating materialized view storage table", + icebergConfig.getMaterializedViewsStorageSchema().orElse(null), + false)) + // Materialized view should allow configuring all the supported iceberg table properties for the storage table + .addAll(tableProperties.getTableProperties()) + .build(); } public List> getMaterializedViewProperties() diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergMergeSink.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergMergeSink.java index e51f334e71d3..e1e5e4111a99 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergMergeSink.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergMergeSink.java @@ -13,15 +13,14 @@ */ package io.trino.plugin.iceberg; -import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import io.airlift.json.JsonCodec; import io.airlift.slice.Slice; import io.trino.filesystem.TrinoFileSystem; -import io.trino.plugin.iceberg.delete.IcebergPositionDeletePageSink; +import io.trino.plugin.iceberg.delete.PositionDeleteWriter; import io.trino.spi.Page; -import io.trino.spi.PageBuilder; -import io.trino.spi.block.ColumnarRow; +import io.trino.spi.block.Block; +import io.trino.spi.block.RowBlock; import io.trino.spi.connector.ConnectorMergeSink; import io.trino.spi.connector.ConnectorPageSink; import io.trino.spi.connector.ConnectorSession; @@ -44,7 +43,6 @@ import java.util.concurrent.CompletableFuture; import static io.trino.plugin.base.util.Closables.closeAllSuppress; -import static io.trino.spi.block.ColumnarRow.toColumnarRow; import static io.trino.spi.connector.MergePage.createDeleteAndInsertPages; import static io.trino.spi.type.BigintType.BIGINT; import static io.trino.spi.type.IntegerType.INTEGER; @@ -101,16 +99,19 @@ public void storeMergedRows(Page page) mergePage.getInsertionsPage().ifPresent(insertPageSink::appendPage); mergePage.getDeletionsPage().ifPresent(deletions -> { - ColumnarRow rowIdRow = toColumnarRow(deletions.getBlock(deletions.getChannelCount() - 1)); - - for (int position = 0; position < rowIdRow.getPositionCount(); position++) { - Slice filePath = VarcharType.VARCHAR.getSlice(rowIdRow.getField(0), position); - long rowPosition = BIGINT.getLong(rowIdRow.getField(1), position); + List fields = RowBlock.getRowFieldsFromBlock(deletions.getBlock(deletions.getChannelCount() - 1)); + Block fieldPathBlock = fields.get(0); + Block rowPositionBlock = fields.get(1); + Block partitionSpecIdBlock = fields.get(2); + Block partitionDataBlock = fields.get(3); + for (int position = 0; position < fieldPathBlock.getPositionCount(); position++) { + Slice filePath = VarcharType.VARCHAR.getSlice(fieldPathBlock, position); + long rowPosition = BIGINT.getLong(rowPositionBlock, position); int index = position; - FileDeletion deletion = fileDeletions.computeIfAbsent(filePath, ignored -> { - int partitionSpecId = INTEGER.getInt(rowIdRow.getField(2), index); - String partitionData = VarcharType.VARCHAR.getSlice(rowIdRow.getField(3), index).toStringUtf8(); + FileDeletion deletion = fileDeletions.computeIfAbsent(filePath, ignore -> { + int partitionSpecId = INTEGER.getInt(partitionSpecIdBlock, index); + String partitionData = VarcharType.VARCHAR.getSlice(partitionDataBlock, index).toStringUtf8(); return new FileDeletion(partitionSpecId, partitionData); }); @@ -125,12 +126,12 @@ public CompletableFuture> finish() List fragments = new ArrayList<>(insertPageSink.finish().join()); fileDeletions.forEach((dataFilePath, deletion) -> { - ConnectorPageSink sink = createPositionDeletePageSink( + PositionDeleteWriter writer = createPositionDeleteWriter( dataFilePath.toStringUtf8(), partitionsSpecs.get(deletion.partitionSpecId()), deletion.partitionDataJson()); - fragments.addAll(writePositionDeletes(sink, deletion.rowsToDelete())); + fragments.addAll(writePositionDeletes(writer, deletion.rowsToDelete())); }); return completedFuture(fragments); @@ -142,7 +143,7 @@ public void abort() insertPageSink.abort(); } - private ConnectorPageSink createPositionDeletePageSink(String dataFilePath, PartitionSpec partitionSpec, String partitionDataJson) + private PositionDeleteWriter createPositionDeleteWriter(String dataFilePath, PartitionSpec partitionSpec, String partitionDataJson) { Optional partitionData = Optional.empty(); if (partitionSpec.isPartitioned()) { @@ -152,7 +153,7 @@ private ConnectorPageSink createPositionDeletePageSink(String dataFilePath, Part partitionData = Optional.of(PartitionData.fromJson(partitionDataJson, columnTypes)); } - return new IcebergPositionDeletePageSink( + return new PositionDeleteWriter( dataFilePath, partitionSpec, partitionData, @@ -165,37 +166,17 @@ private ConnectorPageSink createPositionDeletePageSink(String dataFilePath, Part storageProperties); } - private static Collection writePositionDeletes(ConnectorPageSink sink, ImmutableLongBitmapDataProvider rowsToDelete) + private static Collection writePositionDeletes(PositionDeleteWriter writer, ImmutableLongBitmapDataProvider rowsToDelete) { try { - return doWritePositionDeletes(sink, rowsToDelete); + return writer.write(rowsToDelete); } catch (Throwable t) { - closeAllSuppress(t, sink::abort); + closeAllSuppress(t, writer::abort); throw t; } } - private static Collection doWritePositionDeletes(ConnectorPageSink sink, ImmutableLongBitmapDataProvider rowsToDelete) - { - PageBuilder pageBuilder = new PageBuilder(ImmutableList.of(BIGINT)); - - rowsToDelete.forEach(rowPosition -> { - BIGINT.writeLong(pageBuilder.getBlockBuilder(0), rowPosition); - pageBuilder.declarePosition(); - if (pageBuilder.isFull()) { - sink.appendPage(pageBuilder.build()); - pageBuilder.reset(); - } - }); - - if (!pageBuilder.isEmpty()) { - sink.appendPage(pageBuilder.build()); - } - - return sink.finish().join(); - } - private static class FileDeletion { private final int partitionSpecId; diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergMetadata.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergMetadata.java index cf57078bdbc6..f050abeb1279 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergMetadata.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergMetadata.java @@ -13,6 +13,7 @@ */ package io.trino.plugin.iceberg; +import com.google.common.base.Joiner; import com.google.common.base.Splitter; import com.google.common.base.Suppliers; import com.google.common.base.VerifyException; @@ -32,20 +33,28 @@ import io.trino.filesystem.FileIterator; import io.trino.filesystem.Location; import io.trino.filesystem.TrinoFileSystem; -import io.trino.filesystem.TrinoFileSystemFactory; import io.trino.plugin.base.classloader.ClassLoaderSafeSystemTable; import io.trino.plugin.base.projection.ApplyProjectionUtil; import io.trino.plugin.base.projection.ApplyProjectionUtil.ProjectedColumnRepresentation; +import io.trino.plugin.hive.HiveStorageFormat; import io.trino.plugin.hive.HiveWrittenPartitions; +import io.trino.plugin.hive.metastore.HiveMetastore; +import io.trino.plugin.hive.metastore.HiveMetastoreFactory; +import io.trino.plugin.hive.metastore.TableInfo; import io.trino.plugin.iceberg.aggregation.DataSketchStateSerializer; import io.trino.plugin.iceberg.aggregation.IcebergThetaSketchForStats; import io.trino.plugin.iceberg.catalog.TrinoCatalog; +import io.trino.plugin.iceberg.procedure.IcebergAddFilesFromTableHandle; +import io.trino.plugin.iceberg.procedure.IcebergAddFilesHandle; import io.trino.plugin.iceberg.procedure.IcebergDropExtendedStatsHandle; import io.trino.plugin.iceberg.procedure.IcebergExpireSnapshotsHandle; import io.trino.plugin.iceberg.procedure.IcebergOptimizeHandle; +import io.trino.plugin.iceberg.procedure.IcebergOptimizeManifestsHandle; import io.trino.plugin.iceberg.procedure.IcebergRemoveOrphanFilesHandle; +import io.trino.plugin.iceberg.procedure.IcebergRollbackToSnapshotHandle; import io.trino.plugin.iceberg.procedure.IcebergTableExecuteHandle; import io.trino.plugin.iceberg.procedure.IcebergTableProcedureId; +import io.trino.plugin.iceberg.procedure.MigrationUtils.RecursiveDirectory; import io.trino.plugin.iceberg.util.DataFileWithDeleteFiles; import io.trino.spi.ErrorCode; import io.trino.spi.TrinoException; @@ -99,9 +108,12 @@ import io.trino.spi.statistics.ComputedStatistics; import io.trino.spi.statistics.TableStatistics; import io.trino.spi.statistics.TableStatisticsMetadata; +import io.trino.spi.type.LongTimestamp; import io.trino.spi.type.LongTimestampWithTimeZone; +import io.trino.spi.type.TimestampType; import io.trino.spi.type.TimestampWithTimeZoneType; import io.trino.spi.type.TypeManager; +import io.trino.spi.type.VarcharType; import org.apache.datasketches.theta.CompactSketch; import org.apache.iceberg.AppendFiles; import org.apache.iceberg.BaseTable; @@ -110,6 +122,7 @@ import org.apache.iceberg.DataFiles; import org.apache.iceberg.DeleteFile; import org.apache.iceberg.DeleteFiles; +import org.apache.iceberg.FileFormat; import org.apache.iceberg.FileMetadata; import org.apache.iceberg.FileScanTask; import org.apache.iceberg.IsolationLevel; @@ -122,13 +135,17 @@ import org.apache.iceberg.PartitionSpecParser; import org.apache.iceberg.ReplaceSortOrder; import org.apache.iceberg.RewriteFiles; +import org.apache.iceberg.RewriteManifests; import org.apache.iceberg.RowDelta; import org.apache.iceberg.Schema; import org.apache.iceberg.SchemaParser; import org.apache.iceberg.Snapshot; +import org.apache.iceberg.SnapshotRef; +import org.apache.iceberg.SnapshotUpdate; import org.apache.iceberg.SortField; import org.apache.iceberg.SortOrder; import org.apache.iceberg.StatisticsFile; +import org.apache.iceberg.StructLike; import org.apache.iceberg.Table; import org.apache.iceberg.TableProperties; import org.apache.iceberg.TableScan; @@ -137,11 +154,16 @@ import org.apache.iceberg.UpdateProperties; import org.apache.iceberg.UpdateSchema; import org.apache.iceberg.UpdateStatistics; +import org.apache.iceberg.exceptions.AlreadyExistsException; +import org.apache.iceberg.exceptions.CommitFailedException; +import org.apache.iceberg.exceptions.CommitStateUnknownException; +import org.apache.iceberg.exceptions.NotFoundException; import org.apache.iceberg.exceptions.ValidationException; import org.apache.iceberg.expressions.Expressions; import org.apache.iceberg.expressions.Term; import org.apache.iceberg.io.CloseableIterable; import org.apache.iceberg.types.Type; +import org.apache.iceberg.types.TypeUtil; import org.apache.iceberg.types.Types; import org.apache.iceberg.types.Types.IntegerType; import org.apache.iceberg.types.Types.NestedField; @@ -151,6 +173,9 @@ import java.io.IOException; import java.io.UncheckedIOException; import java.time.Instant; +import java.time.LocalDate; +import java.time.LocalDateTime; +import java.time.ZoneOffset; import java.util.ArrayDeque; import java.util.ArrayList; import java.util.Collection; @@ -165,30 +190,51 @@ import java.util.Optional; import java.util.OptionalLong; import java.util.Set; +import java.util.concurrent.Callable; import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.Executor; +import java.util.concurrent.ExecutorService; import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; import java.util.function.Consumer; +import java.util.function.Function; import java.util.function.Predicate; import java.util.function.Supplier; import java.util.function.UnaryOperator; import java.util.regex.Matcher; import java.util.regex.Pattern; +import java.util.stream.IntStream; +import java.util.stream.Stream; import static com.google.common.base.MoreObjects.firstNonNull; import static com.google.common.base.Preconditions.checkArgument; +import static com.google.common.base.Preconditions.checkState; import static com.google.common.base.Verify.verify; import static com.google.common.base.Verify.verifyNotNull; import static com.google.common.collect.ImmutableList.toImmutableList; import static com.google.common.collect.ImmutableMap.toImmutableMap; import static com.google.common.collect.ImmutableSet.toImmutableSet; import static com.google.common.collect.Iterables.getLast; +import static com.google.common.collect.Iterables.getOnlyElement; import static com.google.common.collect.Maps.transformValues; import static com.google.common.collect.Sets.difference; +import static io.trino.filesystem.Locations.isS3Tables; import static io.trino.plugin.base.projection.ApplyProjectionUtil.extractSupportedProjectedColumns; import static io.trino.plugin.base.projection.ApplyProjectionUtil.replaceWithNewVariables; +import static io.trino.plugin.base.util.ExecutorUtil.processWithAdditionalThreads; import static io.trino.plugin.base.util.Procedures.checkProcedureArgument; -import static io.trino.plugin.hive.util.HiveUtil.isStructuralType; +import static io.trino.plugin.hive.HiveMetadata.TRANSACTIONAL; +import static io.trino.plugin.hive.HiveSessionProperties.isBucketExecutionEnabled; +import static io.trino.plugin.hive.HiveTimestampPrecision.DEFAULT_PRECISION; +import static io.trino.plugin.hive.ViewReaderUtil.isSomeKindOfAView; +import static io.trino.plugin.hive.util.HiveTypeUtil.getTypeSignature; +import static io.trino.plugin.hive.util.HiveUtil.isDeltaLakeTable; +import static io.trino.plugin.hive.util.HiveUtil.isHudiTable; +import static io.trino.plugin.hive.util.HiveUtil.isIcebergTable; +import static io.trino.plugin.iceberg.ColumnIdentity.createColumnIdentity; import static io.trino.plugin.iceberg.ConstraintExtractor.extractTupleDomain; +import static io.trino.plugin.iceberg.ExpressionConverter.isConvertibleToIcebergExpression; import static io.trino.plugin.iceberg.ExpressionConverter.toIcebergExpression; import static io.trino.plugin.iceberg.IcebergAnalyzeProperties.getColumnNames; import static io.trino.plugin.iceberg.IcebergColumnHandle.TRINO_MERGE_PARTITION_DATA; @@ -196,83 +242,136 @@ import static io.trino.plugin.iceberg.IcebergColumnHandle.TRINO_MERGE_ROW_ID; import static io.trino.plugin.iceberg.IcebergColumnHandle.TRINO_ROW_ID_NAME; import static io.trino.plugin.iceberg.IcebergColumnHandle.fileModifiedTimeColumnHandle; +import static io.trino.plugin.iceberg.IcebergColumnHandle.partitionColumnHandle; import static io.trino.plugin.iceberg.IcebergColumnHandle.pathColumnHandle; +import static io.trino.plugin.iceberg.IcebergErrorCode.ICEBERG_CATALOG_ERROR; import static io.trino.plugin.iceberg.IcebergErrorCode.ICEBERG_COMMIT_ERROR; import static io.trino.plugin.iceberg.IcebergErrorCode.ICEBERG_FILESYSTEM_ERROR; import static io.trino.plugin.iceberg.IcebergErrorCode.ICEBERG_INVALID_METADATA; import static io.trino.plugin.iceberg.IcebergErrorCode.ICEBERG_MISSING_METADATA; +import static io.trino.plugin.iceberg.IcebergErrorCode.ICEBERG_UNSUPPORTED_VIEW_DIALECT; import static io.trino.plugin.iceberg.IcebergMetadataColumn.FILE_MODIFIED_TIME; import static io.trino.plugin.iceberg.IcebergMetadataColumn.FILE_PATH; +import static io.trino.plugin.iceberg.IcebergMetadataColumn.PARTITION; import static io.trino.plugin.iceberg.IcebergMetadataColumn.isMetadataColumnId; +import static io.trino.plugin.iceberg.IcebergPartitionFunction.Transform.BUCKET; import static io.trino.plugin.iceberg.IcebergSessionProperties.getExpireSnapshotMinRetention; import static io.trino.plugin.iceberg.IcebergSessionProperties.getHiveCatalogName; import static io.trino.plugin.iceberg.IcebergSessionProperties.getRemoveOrphanFilesMinRetention; import static io.trino.plugin.iceberg.IcebergSessionProperties.isCollectExtendedStatisticsOnWrite; import static io.trino.plugin.iceberg.IcebergSessionProperties.isExtendedStatisticsEnabled; +import static io.trino.plugin.iceberg.IcebergSessionProperties.isFileBasedConflictDetectionEnabled; import static io.trino.plugin.iceberg.IcebergSessionProperties.isMergeManifestsOnWrite; import static io.trino.plugin.iceberg.IcebergSessionProperties.isProjectionPushdownEnabled; import static io.trino.plugin.iceberg.IcebergSessionProperties.isStatisticsEnabled; +import static io.trino.plugin.iceberg.IcebergTableName.isDataTable; +import static io.trino.plugin.iceberg.IcebergTableName.isIcebergTableName; +import static io.trino.plugin.iceberg.IcebergTableName.isMaterializedViewStorage; +import static io.trino.plugin.iceberg.IcebergTableName.tableNameFrom; +import static io.trino.plugin.iceberg.IcebergTableProperties.DATA_LOCATION_PROPERTY; +import static io.trino.plugin.iceberg.IcebergTableProperties.EXTRA_PROPERTIES_PROPERTY; import static io.trino.plugin.iceberg.IcebergTableProperties.FILE_FORMAT_PROPERTY; import static io.trino.plugin.iceberg.IcebergTableProperties.FORMAT_VERSION_PROPERTY; +import static io.trino.plugin.iceberg.IcebergTableProperties.MAX_COMMIT_RETRY; +import static io.trino.plugin.iceberg.IcebergTableProperties.OBJECT_STORE_LAYOUT_ENABLED_PROPERTY; +import static io.trino.plugin.iceberg.IcebergTableProperties.ORC_BLOOM_FILTER_COLUMNS_PROPERTY; +import static io.trino.plugin.iceberg.IcebergTableProperties.PARQUET_BLOOM_FILTER_COLUMNS_PROPERTY; import static io.trino.plugin.iceberg.IcebergTableProperties.PARTITIONING_PROPERTY; import static io.trino.plugin.iceberg.IcebergTableProperties.SORTED_BY_PROPERTY; import static io.trino.plugin.iceberg.IcebergTableProperties.getPartitioning; +import static io.trino.plugin.iceberg.IcebergTableProperties.getTableLocation; +import static io.trino.plugin.iceberg.IcebergUtil.buildPath; import static io.trino.plugin.iceberg.IcebergUtil.canEnforceColumnConstraintInSpecs; +import static io.trino.plugin.iceberg.IcebergUtil.checkFormatForProperty; import static io.trino.plugin.iceberg.IcebergUtil.commit; +import static io.trino.plugin.iceberg.IcebergUtil.createColumnHandle; import static io.trino.plugin.iceberg.IcebergUtil.deserializePartitionValue; import static io.trino.plugin.iceberg.IcebergUtil.fileName; import static io.trino.plugin.iceberg.IcebergUtil.firstSnapshot; import static io.trino.plugin.iceberg.IcebergUtil.firstSnapshotAfter; import static io.trino.plugin.iceberg.IcebergUtil.getColumnHandle; import static io.trino.plugin.iceberg.IcebergUtil.getColumnMetadatas; -import static io.trino.plugin.iceberg.IcebergUtil.getColumns; import static io.trino.plugin.iceberg.IcebergUtil.getFileFormat; import static io.trino.plugin.iceberg.IcebergUtil.getIcebergTableProperties; import static io.trino.plugin.iceberg.IcebergUtil.getPartitionKeys; +import static io.trino.plugin.iceberg.IcebergUtil.getPartitionValues; +import static io.trino.plugin.iceberg.IcebergUtil.getProjectedColumns; import static io.trino.plugin.iceberg.IcebergUtil.getSnapshotIdAsOfTime; import static io.trino.plugin.iceberg.IcebergUtil.getTableComment; +import static io.trino.plugin.iceberg.IcebergUtil.getTopLevelColumns; import static io.trino.plugin.iceberg.IcebergUtil.newCreateTableTransaction; import static io.trino.plugin.iceberg.IcebergUtil.schemaFromMetadata; +import static io.trino.plugin.iceberg.IcebergUtil.validateOrcBloomFilterColumns; +import static io.trino.plugin.iceberg.IcebergUtil.validateParquetBloomFilterColumns; +import static io.trino.plugin.iceberg.IcebergUtil.verifyExtraProperties; import static io.trino.plugin.iceberg.PartitionFields.parsePartitionFields; -import static io.trino.plugin.iceberg.PartitionFields.toPartitionFields; import static io.trino.plugin.iceberg.SortFieldUtils.parseSortFields; -import static io.trino.plugin.iceberg.TableStatisticsReader.TRINO_STATS_COLUMN_ID_PATTERN; -import static io.trino.plugin.iceberg.TableStatisticsReader.TRINO_STATS_PREFIX; +import static io.trino.plugin.iceberg.StructLikeWrapperWithFieldIdToIndex.createStructLikeWrapper; +import static io.trino.plugin.iceberg.TableStatisticsReader.readNdvs; import static io.trino.plugin.iceberg.TableStatisticsWriter.StatsUpdateMode.INCREMENTAL_UPDATE; import static io.trino.plugin.iceberg.TableStatisticsWriter.StatsUpdateMode.REPLACE; import static io.trino.plugin.iceberg.TableType.DATA; +import static io.trino.plugin.iceberg.TypeConverter.toIcebergType; import static io.trino.plugin.iceberg.TypeConverter.toIcebergTypeForNewColumn; import static io.trino.plugin.iceberg.catalog.hms.TrinoHiveCatalog.DEPENDS_ON_TABLES; import static io.trino.plugin.iceberg.catalog.hms.TrinoHiveCatalog.TRINO_QUERY_START_TIME; +import static io.trino.plugin.iceberg.procedure.IcebergTableProcedureId.ADD_FILES; +import static io.trino.plugin.iceberg.procedure.IcebergTableProcedureId.ADD_FILES_FROM_TABLE; import static io.trino.plugin.iceberg.procedure.IcebergTableProcedureId.DROP_EXTENDED_STATS; import static io.trino.plugin.iceberg.procedure.IcebergTableProcedureId.EXPIRE_SNAPSHOTS; import static io.trino.plugin.iceberg.procedure.IcebergTableProcedureId.OPTIMIZE; +import static io.trino.plugin.iceberg.procedure.IcebergTableProcedureId.OPTIMIZE_MANIFESTS; import static io.trino.plugin.iceberg.procedure.IcebergTableProcedureId.REMOVE_ORPHAN_FILES; +import static io.trino.plugin.iceberg.procedure.IcebergTableProcedureId.ROLLBACK_TO_SNAPSHOT; +import static io.trino.plugin.iceberg.procedure.MigrationUtils.addFiles; +import static io.trino.plugin.iceberg.procedure.MigrationUtils.addFilesFromTable; import static io.trino.spi.StandardErrorCode.COLUMN_ALREADY_EXISTS; +import static io.trino.spi.StandardErrorCode.COLUMN_NOT_FOUND; import static io.trino.spi.StandardErrorCode.INVALID_ANALYZE_PROPERTY; import static io.trino.spi.StandardErrorCode.INVALID_ARGUMENTS; +import static io.trino.spi.StandardErrorCode.INVALID_TABLE_PROPERTY; import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED; +import static io.trino.spi.StandardErrorCode.PERMISSION_DENIED; +import static io.trino.spi.StandardErrorCode.TABLE_ALREADY_EXISTS; +import static io.trino.spi.StandardErrorCode.TABLE_NOT_FOUND; +import static io.trino.spi.StandardErrorCode.TYPE_MISMATCH; import static io.trino.spi.connector.MaterializedViewFreshness.Freshness.FRESH; import static io.trino.spi.connector.MaterializedViewFreshness.Freshness.STALE; import static io.trino.spi.connector.MaterializedViewFreshness.Freshness.UNKNOWN; import static io.trino.spi.connector.RetryMode.NO_RETRIES; import static io.trino.spi.connector.RowChangeParadigm.DELETE_ROW_AND_INSERT_ROW; +import static io.trino.spi.predicate.TupleDomain.withColumnDomains; import static io.trino.spi.type.BigintType.BIGINT; import static io.trino.spi.type.DateTimeEncoding.unpackMillisUtc; -import static io.trino.spi.type.UuidType.UUID; +import static io.trino.spi.type.DateType.DATE; +import static io.trino.spi.type.Timestamps.MICROSECONDS_PER_MILLISECOND; +import static java.lang.Boolean.parseBoolean; +import static java.lang.Math.floorDiv; import static java.lang.String.format; import static java.util.Locale.ENGLISH; import static java.util.Objects.requireNonNull; import static java.util.function.Function.identity; import static java.util.stream.Collectors.joining; +import static org.apache.iceberg.MetadataTableType.ALL_ENTRIES; +import static org.apache.iceberg.MetadataTableType.ENTRIES; import static org.apache.iceberg.ReachableFileUtil.metadataFileLocations; -import static org.apache.iceberg.ReachableFileUtil.versionHintLocation; +import static org.apache.iceberg.ReachableFileUtil.statisticsFilesLocations; import static org.apache.iceberg.SnapshotSummary.DELETED_RECORDS_PROP; import static org.apache.iceberg.SnapshotSummary.REMOVED_EQ_DELETES_PROP; import static org.apache.iceberg.SnapshotSummary.REMOVED_POS_DELETES_PROP; +import static org.apache.iceberg.SnapshotSummary.TOTAL_DATA_FILES_PROP; +import static org.apache.iceberg.SnapshotSummary.TOTAL_DELETE_FILES_PROP; +import static org.apache.iceberg.SnapshotSummary.TOTAL_RECORDS_PROP; +import static org.apache.iceberg.TableProperties.COMMIT_NUM_RETRIES; import static org.apache.iceberg.TableProperties.DELETE_ISOLATION_LEVEL; import static org.apache.iceberg.TableProperties.DELETE_ISOLATION_LEVEL_DEFAULT; import static org.apache.iceberg.TableProperties.FORMAT_VERSION; +import static org.apache.iceberg.TableProperties.MANIFEST_TARGET_SIZE_BYTES; +import static org.apache.iceberg.TableProperties.MANIFEST_TARGET_SIZE_BYTES_DEFAULT; +import static org.apache.iceberg.TableProperties.OBJECT_STORE_ENABLED; +import static org.apache.iceberg.TableProperties.ORC_BLOOM_FILTER_COLUMNS; +import static org.apache.iceberg.TableProperties.PARQUET_BLOOM_FILTER_COLUMN_ENABLED_PREFIX; +import static org.apache.iceberg.TableProperties.WRITE_DATA_LOCATION; import static org.apache.iceberg.TableProperties.WRITE_LOCATION_PROVIDER_IMPL; import static org.apache.iceberg.types.TypeUtil.indexParents; import static org.apache.iceberg.util.SnapshotUtil.schemaFor; @@ -286,10 +385,18 @@ public class IcebergMetadata private static final int CLEANING_UP_PROCEDURES_MAX_SUPPORTED_TABLE_VERSION = 2; private static final String RETENTION_THRESHOLD = "retention_threshold"; private static final String UNKNOWN_SNAPSHOT_TOKEN = "UNKNOWN"; - public static final Set UPDATABLE_TABLE_PROPERTIES = ImmutableSet.of(FILE_FORMAT_PROPERTY, FORMAT_VERSION_PROPERTY, PARTITIONING_PROPERTY, SORTED_BY_PROPERTY); - - public static final String ORC_BLOOM_FILTER_COLUMNS_KEY = "orc.bloom.filter.columns"; - public static final String ORC_BLOOM_FILTER_FPP_KEY = "orc.bloom.filter.fpp"; + public static final Set UPDATABLE_TABLE_PROPERTIES = ImmutableSet.builder() + .add(EXTRA_PROPERTIES_PROPERTY) + .add(FILE_FORMAT_PROPERTY) + .add(FORMAT_VERSION_PROPERTY) + .add(MAX_COMMIT_RETRY) + .add(OBJECT_STORE_LAYOUT_ENABLED_PROPERTY) + .add(DATA_LOCATION_PROPERTY) + .add(ORC_BLOOM_FILTER_COLUMNS_PROPERTY) + .add(PARQUET_BLOOM_FILTER_COLUMNS_PROPERTY) + .add(PARTITIONING_PROPERTY) + .add(SORTED_BY_PROPERTY) + .build(); public static final String NUMBER_OF_DISTINCT_VALUES_NAME = "NUMBER_OF_DISTINCT_VALUES"; private static final FunctionName NUMBER_OF_DISTINCT_VALUES_FUNCTION = new FunctionName(IcebergThetaSketchForStats.NAME); @@ -301,20 +408,31 @@ public class IcebergMetadata private final CatalogHandle trinoCatalogHandle; private final JsonCodec commitTaskCodec; private final TrinoCatalog catalog; - private final TrinoFileSystemFactory fileSystemFactory; + private final IcebergFileSystemFactory fileSystemFactory; private final TableStatisticsWriter tableStatisticsWriter; + private final Optional metastoreFactory; + private final boolean addFilesProcedureEnabled; + private final Predicate allowedExtraProperties; + private final ExecutorService icebergScanExecutor; + private final Executor metadataFetchingExecutor; - private final Map tableStatisticsCache = new ConcurrentHashMap<>(); + private final Map> tableStatisticsCache = new ConcurrentHashMap<>(); private Transaction transaction; + private Optional fromSnapshotForRefresh = Optional.empty(); public IcebergMetadata( TypeManager typeManager, CatalogHandle trinoCatalogHandle, JsonCodec commitTaskCodec, TrinoCatalog catalog, - TrinoFileSystemFactory fileSystemFactory, - TableStatisticsWriter tableStatisticsWriter) + IcebergFileSystemFactory fileSystemFactory, + TableStatisticsWriter tableStatisticsWriter, + Optional metastoreFactory, + boolean addFilesProcedureEnabled, + Predicate allowedExtraProperties, + ExecutorService icebergScanExecutor, + Executor metadataFetchingExecutor) { this.typeManager = requireNonNull(typeManager, "typeManager is null"); this.trinoCatalogHandle = requireNonNull(trinoCatalogHandle, "trinoCatalogHandle is null"); @@ -322,6 +440,11 @@ public IcebergMetadata( this.catalog = requireNonNull(catalog, "catalog is null"); this.fileSystemFactory = requireNonNull(fileSystemFactory, "fileSystemFactory is null"); this.tableStatisticsWriter = requireNonNull(tableStatisticsWriter, "tableStatisticsWriter is null"); + this.metastoreFactory = requireNonNull(metastoreFactory, "metastoreFactory is null"); + this.addFilesProcedureEnabled = addFilesProcedureEnabled; + this.allowedExtraProperties = requireNonNull(allowedExtraProperties, "allowedExtraProperties is null"); + this.icebergScanExecutor = requireNonNull(icebergScanExecutor, "icebergScanExecutor is null"); + this.metadataFetchingExecutor = requireNonNull(metadataFetchingExecutor, "metadataFetchingExecutor is null"); } @Override @@ -348,12 +471,6 @@ public Optional getSchemaOwner(ConnectorSession session, String return catalog.getNamespacePrincipal(session, schemaName); } - @Override - public IcebergTableHandle getTableHandle(ConnectorSession session, SchemaTableName tableName) - { - throw new UnsupportedOperationException("This method is not supported because getTableHandle with versions is implemented instead"); - } - @Override public ConnectorTableHandle getTableHandle( ConnectorSession session, @@ -365,14 +482,32 @@ public ConnectorTableHandle getTableHandle( throw new TrinoException(NOT_SUPPORTED, "Read table with start version is not supported"); } - if (!IcebergTableName.isDataTable(tableName.getTableName())) { + if (!isIcebergTableName(tableName.getTableName())) { + return null; + } + + if (isMaterializedViewStorage(tableName.getTableName())) { + verify(endVersion.isEmpty(), "Materialized views do not support versioned queries"); + + SchemaTableName materializedViewName = new SchemaTableName(tableName.getSchemaName(), tableNameFrom(tableName.getTableName())); + if (getMaterializedView(session, materializedViewName).isEmpty()) { + throw new TableNotFoundException(tableName); + } + + BaseTable storageTable = catalog.getMaterializedViewStorageTable(session, materializedViewName) + .orElseThrow(() -> new TrinoException(TABLE_NOT_FOUND, "Storage table metadata not found for materialized view " + tableName)); + + return tableHandleForCurrentSnapshot(session, tableName, storageTable); + } + + if (!isDataTable(tableName.getTableName())) { // Pretend the table does not exist to produce better error message in case of table redirects to Hive return null; } BaseTable table; try { - table = (BaseTable) catalog.loadTable(session, new SchemaTableName(tableName.getSchemaName(), tableName.getTableName())); + table = catalog.loadTable(session, new SchemaTableName(tableName.getSchemaName(), tableName.getTableName())); } catch (TableNotFoundException e) { return null; @@ -386,23 +521,39 @@ public ConnectorTableHandle getTableHandle( throw e; } - Optional tableSnapshotId; - Schema tableSchema; - Optional partitionSpec; if (endVersion.isPresent()) { - long snapshotId = getSnapshotIdFromVersion(table, endVersion.get()); - tableSnapshotId = Optional.of(snapshotId); - tableSchema = schemaFor(table, snapshotId); - partitionSpec = Optional.empty(); - } - else { - tableSnapshotId = Optional.ofNullable(table.currentSnapshot()).map(Snapshot::snapshotId); - tableSchema = table.schema(); - partitionSpec = Optional.of(table.spec()); + long snapshotId = getSnapshotIdFromVersion(session, table, endVersion.get()); + return tableHandleForSnapshot( + session, + tableName, + table, + Optional.of(snapshotId), + schemaFor(table, snapshotId), + Optional.empty()); } + return tableHandleForCurrentSnapshot(session, tableName, table); + } + private IcebergTableHandle tableHandleForCurrentSnapshot(ConnectorSession session, SchemaTableName tableName, BaseTable table) + { + return tableHandleForSnapshot( + session, + tableName, + table, + Optional.ofNullable(table.currentSnapshot()).map(Snapshot::snapshotId), + table.schema(), + Optional.of(table.spec())); + } + + private IcebergTableHandle tableHandleForSnapshot( + ConnectorSession session, + SchemaTableName tableName, + BaseTable table, + Optional tableSnapshotId, + Schema tableSchema, + Optional partitionSpec) + { Map tableProperties = table.properties(); - String nameMappingJson = tableProperties.get(TableProperties.DEFAULT_NAME_MAPPING); return new IcebergTableHandle( trinoCatalogHandle, tableName.getSchemaName(), @@ -416,36 +567,103 @@ public ConnectorTableHandle getTableHandle( TupleDomain.all(), OptionalLong.empty(), ImmutableSet.of(), - Optional.ofNullable(nameMappingJson), + Optional.ofNullable(tableProperties.get(TableProperties.DEFAULT_NAME_MAPPING)), table.location(), table.properties(), + getTablePartitioning(session, table), false, - Optional.empty()); + Optional.empty(), + ImmutableSet.of(), + Optional.of(false)); + } + + private Optional getTablePartitioning(ConnectorSession session, Table icebergTable) + { + if (!isBucketExecutionEnabled(session) || icebergTable.specs().size() != 1) { + return Optional.empty(); + } + PartitionSpec partitionSpec = icebergTable.spec(); + if (partitionSpec.fields().isEmpty()) { + return Optional.empty(); + } + + Schema schema = icebergTable.schema(); + + IcebergPartitioningHandle partitioningHandle = IcebergPartitioningHandle.create(partitionSpec, typeManager, List.of()); + + Map columnById = getProjectedColumns(schema, typeManager).stream() + .collect(toImmutableMap(IcebergColumnHandle::getId, identity())); + List partitionColumns = partitionSpec.fields().stream() + .map(PartitionField::sourceId) + .distinct() + .sorted() + .map(columnById::get) + .collect(toImmutableList()); + + // Partitioning is only activated if it is actually necessary for the query. + // This happens in applyPartitioning + return Optional.of(new IcebergTablePartitioning( + false, + partitioningHandle, + partitionColumns, + IntStream.range(0, partitioningHandle.partitionFunctions().size()).boxed().collect(toImmutableList()))); } - private static long getSnapshotIdFromVersion(Table table, ConnectorTableVersion version) + private static long getSnapshotIdFromVersion(ConnectorSession session, Table table, ConnectorTableVersion version) { io.trino.spi.type.Type versionType = version.getVersionType(); return switch (version.getPointerType()) { - case TEMPORAL -> getTemporalSnapshotIdFromVersion(table, version, versionType); + case TEMPORAL -> getTemporalSnapshotIdFromVersion(session, table, version, versionType); case TARGET_ID -> getTargetSnapshotIdFromVersion(table, version, versionType); }; } private static long getTargetSnapshotIdFromVersion(Table table, ConnectorTableVersion version, io.trino.spi.type.Type versionType) { - if (versionType != BIGINT) { + long snapshotId; + if (versionType == BIGINT) { + snapshotId = (long) version.getVersion(); + } + else if (versionType instanceof VarcharType) { + String refName = ((Slice) version.getVersion()).toStringUtf8(); + SnapshotRef ref = table.refs().get(refName); + if (ref == null) { + throw new TrinoException(INVALID_ARGUMENTS, "Cannot find snapshot with reference name: " + refName); + } + snapshotId = ref.snapshotId(); + } + else { throw new TrinoException(NOT_SUPPORTED, "Unsupported type for table version: " + versionType.getDisplayName()); } - long snapshotId = (long) version.getVersion(); + if (table.snapshot(snapshotId) == null) { throw new TrinoException(INVALID_ARGUMENTS, "Iceberg snapshot ID does not exists: " + snapshotId); } return snapshotId; } - private static long getTemporalSnapshotIdFromVersion(Table table, ConnectorTableVersion version, io.trino.spi.type.Type versionType) + private static long getTemporalSnapshotIdFromVersion(ConnectorSession session, Table table, ConnectorTableVersion version, io.trino.spi.type.Type versionType) { + if (versionType.equals(DATE)) { + // Retrieve the latest snapshot made before or at the beginning of the day of the specified date in the session's time zone + long epochMillis = LocalDate.ofEpochDay((Long) version.getVersion()) + .atStartOfDay() + .atZone(session.getTimeZoneKey().getZoneId()) + .toInstant() + .toEpochMilli(); + return getSnapshotIdAsOfTime(table, epochMillis); + } + if (versionType instanceof TimestampType timestampVersionType) { + long epochMicrosUtc = timestampVersionType.isShort() + ? (long) version.getVersion() + : ((LongTimestamp) version.getVersion()).getEpochMicros(); + long epochMillisUtc = floorDiv(epochMicrosUtc, MICROSECONDS_PER_MILLISECOND); + long epochMillis = LocalDateTime.ofInstant(Instant.ofEpochMilli(epochMillisUtc), ZoneOffset.UTC) + .atZone(session.getTimeZoneKey().getZoneId()) + .toInstant() + .toEpochMilli(); + return getSnapshotIdAsOfTime(table, epochMillis); + } if (versionType instanceof TimestampWithTimeZoneType timeZonedVersionType) { long epochMillis = timeZonedVersionType.isShort() ? unpackMillisUtc((long) version.getVersion()) @@ -464,12 +682,12 @@ public Optional getSystemTable(ConnectorSession session, SchemaTabl private Optional getRawSystemTable(ConnectorSession session, SchemaTableName tableName) { - if (IcebergTableName.isDataTable(tableName.getTableName())) { + if (!isIcebergTableName(tableName.getTableName()) || isDataTable(tableName.getTableName()) || isMaterializedViewStorage(tableName.getTableName())) { return Optional.empty(); } // Only when dealing with an actual system table proceed to retrieve the base table for the system table - String name = IcebergTableName.tableNameFrom(tableName.getTableName()); + String name = tableNameFrom(tableName.getTableName()); Table table; try { table = catalog.loadTable(session, new SchemaTableName(tableName.getSchemaName(), name)); @@ -482,20 +700,20 @@ private Optional getRawSystemTable(ConnectorSession session, Schema return Optional.empty(); } - Optional tableType = IcebergTableName.tableTypeFrom(tableName.getTableName()); - if (tableType.isEmpty()) { - return Optional.empty(); - } - SchemaTableName systemTableName = new SchemaTableName(tableName.getSchemaName(), IcebergTableName.tableNameWithType(name, tableType.get())); - return switch (tableType.get()) { - case DATA -> throw new VerifyException("Unexpected DATA table type"); // Handled above. - case HISTORY -> Optional.of(new HistoryTable(systemTableName, table)); - case SNAPSHOTS -> Optional.of(new SnapshotsTable(systemTableName, typeManager, table)); - case PARTITIONS -> Optional.of(new PartitionTable(systemTableName, typeManager, table, getCurrentSnapshotId(table))); - case MANIFESTS -> Optional.of(new ManifestsTable(systemTableName, table, getCurrentSnapshotId(table))); - case FILES -> Optional.of(new FilesTable(systemTableName, typeManager, table, getCurrentSnapshotId(table))); - case PROPERTIES -> Optional.of(new PropertiesTable(systemTableName, table)); - case REFS -> Optional.of(new RefsTable(systemTableName, table)); + TableType tableType = IcebergTableName.tableTypeFrom(tableName.getTableName()); + return switch (tableType) { + case DATA, MATERIALIZED_VIEW_STORAGE -> throw new VerifyException("Unexpected table type: " + tableType); // Handled above. + case HISTORY -> Optional.of(new HistoryTable(tableName, table)); + case METADATA_LOG_ENTRIES -> Optional.of(new MetadataLogEntriesTable(tableName, table, icebergScanExecutor)); + case SNAPSHOTS -> Optional.of(new SnapshotsTable(tableName, typeManager, table, icebergScanExecutor)); + case PARTITIONS -> Optional.of(new PartitionsTable(tableName, typeManager, table, getCurrentSnapshotId(table), icebergScanExecutor)); + case ALL_MANIFESTS -> Optional.of(new AllManifestsTable(tableName, table, icebergScanExecutor)); + case MANIFESTS -> Optional.of(new ManifestsTable(tableName, table, getCurrentSnapshotId(table))); + case FILES -> Optional.of(new FilesTable(tableName, typeManager, table, getCurrentSnapshotId(table), icebergScanExecutor)); + case ALL_ENTRIES -> Optional.of(new EntriesTable(typeManager, tableName, table, ALL_ENTRIES, icebergScanExecutor)); + case ENTRIES -> Optional.of(new EntriesTable(typeManager, tableName, table, ENTRIES, icebergScanExecutor)); + case PROPERTIES -> Optional.of(new PropertiesTable(tableName, table)); + case REFS -> Optional.of(new RefsTable(tableName, table, icebergScanExecutor)); }; } @@ -520,44 +738,48 @@ public ConnectorTableProperties getTableProperties(ConnectorSession session, Con DiscretePredicates discretePredicates = null; if (!partitionSourceIds.isEmpty()) { // Extract identity partition columns - Map columns = getColumns(icebergTable.schema(), typeManager).stream() - .filter(column -> partitionSourceIds.contains(column.getId())) + Map columns = getProjectedColumns(icebergTable.schema(), typeManager, partitionSourceIds).stream() .collect(toImmutableMap(IcebergColumnHandle::getId, identity())); - Supplier> lazyFiles = Suppliers.memoize(() -> { + Supplier> lazyUniquePartitions = Suppliers.memoize(() -> { TableScan tableScan = icebergTable.newScan() .useSnapshot(table.getSnapshotId().get()) .filter(toIcebergExpression(enforcedPredicate)); - try (CloseableIterable iterator = tableScan.planFiles()) { - return ImmutableList.copyOf(iterator); + try (CloseableIterable fileScanTasks = tableScan.planFiles()) { + Map partitions = new HashMap<>(); + for (FileScanTask fileScanTask : fileScanTasks) { + StructLikeWrapperWithFieldIdToIndex structLikeWrapperWithFieldIdToIndex = createStructLikeWrapper(fileScanTask); + partitions.putIfAbsent(structLikeWrapperWithFieldIdToIndex, fileScanTask.spec()); + } + return partitions; } catch (IOException e) { throw new UncheckedIOException(e); } }); - Iterable files = () -> lazyFiles.get().iterator(); - - Iterable> discreteTupleDomain = Iterables.transform(files, fileScan -> { - // Extract partition values in the data file - Map> partitionColumnValueStrings = getPartitionKeys(fileScan); - Map partitionValues = partitionSourceIds.stream() - .filter(partitionColumnValueStrings::containsKey) - .collect(toImmutableMap( - columns::get, - columnId -> { - IcebergColumnHandle column = columns.get(columnId); - Object prestoValue = deserializePartitionValue( - column.getType(), - partitionColumnValueStrings.get(columnId).orElse(null), - column.getName()); - - return NullableValue.of(column.getType(), prestoValue); - })); - - return TupleDomain.fromFixedValues(partitionValues); - }); + Iterable> discreteTupleDomain = Iterables.transform( + () -> lazyUniquePartitions.get().entrySet().iterator(), + entry -> { + // Extract partition values + Map> partitionColumnValueStrings = getPartitionKeys(entry.getKey().getStructLikeWrapper().get(), entry.getValue()); + Map partitionValues = partitionSourceIds.stream() + .filter(partitionColumnValueStrings::containsKey) + .collect(toImmutableMap( + columns::get, + columnId -> { + IcebergColumnHandle column = columns.get(columnId); + Object prestoValue = deserializePartitionValue( + column.getType(), + partitionColumnValueStrings.get(columnId).orElse(null), + column.getName()); + + return new NullableValue(column.getType(), prestoValue); + })); + + return TupleDomain.fromFixedValues(partitionValues); + }); discretePredicates = new DiscretePredicates( columns.values().stream() @@ -572,12 +794,64 @@ public ConnectorTableProperties getTableProperties(ConnectorSession session, Con // can be further optimized by intersecting with partition values at the cost of iterating // over all tableScan.planFiles() and caching partition values in table handle. enforcedPredicate.transformKeys(ColumnHandle.class::cast), - // TODO: implement table partitioning - Optional.empty(), + table.getTablePartitioning().flatMap(IcebergTablePartitioning::toConnectorTablePartitioning), Optional.ofNullable(discretePredicates), + // todo support sorting properties ImmutableList.of()); } + public Optional applyPartitioning(ConnectorSession session, ConnectorTableHandle tableHandle, Optional partitioningHandle, List partitioningColumns) + { + IcebergTableHandle icebergTableHandle = checkValidTableHandle(tableHandle); + if (icebergTableHandle.getPartitionSpecJson().isEmpty()) { + return Optional.empty(); + } + + Optional connectorTablePartitioning = icebergTableHandle.getTablePartitioning(); + if (connectorTablePartitioning.isEmpty()) { + return Optional.empty(); + } + IcebergTablePartitioning tablePartitioning = connectorTablePartitioning.get(); + + // Check if the table can be partitioned on the requested columns + if (!new HashSet<>(tablePartitioning.partitioningColumns()).containsAll(partitioningColumns)) { + return Optional.empty(); + } + + Map newPartitioningColumnIndex = IntStream.range(0, partitioningColumns.size()).boxed() + .collect(toImmutableMap(partitioningColumns::get, identity())); + ImmutableList.Builder newPartitionFunctions = ImmutableList.builder(); + ImmutableList.Builder newPartitionStructFields = ImmutableList.builder(); + for (int functionIndex = 0; functionIndex < tablePartitioning.partitioningHandle().partitionFunctions().size(); functionIndex++) { + IcebergPartitionFunction function = tablePartitioning.partitioningHandle().partitionFunctions().get(functionIndex); + int oldColumnIndex = function.dataPath().get(0); + Integer newColumnIndex = newPartitioningColumnIndex.get(tablePartitioning.partitioningColumns().get(oldColumnIndex)); + if (newColumnIndex != null) { + // Change the index of the top level column to the location in the new partitioning columns + newPartitionFunctions.add(function.withTopLevelColumnIndex(newColumnIndex)); + // Some partition functions may be dropped so update the struct fields used in split partitioning must be updated + newPartitionStructFields.add(tablePartitioning.partitionStructFields().get(functionIndex)); + } + } + + IcebergPartitioningHandle newPartitioningHandle = new IcebergPartitioningHandle(false, newPartitionFunctions.build()); + if (partitioningHandle.isPresent() && !partitioningHandle.get().equals(newPartitioningHandle)) { + // todo if bucketing is a power of two, we can adapt the bucketing + return Optional.empty(); + } + if (newPartitioningHandle.partitionFunctions().stream().map(IcebergPartitionFunction::transform).noneMatch(BUCKET::equals)) { + // The table is only using value-based partitioning, and this can hurt performance if there is a filter + // on the partitioning columns. This is something we may be able to support with statistics in the future. + return Optional.empty(); + } + + return Optional.of(icebergTableHandle.withTablePartitioning(Optional.of(new IcebergTablePartitioning( + true, + newPartitioningHandle, + partitioningColumns.stream().map(IcebergColumnHandle.class::cast).collect(toImmutableList()), + newPartitionStructFields.build())))); + } + @Override public SchemaTableName getTableName(ConnectorSession session, ConnectorTableHandle table) { @@ -593,7 +867,7 @@ public ConnectorTableMetadata getTableMetadata(ConnectorSession session, Connect IcebergTableHandle tableHandle = checkValidTableHandle(table); // This method does not calculate column metadata for the projected columns checkArgument(tableHandle.getProjectedColumns().isEmpty(), "Unexpected projected columns"); - Table icebergTable = catalog.loadTable(session, tableHandle.getSchemaTableName()); + BaseTable icebergTable = catalog.loadTable(session, tableHandle.getSchemaTableName()); List columns = getColumnMetadatas(SchemaParser.fromJson(tableHandle.getTableSchemaJson()), typeManager); return new ConnectorTableMetadata(tableHandle.getSchemaTableName(), columns, getIcebergTableProperties(icebergTable), getTableComment(icebergTable)); } @@ -601,7 +875,9 @@ public ConnectorTableMetadata getTableMetadata(ConnectorSession session, Connect @Override public List listTables(ConnectorSession session, Optional schemaName) { - return catalog.listTables(session, schemaName); + return catalog.listTables(session, schemaName).stream() + .map((TableInfo::tableName)) + .toList(); } @Override @@ -609,9 +885,10 @@ public Map getColumnHandles(ConnectorSession session, Conn { IcebergTableHandle table = checkValidTableHandle(tableHandle); ImmutableMap.Builder columnHandles = ImmutableMap.builder(); - for (IcebergColumnHandle columnHandle : getColumns(SchemaParser.fromJson(table.getTableSchemaJson()), typeManager)) { + for (IcebergColumnHandle columnHandle : getTopLevelColumns(SchemaParser.fromJson(table.getTableSchemaJson()), typeManager)) { columnHandles.put(columnHandle.getName(), columnHandle); } + columnHandles.put(PARTITION.getColumnName(), partitionColumnHandle()); columnHandles.put(FILE_PATH.getColumnName(), pathColumnHandle()); columnHandles.put(FILE_MODIFIED_TIME.getColumnName(), fileModifiedTimeColumnHandle()); return columnHandles.buildOrThrow(); @@ -640,7 +917,9 @@ public Iterator streamTableColumns(ConnectorSession sessio requireNonNull(prefix, "prefix is null"); List schemaTableNames; if (prefix.getTable().isEmpty()) { - schemaTableNames = catalog.listTables(session, prefix.getSchema()); + schemaTableNames = catalog.listTables(session, prefix.getSchema()).stream() + .map(TableInfo::tableName) + .collect(toImmutableList()); } else { schemaTableNames = ImmutableList.of(prefix.toSchemaTableName()); @@ -665,26 +944,40 @@ public Iterator streamTableColumns(ConnectorSession sessio tableMetadatas.add(TableColumnsMetadata.forTable(tableName, columns)); }); - for (SchemaTableName tableName : remainingTables) { - try { - Table icebergTable = catalog.loadTable(session, tableName); - List columns = getColumnMetadatas(icebergTable.schema(), typeManager); - tableMetadatas.add(TableColumnsMetadata.forTable(tableName, columns)); - } - catch (TableNotFoundException e) { - // Table disappeared during listing operation - continue; - } - catch (UnknownTableTypeException e) { - // Skip unsupported table type in case that the table redirects are not enabled - continue; - } - catch (RuntimeException e) { - // Table can be being removed and this may cause all sorts of exceptions. Log, because we're catching broadly. - log.warn(e, "Failed to access metadata of table %s during streaming table columns for %s", tableName, prefix); - continue; - } + List>> tasks = remainingTables.stream() + .map(tableName -> (Callable>) () -> { + try { + Table icebergTable = catalog.loadTable(session, tableName); + List columns = getColumnMetadatas(icebergTable.schema(), typeManager); + return Optional.of(TableColumnsMetadata.forTable(tableName, columns)); + } + catch (TableNotFoundException e) { + // Table disappeared during listing operation + return Optional.empty(); + } + catch (UnknownTableTypeException e) { + // Skip unsupported table type in case that the table redirects are not enabled + return Optional.empty(); + } + catch (RuntimeException e) { + // Table can be being removed and this may cause all sorts of exceptions. Log, because we're catching broadly. + log.warn(e, "Failed to access metadata of table %s during streaming table columns for %s", tableName, prefix); + return Optional.empty(); + } + }) + .collect(toImmutableList()); + + try { + List taskResults = processWithAdditionalThreads(tasks, metadataFetchingExecutor).stream() + .flatMap(Optional::stream) // Flatten the Optionals into a stream + .collect(toImmutableList()); + + tableMetadatas.addAll(taskResults); + } + catch (ExecutionException e) { + throw new RuntimeException(e.getCause()); } + return tableMetadatas.build(); }) .flatMap(List::stream) @@ -718,8 +1011,26 @@ public void createSchema(ConnectorSession session, String schemaName, Map nestedNamespaces = getChildNamespaces(session, schemaName); + if (!nestedNamespaces.isEmpty()) { + throw new TrinoException( + ICEBERG_CATALOG_ERROR, + format("Cannot drop non-empty schema: %s, contains %s nested schema(s)", schemaName, Joiner.on(", ").join(nestedNamespaces))); + } + + for (SchemaTableName materializedView : listMaterializedViews(session, Optional.of(schemaName))) { + dropMaterializedView(session, materializedView); + } + for (SchemaTableName viewName : listViews(session, Optional.of(schemaName))) { + dropView(session, viewName); + } + for (SchemaTableName tableName : listTables(session, Optional.of(schemaName))) { + dropTable(session, getTableHandle(session, tableName, Optional.empty(), Optional.empty())); + } + } catalog.dropNamespace(session, schemaName); } @@ -783,14 +1094,20 @@ public ConnectorOutputTableHandle beginCreateTable(ConnectorSession session, Con if (!schemaExists(session, schemaName)) { throw new SchemaNotFoundException(schemaName); } - transaction = newCreateTableTransaction(catalog, tableMetadata, session); + + String tableLocation = getTableLocation(tableMetadata.getProperties()) + .orElseGet(() -> catalog.defaultTableLocation(session, tableMetadata.getTable())); + transaction = newCreateTableTransaction(catalog, tableMetadata, session, tableLocation, allowedExtraProperties); Location location = Location.of(transaction.table().location()); - TrinoFileSystem fileSystem = fileSystemFactory.create(session); try { - if (fileSystem.listFiles(location).hasNext()) { - throw new TrinoException(ICEBERG_FILESYSTEM_ERROR, format("" + - "Cannot create a table on a non-empty location: %s, set 'iceberg.unique-table-location=true' in your Iceberg catalog properties " + - "to use unique table locations for every table.", location)); + // S3 Tables internally assigns a unique location for each table + if (!isS3Tables(location.toString())) { + TrinoFileSystem fileSystem = fileSystemFactory.create(session.getIdentity(), transaction.table().io().properties()); + if (fileSystem.listFiles(location).hasNext()) { + throw new TrinoException(ICEBERG_FILESYSTEM_ERROR, format("" + + "Cannot create a table on a non-empty location: %s, set 'iceberg.unique-table-location=true' in your Iceberg catalog properties " + + "to use unique table locations for every table.", location)); + } } return newWritableTableHandle(tableMetadata.getTable(), transaction.table(), retryMode); } @@ -802,16 +1119,22 @@ public ConnectorOutputTableHandle beginCreateTable(ConnectorSession session, Con @Override public Optional finishCreateTable(ConnectorSession session, ConnectorOutputTableHandle tableHandle, Collection fragments, Collection computedStatistics) { - if (fragments.isEmpty()) { - // Commit the transaction if the table is being created without data - AppendFiles appendFiles = transaction.newFastAppend(); - commit(appendFiles, session); - transaction.commitTransaction(); - transaction = null; - return Optional.empty(); - } + IcebergWritableTableHandle icebergTableHandle = (IcebergWritableTableHandle) tableHandle; + try { + if (fragments.isEmpty()) { + // Commit the transaction if the table is being created without data + AppendFiles appendFiles = transaction.newFastAppend(); + commitUpdateAndTransaction(appendFiles, session, transaction, "create table"); + transaction = null; + return Optional.empty(); + } - return finishInsert(session, (IcebergWritableTableHandle) tableHandle, fragments, computedStatistics); + return finishInsert(session, icebergTableHandle, fragments, computedStatistics); + } + catch (AlreadyExistsException e) { + // May happen when table has been already created concurrently. + throw new TrinoException(TABLE_ALREADY_EXISTS, format("Table %s already exists", icebergTableHandle.getName()), e); + } } @Override @@ -831,13 +1154,35 @@ private Optional getWriteLayout(Schema tableSchema, Partit return Optional.empty(); } - validateNotPartitionedByNestedField(tableSchema, partitionSpec); - Map columnById = getColumns(tableSchema, typeManager).stream() - .collect(toImmutableMap(IcebergColumnHandle::getId, identity())); + StructType schemaAsStruct = tableSchema.asStruct(); + Map indexById = TypeUtil.indexById(schemaAsStruct); + Map indexParents = indexParents(schemaAsStruct); + Map> indexPaths = indexById.entrySet().stream() + .collect(toImmutableMap(Map.Entry::getKey, entry -> ImmutableList.copyOf(buildPath(indexParents, entry.getKey())))); List partitioningColumns = partitionSpec.fields().stream() .sorted(Comparator.comparing(PartitionField::sourceId)) - .map(field -> requireNonNull(columnById.get(field.sourceId()), () -> "Cannot find source column for partitioning field " + field)) + .map(field -> { + boolean isBaseColumn = !indexParents.containsKey(field.sourceId()); + int sourceId; + if (isBaseColumn) { + sourceId = field.sourceId(); + } + else { + sourceId = getRootFieldId(indexParents, field.sourceId()); + } + Type sourceType = tableSchema.findType(sourceId); + // The source column, must be a primitive type and cannot be contained in a map or list, but may be nested in a struct. + // https://iceberg.apache.org/spec/#partitioning + if (sourceType.isMapType()) { + throw new TrinoException(NOT_SUPPORTED, "Partitioning field [" + field.name() + "] cannot be contained in a map"); + } + if (sourceType.isListType()) { + throw new TrinoException(NOT_SUPPORTED, "Partitioning field [" + field.name() + "] cannot be contained in a array"); + } + verify(indexById.containsKey(sourceId), "Cannot find source column for partition field " + field); + return createColumnHandle(typeManager, sourceId, indexById, indexPaths); + }) .distinct() .collect(toImmutableList()); List partitioningColumnNames = partitioningColumns.stream() @@ -848,10 +1193,19 @@ private Optional getWriteLayout(Schema tableSchema, Partit // Do not set partitioningHandle, to let engine determine whether to repartition data or not, on stat-based basis. return Optional.of(new ConnectorTableLayout(partitioningColumnNames)); } - IcebergPartitioningHandle partitioningHandle = new IcebergPartitioningHandle(toPartitionFields(partitionSpec), partitioningColumns); + IcebergPartitioningHandle partitioningHandle = IcebergPartitioningHandle.create(partitionSpec, typeManager, List.of()); return Optional.of(new ConnectorTableLayout(partitioningHandle, partitioningColumnNames, true)); } + private static int getRootFieldId(Map indexParents, int fieldId) + { + int rootFieldId = fieldId; + while (indexParents.containsKey(rootFieldId)) { + rootFieldId = indexParents.get(rootFieldId); + } + return rootFieldId; + } + @Override public ConnectorInsertTableHandle beginInsert(ConnectorSession session, ConnectorTableHandle tableHandle, List columns, RetryMode retryMode) { @@ -859,13 +1213,25 @@ public ConnectorInsertTableHandle beginInsert(ConnectorSession session, Connecto Table icebergTable = catalog.loadTable(session, table.getSchemaTableName()); validateNotModifyingOldSnapshot(table, icebergTable); - validateNotPartitionedByNestedField(icebergTable.schema(), icebergTable.spec()); beginTransaction(icebergTable); return newWritableTableHandle(table.getSchemaTableName(), icebergTable, retryMode); } + private List getChildNamespaces(ConnectorSession session, String parentNamespace) + { + Optional namespaceSeparator = catalog.getNamespaceSeparator(); + + if (namespaceSeparator.isEmpty()) { + return ImmutableList.of(); + } + + return catalog.listNamespaces(session).stream() + .filter(namespace -> namespace.startsWith(parentNamespace + namespaceSeparator.get())) + .collect(toImmutableList()); + } + private IcebergWritableTableHandle newWritableTableHandle(SchemaTableName name, Table table, RetryMode retryMode) { return new IcebergWritableTableHandle( @@ -874,11 +1240,12 @@ private IcebergWritableTableHandle newWritableTableHandle(SchemaTableName name, transformValues(table.specs(), PartitionSpecParser::toJson), table.spec().specId(), getSupportedSortFields(table.schema(), table.sortOrder()), - getColumns(table.schema(), typeManager), + getProjectedColumns(table.schema(), typeManager), table.location(), getFileFormat(table), table.properties(), - retryMode); + retryMode, + table.io().properties()); } private static List getSupportedSortFields(Schema schema, SortOrder sortOrder) @@ -890,7 +1257,7 @@ private static List getSupportedSortFields(Schema schema, SortOr .map(Types.NestedField::fieldId) .collect(toImmutableSet()); - ImmutableList.Builder sortFields = ImmutableList.builder(); + ImmutableList.Builder sortFields = ImmutableList.builder(); for (SortField sortField : sortOrder.fields()) { if (!sortField.transform().isIdentity()) { continue; @@ -909,7 +1276,8 @@ private static List getSupportedSortFields(Schema schema, SortOr public Optional finishInsert(ConnectorSession session, ConnectorInsertTableHandle insertHandle, Collection fragments, Collection computedStatistics) { List commitTasks = fragments.stream() - .map(slice -> commitTaskCodec.fromJson(slice.getBytes())) + .map(Slice::getBytes) + .map(commitTaskCodec::fromJson) .collect(toImmutableList()); if (commitTasks.isEmpty()) { @@ -934,6 +1302,7 @@ public Optional finishInsert(ConnectorSession session, .withFileSizeInBytes(task.getFileSizeInBytes()) .withFormat(table.getFileFormat().toIceberg()) .withMetrics(task.getMetrics().metrics()); + task.getFileSplitOffsets().ifPresent(builder::withSplitOffsets); if (!icebergTable.spec().fields().isEmpty()) { String partitionDataJson = task.getPartitionDataJson() @@ -950,17 +1319,19 @@ public Optional finishInsert(ConnectorSession session, cleanExtraOutputFiles(session, writtenFiles.build()); } - commit(appendFiles, session); - transaction.commitTransaction(); + commitUpdateAndTransaction(appendFiles, session, transaction, "insert"); // TODO (https://github.com/trinodb/trino/issues/15439) this may not exactly be the snapshot we committed, if there is another writer long newSnapshotId = transaction.table().currentSnapshot().snapshotId(); transaction = null; // TODO (https://github.com/trinodb/trino/issues/15439): it would be good to publish data and stats atomically beforeWriteSnapshotId.ifPresent(previous -> - verify(previous != newSnapshotId, "Failed to get new snapshot ID ")); + verify(previous != newSnapshotId, "Failed to get new snapshot ID")); - if (!computedStatistics.isEmpty()) { + if (isS3Tables(icebergTable.location())) { + log.debug("S3 Tables does not support statistics: %s", table.getName()); + } + else if (!computedStatistics.isEmpty()) { try { beginTransaction(catalog.loadTable(session, table.getName())); Table reloadedTable = transaction.table(); @@ -972,10 +1343,10 @@ public Optional finishInsert(ConnectorSession session, INCREMENTAL_UPDATE, collectedStatistics); transaction.updateStatistics() - .setStatistics(newSnapshotId, statisticsFile) + .setStatistics(statisticsFile) .commit(); - transaction.commitTransaction(); + commitTransaction(transaction, "update statistics on insert"); } catch (Exception e) { // Write was committed, so at this point we cannot fail the query @@ -992,7 +1363,7 @@ public Optional finishInsert(ConnectorSession session, private void cleanExtraOutputFiles(ConnectorSession session, Set writtenFiles) { - TrinoFileSystem fileSystem = fileSystemFactory.create(session); + TrinoFileSystem fileSystem = fileSystemFactory.create(session.getIdentity(), transaction.table().io().properties()); Set locations = getOutputFilesLocations(writtenFiles); Set fileNames = getOutputFilesFileNames(writtenFiles); for (String location : locations) { @@ -1097,9 +1468,13 @@ public Optional getTableHandleForExecute( return switch (procedureId) { case OPTIMIZE -> getTableHandleForOptimize(tableHandle, icebergTable, executeProperties, retryMode); + case OPTIMIZE_MANIFESTS -> getTableHandleForOptimizeManifests(session, tableHandle); case DROP_EXTENDED_STATS -> getTableHandleForDropExtendedStats(session, tableHandle); + case ROLLBACK_TO_SNAPSHOT -> getTableHandleForRollbackToSnapshot(session, tableHandle, executeProperties); case EXPIRE_SNAPSHOTS -> getTableHandleForExpireSnapshots(session, tableHandle, executeProperties); case REMOVE_ORPHAN_FILES -> getTableHandleForRemoveOrphanFiles(session, tableHandle, executeProperties); + case ADD_FILES -> getTableHandleForAddFiles(session, tableHandle, executeProperties); + case ADD_FILES_FROM_TABLE -> getTableHandleForAddFilesFromTable(session, tableHandle, executeProperties); }; } @@ -1118,7 +1493,7 @@ private Optional getTableHandleForOptimize( tableHandle.getSnapshotId(), tableHandle.getTableSchemaJson(), tableHandle.getPartitionSpecJson().orElseThrow(() -> new VerifyException("Partition spec missing in the table handle")), - getColumns(SchemaParser.fromJson(tableHandle.getTableSchemaJson()), typeManager), + getProjectedColumns(SchemaParser.fromJson(tableHandle.getTableSchemaJson()), typeManager), icebergTable.sortOrder().fields().stream() .map(TrinoSortField::fromIceberg) .collect(toImmutableList()), @@ -1126,7 +1501,20 @@ private Optional getTableHandleForOptimize( tableHandle.getStorageProperties(), maxScannedFileSize, retryMode != NO_RETRIES), - tableHandle.getTableLocation())); + tableHandle.getTableLocation(), + icebergTable.io().properties())); + } + + private Optional getTableHandleForOptimizeManifests(ConnectorSession session, IcebergTableHandle tableHandle) + { + Table icebergTable = catalog.loadTable(session, tableHandle.getSchemaTableName()); + + return Optional.of(new IcebergTableExecuteHandle( + tableHandle.getSchemaTableName(), + OPTIMIZE_MANIFESTS, + new IcebergOptimizeManifestsHandle(), + icebergTable.location(), + icebergTable.io().properties())); } private Optional getTableHandleForDropExtendedStats(ConnectorSession session, IcebergTableHandle tableHandle) @@ -1137,7 +1525,8 @@ private Optional getTableHandleForDropExtendedStats tableHandle.getSchemaTableName(), DROP_EXTENDED_STATS, new IcebergDropExtendedStatsHandle(), - icebergTable.location())); + icebergTable.location(), + icebergTable.io().properties())); } private Optional getTableHandleForExpireSnapshots(ConnectorSession session, IcebergTableHandle tableHandle, Map executeProperties) @@ -1149,7 +1538,8 @@ private Optional getTableHandleForExpireSnapshots(C tableHandle.getSchemaTableName(), EXPIRE_SNAPSHOTS, new IcebergExpireSnapshotsHandle(retentionThreshold), - icebergTable.location())); + icebergTable.location(), + icebergTable.io().properties())); } private Optional getTableHandleForRemoveOrphanFiles(ConnectorSession session, IcebergTableHandle tableHandle, Map executeProperties) @@ -1161,7 +1551,119 @@ private Optional getTableHandleForRemoveOrphanFiles tableHandle.getSchemaTableName(), REMOVE_ORPHAN_FILES, new IcebergRemoveOrphanFilesHandle(retentionThreshold), - icebergTable.location())); + icebergTable.location(), + icebergTable.io().properties())); + } + + private Optional getTableHandleForRollbackToSnapshot(ConnectorSession session, IcebergTableHandle tableHandle, Map executeProperties) + { + long snapshotId = (long) executeProperties.get("snapshot_id"); + Table icebergTable = catalog.loadTable(session, tableHandle.getSchemaTableName()); + + return Optional.of(new IcebergTableExecuteHandle( + tableHandle.getSchemaTableName(), + ROLLBACK_TO_SNAPSHOT, + new IcebergRollbackToSnapshotHandle(snapshotId), + icebergTable.location(), + icebergTable.io().properties())); + } + + private Optional getTableHandleForAddFiles(ConnectorSession session, IcebergTableHandle tableHandle, Map executeProperties) + { + if (!addFilesProcedureEnabled) { + throw new TrinoException(PERMISSION_DENIED, "add_files procedure is disabled"); + } + + String location = (String) requireProcedureArgument(executeProperties, "location"); + HiveStorageFormat format = (HiveStorageFormat) requireProcedureArgument(executeProperties, "format"); + RecursiveDirectory recursiveDirectory = (RecursiveDirectory) executeProperties.getOrDefault("recursive_directory", "fail"); + + Table icebergTable = catalog.loadTable(session, tableHandle.getSchemaTableName()); + + return Optional.of(new IcebergTableExecuteHandle( + tableHandle.getSchemaTableName(), + ADD_FILES, + new IcebergAddFilesHandle(location, format, recursiveDirectory), + icebergTable.location(), + icebergTable.io().properties())); + } + + private Optional getTableHandleForAddFilesFromTable(ConnectorSession session, IcebergTableHandle tableHandle, Map executeProperties) + { + String schemaName = (String) requireProcedureArgument(executeProperties, "schema_name"); + String tableName = (String) requireProcedureArgument(executeProperties, "table_name"); + @SuppressWarnings("unchecked") + Map partitionFilter = (Map) executeProperties.get("partition_filter"); + RecursiveDirectory recursiveDirectory = (RecursiveDirectory) executeProperties.getOrDefault("recursive_directory", "fail"); + + HiveMetastore metastore = metastoreFactory.orElseThrow(() -> new TrinoException(NOT_SUPPORTED, "This catalog does not support add_files_from_table procedure")) + .createMetastore(Optional.of(session.getIdentity())); + SchemaTableName sourceName = new SchemaTableName(schemaName, tableName); + io.trino.plugin.hive.metastore.Table sourceTable = metastore.getTable(schemaName, tableName).orElseThrow(() -> new TableNotFoundException(sourceName)); + + Table icebergTable = catalog.loadTable(session, tableHandle.getSchemaTableName()); + + checkProcedureArgument( + icebergTable.schemas().size() >= sourceTable.getDataColumns().size(), + "Target table should have at least %d columns but got %d", sourceTable.getDataColumns().size(), icebergTable.schemas().size()); + checkProcedureArgument( + icebergTable.spec().fields().size() == sourceTable.getPartitionColumns().size(), + "Numbers of partition columns should be equivalent. target: %d, source: %d", icebergTable.spec().fields().size(), sourceTable.getPartitionColumns().size()); + + // TODO Add files from all partitions when partition filter is not provided + checkProcedureArgument( + sourceTable.getPartitionColumns().isEmpty() || partitionFilter != null, + "partition_filter argument must be provided for partitioned tables"); + + String transactionalProperty = sourceTable.getParameters().get(TRANSACTIONAL); + if (parseBoolean(transactionalProperty)) { + throw new TrinoException(NOT_SUPPORTED, "Adding files from transactional tables is unsupported"); + } + if (!"MANAGED_TABLE".equalsIgnoreCase(sourceTable.getTableType()) && !"EXTERNAL_TABLE".equalsIgnoreCase(sourceTable.getTableType())) { + throw new TrinoException(NOT_SUPPORTED, "The procedure doesn't support adding files from %s table type".formatted(sourceTable.getTableType())); + } + if (isSomeKindOfAView(sourceTable) || isIcebergTable(sourceTable) || isDeltaLakeTable(sourceTable) || isHudiTable(sourceTable)) { + throw new TrinoException(NOT_SUPPORTED, "Adding files from non-Hive tables is unsupported"); + } + if (sourceTable.getPartitionColumns().isEmpty() && partitionFilter != null && !partitionFilter.isEmpty()) { + throw new TrinoException(NOT_SUPPORTED, "Partition filter is not supported for non-partitioned tables"); + } + + Set missingDataColumns = new HashSet<>(); + Stream.of(sourceTable.getDataColumns(), sourceTable.getPartitionColumns()) + .flatMap(List::stream) + .forEach(sourceColumn -> { + Types.NestedField targetColumn = icebergTable.schema().caseInsensitiveFindField(sourceColumn.getName()); + if (targetColumn == null) { + if (sourceTable.getPartitionColumns().contains(sourceColumn)) { + throw new TrinoException(COLUMN_NOT_FOUND, "Partition column '%s' does not exist".formatted(sourceColumn.getName())); + } + missingDataColumns.add(sourceColumn.getName()); + return; + } + ColumnIdentity columnIdentity = createColumnIdentity(targetColumn); + org.apache.iceberg.types.Type sourceColumnType = toIcebergType(typeManager.getType(getTypeSignature(sourceColumn.getType(), DEFAULT_PRECISION)), columnIdentity); + if (!targetColumn.type().equals(sourceColumnType)) { + throw new TrinoException(TYPE_MISMATCH, "Target '%s' column is '%s' type, but got source '%s' type".formatted(targetColumn.name(), targetColumn.type(), sourceColumnType)); + } + }); + if (missingDataColumns.size() == sourceTable.getDataColumns().size()) { + throw new TrinoException(COLUMN_NOT_FOUND, "All columns in the source table do not exist in the target table"); + } + + return Optional.of(new IcebergTableExecuteHandle( + tableHandle.getSchemaTableName(), + ADD_FILES_FROM_TABLE, + new IcebergAddFilesFromTableHandle(sourceTable, partitionFilter, recursiveDirectory), + icebergTable.location(), + icebergTable.io().properties())); + } + + private static Object requireProcedureArgument(Map properties, String name) + { + Object value = properties.get(name); + checkProcedureArgument(value != null, "Required procedure argument '%s' is missing", name); + return value; } @Override @@ -1171,9 +1673,13 @@ public Optional getLayoutForTableExecute(ConnectorSession switch (executeHandle.getProcedureId()) { case OPTIMIZE: return getLayoutForOptimize(session, executeHandle); + case OPTIMIZE_MANIFESTS: case DROP_EXTENDED_STATS: + case ROLLBACK_TO_SNAPSHOT: case EXPIRE_SNAPSHOTS: case REMOVE_ORPHAN_FILES: + case ADD_FILES: + case ADD_FILES_FROM_TABLE: // handled via executeTableExecute } throw new IllegalArgumentException("Unknown procedure '" + executeHandle.getProcedureId() + "'"); @@ -1198,9 +1704,13 @@ public BeginTableExecuteResult OPTIMIZE_MAX_SUPPORTED_TABLE_VERSION) { throw new TrinoException(NOT_SUPPORTED, format( "%s is not supported for Iceberg table format version > %d. Table %s format version is %s.", @@ -1242,9 +1751,13 @@ public void finishTableExecute(ConnectorSession session, ConnectorTableExecuteHa case OPTIMIZE: finishOptimize(session, executeHandle, fragments, splitSourceInfo); return; + case OPTIMIZE_MANIFESTS: case DROP_EXTENDED_STATS: + case ROLLBACK_TO_SNAPSHOT: case EXPIRE_SNAPSHOTS: case REMOVE_ORPHAN_FILES: + case ADD_FILES: + case ADD_FILES_FROM_TABLE: // handled via executeTableExecute } throw new IllegalArgumentException("Unknown procedure '" + executeHandle.getProcedureId() + "'"); @@ -1254,6 +1767,7 @@ private void finishOptimize(ConnectorSession session, IcebergTableExecuteHandle { IcebergOptimizeHandle optimizeHandle = (IcebergOptimizeHandle) executeHandle.getProcedureHandle(); Table icebergTable = transaction.table(); + Optional beforeWriteSnapshotId = getCurrentSnapshotId(icebergTable); // files to be deleted ImmutableSet.Builder scannedDataFilesBuilder = ImmutableSet.builder(); @@ -1267,7 +1781,8 @@ private void finishOptimize(ConnectorSession session, IcebergTableExecuteHandle Set fullyAppliedDeleteFiles = scannedDeleteFilesBuilder.build(); List commitTasks = fragments.stream() - .map(slice -> commitTaskCodec.fromJson(slice.getBytes())) + .map(Slice::getBytes) + .map(commitTaskCodec::fromJson) .collect(toImmutableList()); Type[] partitionColumnTypes = icebergTable.spec().fields().stream() @@ -1282,6 +1797,7 @@ private void finishOptimize(ConnectorSession session, IcebergTableExecuteHandle .withFileSizeInBytes(task.getFileSizeInBytes()) .withFormat(optimizeHandle.getFileFormat().toIceberg()) .withMetrics(task.getMetrics().metrics()); + task.getFileSplitOffsets().ifPresent(builder::withSplitOffsets); if (!icebergTable.spec().fields().isEmpty()) { String partitionDataJson = task.getPartitionDataJson() @@ -1303,18 +1819,71 @@ private void finishOptimize(ConnectorSession session, IcebergTableExecuteHandle cleanExtraOutputFiles( session, newFiles.stream() - .map(dataFile -> dataFile.path().toString()) + .map(ContentFile::location) .collect(toImmutableSet())); } RewriteFiles rewriteFiles = transaction.newRewrite(); - rewriteFiles.rewriteFiles(scannedDataFiles, fullyAppliedDeleteFiles, newFiles, ImmutableSet.of()); + scannedDataFiles.forEach(rewriteFiles::deleteFile); + fullyAppliedDeleteFiles.forEach(rewriteFiles::deleteFile); + newFiles.forEach(rewriteFiles::addFile); + // Table.snapshot method returns null if there is no matching snapshot Snapshot snapshot = requireNonNull(icebergTable.snapshot(optimizeHandle.getSnapshotId().get()), "snapshot is null"); + rewriteFiles.dataSequenceNumber(snapshot.sequenceNumber()); rewriteFiles.validateFromSnapshot(snapshot.snapshotId()); - commit(rewriteFiles, session); - transaction.commitTransaction(); + commitUpdateAndTransaction(rewriteFiles, session, transaction, "optimize"); + + // TODO (https://github.com/trinodb/trino/issues/15439) this may not exactly be the snapshot we committed, if there is another writer + long newSnapshotId = transaction.table().currentSnapshot().snapshotId(); transaction = null; + + // TODO (https://github.com/trinodb/trino/issues/15439): it would be good to publish data and stats atomically + beforeWriteSnapshotId.ifPresent(previous -> + verify(previous != newSnapshotId, "Failed to get new snapshot ID")); + + try { + beginTransaction(catalog.loadTable(session, executeHandle.getSchemaTableName())); + Table reloadedTable = transaction.table(); + StatisticsFile newStatsFile = tableStatisticsWriter.rewriteStatisticsFile(session, reloadedTable, newSnapshotId); + + transaction.updateStatistics() + .setStatistics(newSnapshotId, newStatsFile) + .commit(); + commitTransaction(transaction, "update statistics after optimize"); + } + catch (Exception e) { + // Write was committed, so at this point we cannot fail the query + // TODO (https://github.com/trinodb/trino/issues/15439): it would be good to publish data and stats atomically + log.error(e, "Failed to save table statistics"); + } + transaction = null; + } + + private static void commitUpdateAndTransaction(SnapshotUpdate update, ConnectorSession session, Transaction transaction, String operation) + { + commitUpdate(update, session, operation); + commitTransaction(transaction, operation); + } + + private static void commitUpdate(SnapshotUpdate update, ConnectorSession session, String operation) + { + try { + commit(update, session); + } + catch (UncheckedIOException | ValidationException | CommitFailedException | CommitStateUnknownException e) { + throw new TrinoException(ICEBERG_COMMIT_ERROR, format("Failed to commit during %s: %s", operation, firstNonNull(e.getMessage(), e)), e); + } + } + + private static void commitTransaction(Transaction transaction, String operation) + { + try { + transaction.commitTransaction(); + } + catch (UncheckedIOException | ValidationException | CommitFailedException | CommitStateUnknownException e) { + throw new TrinoException(ICEBERG_COMMIT_ERROR, format("Failed to commit the transaction during %s: %s", operation, firstNonNull(e.getMessage(), e)), e); + } } @Override @@ -1322,47 +1891,85 @@ public void executeTableExecute(ConnectorSession session, ConnectorTableExecuteH { IcebergTableExecuteHandle executeHandle = (IcebergTableExecuteHandle) tableExecuteHandle; switch (executeHandle.getProcedureId()) { + case OPTIMIZE_MANIFESTS: + executeOptimizeManifests(session, executeHandle); + return; case DROP_EXTENDED_STATS: executeDropExtendedStats(session, executeHandle); return; + case ROLLBACK_TO_SNAPSHOT: + executeRollbackToSnapshot(session, executeHandle); + return; case EXPIRE_SNAPSHOTS: executeExpireSnapshots(session, executeHandle); return; case REMOVE_ORPHAN_FILES: executeRemoveOrphanFiles(session, executeHandle); return; + case ADD_FILES: + executeAddFiles(session, executeHandle); + return; + case ADD_FILES_FROM_TABLE: + executeAddFilesFromTable(session, executeHandle); + return; default: throw new IllegalArgumentException("Unknown procedure '" + executeHandle.getProcedureId() + "'"); } } + private void executeOptimizeManifests(ConnectorSession session, IcebergTableExecuteHandle executeHandle) + { + checkArgument(executeHandle.getProcedureHandle() instanceof IcebergOptimizeManifestsHandle, "Unexpected procedure handle %s", executeHandle.getProcedureHandle()); + + BaseTable icebergTable = (BaseTable) catalog.loadTable(session, executeHandle.getSchemaTableName()); + List manifests = icebergTable.currentSnapshot().allManifests(icebergTable.io()); + if (manifests.isEmpty()) { + return; + } + if (manifests.size() == 1 && manifests.get(0).length() < icebergTable.operations().current().propertyAsLong(MANIFEST_TARGET_SIZE_BYTES, MANIFEST_TARGET_SIZE_BYTES_DEFAULT)) { + return; + } + + beginTransaction(icebergTable); + RewriteManifests rewriteManifests = transaction.rewriteManifests(); + rewriteManifests.clusterBy(file -> { + // Use the first partition field as the clustering key + StructLike partition = file.partition(); + return partition.size() > 1 ? partition.get(0, Object.class) : partition; + }).commit(); + commitTransaction(transaction, "optimize manifests"); + transaction = null; + } + private void executeDropExtendedStats(ConnectorSession session, IcebergTableExecuteHandle executeHandle) { checkArgument(executeHandle.getProcedureHandle() instanceof IcebergDropExtendedStatsHandle, "Unexpected procedure handle %s", executeHandle.getProcedureHandle()); - Table icebergTable = catalog.loadTable(session, executeHandle.getSchemaTableName()); + BaseTable icebergTable = catalog.loadTable(session, executeHandle.getSchemaTableName()); beginTransaction(icebergTable); UpdateStatistics updateStatistics = transaction.updateStatistics(); for (StatisticsFile statisticsFile : icebergTable.statisticsFiles()) { updateStatistics.removeStatistics(statisticsFile.snapshotId()); } updateStatistics.commit(); - UpdateProperties updateProperties = transaction.updateProperties(); - for (String key : transaction.table().properties().keySet()) { - if (key.startsWith(TRINO_STATS_PREFIX)) { - updateProperties.remove(key); - } - } - updateProperties.commit(); - transaction.commitTransaction(); + commitTransaction(transaction, "drop extended stats"); transaction = null; } + private void executeRollbackToSnapshot(ConnectorSession session, IcebergTableExecuteHandle executeHandle) + { + checkArgument(executeHandle.getProcedureHandle() instanceof IcebergRollbackToSnapshotHandle, "Unexpected procedure handle %s", executeHandle.getProcedureHandle()); + long snapshotId = ((IcebergRollbackToSnapshotHandle) executeHandle.getProcedureHandle()).snapshotId(); + + Table icebergTable = catalog.loadTable(session, executeHandle.getSchemaTableName()); + icebergTable.manageSnapshots().setCurrentSnapshot(snapshotId).commit(); + } + private void executeExpireSnapshots(ConnectorSession session, IcebergTableExecuteHandle executeHandle) { IcebergExpireSnapshotsHandle expireSnapshotsHandle = (IcebergExpireSnapshotsHandle) executeHandle.getProcedureHandle(); - Table table = catalog.loadTable(session, executeHandle.getSchemaTableName()); + BaseTable table = catalog.loadTable(session, executeHandle.getSchemaTableName()); Duration retention = requireNonNull(expireSnapshotsHandle.getRetentionThreshold(), "retention is null"); validateTableExecuteParameters( table, @@ -1374,36 +1981,11 @@ private void executeExpireSnapshots(ConnectorSession session, IcebergTableExecut IcebergSessionProperties.EXPIRE_SNAPSHOTS_MIN_RETENTION); long expireTimestampMillis = session.getStart().toEpochMilli() - retention.toMillis(); - TrinoFileSystem fileSystem = fileSystemFactory.create(session); - List pathsToDelete = new ArrayList<>(); - // deleteFunction is not accessed from multiple threads unless .executeDeleteWith() is used - Consumer deleteFunction = path -> { - pathsToDelete.add(Location.of(path)); - if (pathsToDelete.size() == DELETE_BATCH_SIZE) { - try { - fileSystem.deleteFiles(pathsToDelete); - pathsToDelete.clear(); - } - catch (IOException e) { - throw new TrinoException(ICEBERG_FILESYSTEM_ERROR, "Failed to delete files during snapshot expiration", e); - } - } - }; - - table.expireSnapshots() - .expireOlderThan(expireTimestampMillis) - .deleteWith(deleteFunction) - .commit(); - try { - fileSystem.deleteFiles(pathsToDelete); - } - catch (IOException e) { - throw new TrinoException(ICEBERG_FILESYSTEM_ERROR, "Failed to delete files during snapshot expiration", e); - } + executeExpireSnapshots(table, session, expireTimestampMillis); } private static void validateTableExecuteParameters( - Table table, + BaseTable table, SchemaTableName schemaTableName, String procedureName, Duration retentionThreshold, @@ -1411,7 +1993,7 @@ private static void validateTableExecuteParameters( String minRetentionParameterName, String sessionMinRetentionParameterName) { - int tableFormatVersion = ((BaseTable) table).operations().current().formatVersion(); + int tableFormatVersion = table.operations().current().formatVersion(); if (tableFormatVersion > CLEANING_UP_PROCEDURES_MAX_SUPPORTED_TABLE_VERSION) { // It is not known if future version won't bring any new kind of metadata or data files // because of the way procedures are implemented it is safer to fail here than to potentially remove @@ -1443,7 +2025,7 @@ public void executeRemoveOrphanFiles(ConnectorSession session, IcebergTableExecu { IcebergRemoveOrphanFilesHandle removeOrphanFilesHandle = (IcebergRemoveOrphanFilesHandle) executeHandle.getProcedureHandle(); - Table table = catalog.loadTable(session, executeHandle.getSchemaTableName()); + BaseTable table = catalog.loadTable(session, executeHandle.getSchemaTableName()); Duration retention = requireNonNull(removeOrphanFilesHandle.getRetentionThreshold(), "retention is null"); validateTableExecuteParameters( table, @@ -1460,10 +2042,10 @@ public void executeRemoveOrphanFiles(ConnectorSession session, IcebergTableExecu } Instant expiration = session.getStart().minusMillis(retention.toMillis()); - removeOrphanFiles(table, session, executeHandle.getSchemaTableName(), expiration); + removeOrphanFiles(table, session, executeHandle.getSchemaTableName(), expiration, executeHandle.getFileIoProperties()); } - private void removeOrphanFiles(Table table, ConnectorSession session, SchemaTableName schemaTableName, Instant expiration) + private void removeOrphanFiles(Table table, ConnectorSession session, SchemaTableName schemaTableName, Instant expiration, Map fileIoProperties) { Set processedManifestFilePaths = new HashSet<>(); // Similarly to issues like https://github.com/trinodb/trino/issues/13759, equivalent paths may have different String @@ -1498,10 +2080,45 @@ private void removeOrphanFiles(Table table, ConnectorSession session, SchemaTabl metadataFileLocations(table, false).stream() .map(IcebergUtil::fileName) .forEach(validMetadataFileNames::add); - validMetadataFileNames.add(fileName(versionHintLocation(table))); - scanAndDeleteInvalidFiles(table, session, schemaTableName, expiration, validDataFileNames.build(), "data"); - scanAndDeleteInvalidFiles(table, session, schemaTableName, expiration, validMetadataFileNames.build(), "metadata"); + statisticsFilesLocations(table).stream() + .map(IcebergUtil::fileName) + .forEach(validMetadataFileNames::add); + + validMetadataFileNames.add("version-hint.text"); + + scanAndDeleteInvalidFiles(table, session, schemaTableName, expiration, validDataFileNames.build(), "data", fileIoProperties); + scanAndDeleteInvalidFiles(table, session, schemaTableName, expiration, validMetadataFileNames.build(), "metadata", fileIoProperties); + } + + public void executeAddFiles(ConnectorSession session, IcebergTableExecuteHandle executeHandle) + { + IcebergAddFilesHandle addFilesHandle = (IcebergAddFilesHandle) executeHandle.getProcedureHandle(); + Table table = catalog.loadTable(session, executeHandle.getSchemaTableName()); + TrinoFileSystem fileSystem = fileSystemFactory.create(session.getIdentity(), table.io().properties()); + addFiles( + session, + fileSystem, + catalog, + executeHandle.getSchemaTableName(), + addFilesHandle.location(), + addFilesHandle.format(), + addFilesHandle.recursiveDirectory()); + } + + public void executeAddFilesFromTable(ConnectorSession session, IcebergTableExecuteHandle executeHandle) + { + IcebergAddFilesFromTableHandle addFilesHandle = (IcebergAddFilesFromTableHandle) executeHandle.getProcedureHandle(); + Table table = catalog.loadTable(session, executeHandle.getSchemaTableName()); + TrinoFileSystem fileSystem = fileSystemFactory.create(session.getIdentity(), table.io().properties()); + addFilesFromTable( + session, + fileSystem, + metastoreFactory.orElseThrow(), + table, + addFilesHandle.table(), + addFilesHandle.partitionFilter(), + addFilesHandle.recursiveDirectory()); } private static ManifestReader> readerForManifest(Table table, ManifestFile manifest) @@ -1512,11 +2129,11 @@ private static ManifestReader> readerForManifest(Table }; } - private void scanAndDeleteInvalidFiles(Table table, ConnectorSession session, SchemaTableName schemaTableName, Instant expiration, Set validFiles, String subfolder) + private void scanAndDeleteInvalidFiles(Table table, ConnectorSession session, SchemaTableName schemaTableName, Instant expiration, Set validFiles, String subfolder, Map fileIoProperties) { try { List filesToDelete = new ArrayList<>(); - TrinoFileSystem fileSystem = fileSystemFactory.create(session); + TrinoFileSystem fileSystem = fileSystemFactory.create(session.getIdentity(), fileIoProperties); FileIterator allFiles = fileSystem.listFiles(Location.of(table.location()).appendPath(subfolder)); while (allFiles.hasNext()) { FileEntry entry = allFiles.next(); @@ -1546,13 +2163,28 @@ private void scanAndDeleteInvalidFiles(Table table, ConnectorSession session, Sc public Optional getInfo(ConnectorTableHandle tableHandle) { IcebergTableHandle icebergTableHandle = (IcebergTableHandle) tableHandle; - Optional partitioned = icebergTableHandle.getPartitionSpecJson() - .map(partitionSpecJson -> PartitionSpecParser.fromJson(SchemaParser.fromJson(icebergTableHandle.getTableSchemaJson()), partitionSpecJson).isPartitioned()); + List partitionFields = icebergTableHandle.getPartitionSpecJson() + .map(partitionSpecJson -> PartitionSpecParser.fromJson(SchemaParser.fromJson(icebergTableHandle.getTableSchemaJson()), partitionSpecJson) + .fields().stream() + .map(field -> field.name() + ": " + field.transform()) + .collect(toImmutableList())) + .orElse(ImmutableList.of()); + + // FIXME: Cannot build summary without loading the table + Map summary = ImmutableMap.of(); + Optional totalRecords = Optional.ofNullable(summary.get(TOTAL_RECORDS_PROP)); + Optional deletedRecords = Optional.ofNullable(summary.get(DELETED_RECORDS_PROP)); + Optional totalDataFiles = Optional.ofNullable(summary.get(TOTAL_DATA_FILES_PROP)); + Optional totalDeleteFiles = Optional.ofNullable(summary.get(TOTAL_DELETE_FILES_PROP)); return Optional.of(new IcebergInputInfo( icebergTableHandle.getSnapshotId(), - partitioned, - getFileFormat(icebergTableHandle.getStorageProperties()).name())); + partitionFields, + getFileFormat(icebergTableHandle.getStorageProperties()).name(), + totalRecords, + deletedRecords, + totalDataFiles, + totalDeleteFiles)); } @Override @@ -1587,6 +2219,44 @@ public void setTableProperties(ConnectorSession session, ConnectorTableHandle ta beginTransaction(icebergTable); UpdateProperties updateProperties = transaction.updateProperties(); + if (properties.containsKey(EXTRA_PROPERTIES_PROPERTY)) { + //noinspection unchecked + Map extraProperties = (Map) properties.get(EXTRA_PROPERTIES_PROPERTY) + .orElseThrow(() -> new IllegalArgumentException("The extra_properties property cannot be empty")); + verifyExtraProperties(properties.keySet(), extraProperties, allowedExtraProperties); + extraProperties.forEach(updateProperties::set); + } + + if (properties.containsKey(PARQUET_BLOOM_FILTER_COLUMNS_PROPERTY)) { + checkFormatForProperty(getFileFormat(icebergTable).toIceberg(), FileFormat.PARQUET, PARQUET_BLOOM_FILTER_COLUMNS_PROPERTY); + //noinspection unchecked + List parquetBloomFilterColumns = (List) properties.get(PARQUET_BLOOM_FILTER_COLUMNS_PROPERTY) + .orElseThrow(() -> new IllegalArgumentException("The parquet_bloom_filter_columns property cannot be empty")); + validateParquetBloomFilterColumns(getColumnMetadatas(SchemaParser.fromJson(table.getTableSchemaJson()), typeManager), parquetBloomFilterColumns); + + Set existingParquetBloomFilterColumns = icebergTable.properties().keySet().stream() + .filter(key -> key.startsWith(PARQUET_BLOOM_FILTER_COLUMN_ENABLED_PREFIX)) + .map(key -> key.substring(PARQUET_BLOOM_FILTER_COLUMN_ENABLED_PREFIX.length())) + .collect(toImmutableSet()); + Set removeParquetBloomFilterColumns = Sets.difference(existingParquetBloomFilterColumns, Set.copyOf(parquetBloomFilterColumns)); + removeParquetBloomFilterColumns.forEach(column -> updateProperties.remove(PARQUET_BLOOM_FILTER_COLUMN_ENABLED_PREFIX + column)); + parquetBloomFilterColumns.forEach(column -> updateProperties.set(PARQUET_BLOOM_FILTER_COLUMN_ENABLED_PREFIX + column, "true")); + } + + if (properties.containsKey(ORC_BLOOM_FILTER_COLUMNS_PROPERTY)) { + checkFormatForProperty(getFileFormat(icebergTable).toIceberg(), FileFormat.ORC, ORC_BLOOM_FILTER_COLUMNS_PROPERTY); + //noinspection unchecked + List orcBloomFilterColumns = (List) properties.get(ORC_BLOOM_FILTER_COLUMNS_PROPERTY) + .orElseThrow(() -> new IllegalArgumentException("The orc_bloom_filter_columns property cannot be empty")); + if (orcBloomFilterColumns.isEmpty()) { + updateProperties.remove(ORC_BLOOM_FILTER_COLUMNS); + } + else { + validateOrcBloomFilterColumns(getColumnMetadatas(SchemaParser.fromJson(table.getTableSchemaJson()), typeManager), orcBloomFilterColumns); + updateProperties.set(ORC_BLOOM_FILTER_COLUMNS, Joiner.on(",").join(orcBloomFilterColumns)); + } + } + if (properties.containsKey(FILE_FORMAT_PROPERTY)) { IcebergFileFormat fileFormat = (IcebergFileFormat) properties.get(FILE_FORMAT_PROPERTY) .orElseThrow(() -> new IllegalArgumentException("The format property cannot be empty")); @@ -1600,6 +2270,30 @@ public void setTableProperties(ConnectorSession session, ConnectorTableHandle ta updateProperties.set(FORMAT_VERSION, Integer.toString(formatVersion)); } + if (properties.containsKey(MAX_COMMIT_RETRY)) { + int formatVersion = (int) properties.get(MAX_COMMIT_RETRY) + .orElseThrow(() -> new IllegalArgumentException("The max_commit_retry property cannot be empty")); + updateProperties.set(COMMIT_NUM_RETRIES, Integer.toString(formatVersion)); + } + + if (properties.containsKey(OBJECT_STORE_LAYOUT_ENABLED_PROPERTY)) { + boolean objectStoreEnabled = (boolean) properties.get(OBJECT_STORE_LAYOUT_ENABLED_PROPERTY) + .orElseThrow(() -> new IllegalArgumentException("The object_store_enabled property cannot be empty")); + updateProperties.set(OBJECT_STORE_ENABLED, Boolean.toString(objectStoreEnabled)); + } + + if (properties.containsKey(DATA_LOCATION_PROPERTY)) { + String dataLocation = (String) properties.get(DATA_LOCATION_PROPERTY) + .orElseThrow(() -> new IllegalArgumentException("The data_location property cannot be empty")); + boolean objectStoreEnabled = (boolean) properties.getOrDefault( + OBJECT_STORE_LAYOUT_ENABLED_PROPERTY, + Optional.of(Boolean.parseBoolean(icebergTable.properties().get(OBJECT_STORE_ENABLED)))).orElseThrow(); + if (!objectStoreEnabled) { + throw new TrinoException(INVALID_TABLE_PROPERTY, "Data location can only be set when object store layout is enabled"); + } + updateProperties.set(WRITE_DATA_LOCATION, dataLocation); + } + try { updateProperties.commit(); } @@ -1628,12 +2322,7 @@ public void setTableProperties(ConnectorSession session, ConnectorTableHandle ta } } - try { - transaction.commitTransaction(); - } - catch (RuntimeException e) { - throw new TrinoException(ICEBERG_COMMIT_ERROR, "Failed to commit new table properties", e); - } + commitTransaction(transaction, "set table properties"); } private static void updatePartitioning(Table icebergTable, Transaction transaction, List partitionColumns) @@ -1648,14 +2337,12 @@ private static void updatePartitioning(Table icebergTable, Transaction transacti } else { PartitionSpec partitionSpec = parsePartitionFields(schema, partitionColumns); - validateNotPartitionedByNestedField(schema, partitionSpec); Set partitionFields = ImmutableSet.copyOf(partitionSpec.fields()); difference(existingPartitionFields, partitionFields).stream() .map(PartitionField::name) .forEach(updatePartitionSpec::removeField); - difference(partitionFields, existingPartitionFields).stream() - .map(partitionField -> toIcebergTerm(schema, partitionField)) - .forEach(updatePartitionSpec::addField); + difference(partitionFields, existingPartitionFields) + .forEach(partitionField -> updatePartitionSpec.addField(partitionField.name(), toIcebergTerm(schema, partitionField))); } try { @@ -1686,9 +2373,9 @@ public void addColumn(ConnectorSession session, ConnectorTableHandle tableHandle // added - instead of relying on addColumn in iceberg library to assign Ids AtomicInteger nextFieldId = new AtomicInteger(icebergTable.schema().highestFieldId() + 2); try { - icebergTable.updateSchema() - .addColumn(column.getName(), toIcebergTypeForNewColumn(column.getType(), nextFieldId), column.getComment()) - .commit(); + UpdateSchema updateSchema = icebergTable.updateSchema(); + updateSchema.addColumn(null, column.getName(), toIcebergTypeForNewColumn(column.getType(), nextFieldId), column.getComment()); + updateSchema.commit(); } catch (RuntimeException e) { throw new TrinoException(ICEBERG_COMMIT_ERROR, "Failed to add column: " + firstNonNull(e.getMessage(), e), e); @@ -1705,7 +2392,18 @@ public void addField(ConnectorSession session, ConnectorTableHandle tableHandle, NestedField parent = icebergTable.schema().caseInsensitiveFindField(parentName); String caseSensitiveParentName = icebergTable.schema().findColumnName(parent.fieldId()); - NestedField field = parent.type().asStructType().caseInsensitiveField(fieldName); + + Types.StructType structType; + if (parent.type().isListType()) { + // list(struct...) + structType = parent.type().asListType().elementType().asStructType(); + } + else { + // just struct + structType = parent.type().asStructType(); + } + + NestedField field = structType.caseInsensitiveField(fieldName); if (field != null) { if (ignoreExisting) { return; @@ -1887,7 +2585,19 @@ public void setFieldType(ConnectorSession session, ConnectorTableHandle tableHan NestedField parent = icebergTable.schema().caseInsensitiveFindField(parentPath); String caseSensitiveParentName = icebergTable.schema().findColumnName(parent.fieldId()); - NestedField field = parent.type().asStructType().caseInsensitiveField(getLast(fieldPath)); + + Types.StructType structType; + if (parent.type().isListType()) { + // list(struct...) + structType = parent.type().asListType().elementType().asStructType(); + caseSensitiveParentName += ".element"; + } + else { + // just struct + structType = parent.type().asStructType(); + } + NestedField field = structType.caseInsensitiveField(getLast(fieldPath)); + // TODO: Add support for changing non-primitive field type if (!field.type().isPrimitiveType()) { throw new TrinoException(NOT_SUPPORTED, "Iceberg doesn't support changing field type from non-primitive types"); @@ -1918,17 +2628,34 @@ public TableStatisticsMetadata getStatisticsCollectionMetadataForWrite(Connector ConnectorTableHandle tableHandle = getTableHandle(session, tableMetadata.getTable(), Optional.empty(), Optional.empty()); if (tableHandle == null) { - // Assume new table (CTAS), collect all stats possible + // Assume new table (CTAS), collect NDV stats on all columns return getStatisticsCollectionMetadata(tableMetadata, Optional.empty(), availableColumnNames -> {}); } - TableStatistics tableStatistics = getTableStatistics(session, checkValidTableHandle(tableHandle)); - if (tableStatistics.getRowCount().getValue() == 0.0) { - // Table has no data (empty, or wiped out). Collect all stats possible + IcebergTableHandle table = checkValidTableHandle(tableHandle); + if (table.getSnapshotId().isEmpty()) { + // Table has no data (empty, or wiped out). Collect NDV stats on all columns return getStatisticsCollectionMetadata(tableMetadata, Optional.empty(), availableColumnNames -> {}); } - Set columnsWithExtendedStatistics = tableStatistics.getColumnStatistics().entrySet().stream() - .filter(entry -> !entry.getValue().getDistinctValuesCount().isUnknown()) - .map(entry -> ((IcebergColumnHandle) entry.getKey()).getName()) + + Table icebergTable = catalog.loadTable(session, table.getSchemaTableName()); + long snapshotId = table.getSnapshotId().orElseThrow(); + Snapshot snapshot = icebergTable.snapshot(snapshotId); + String totalRecords = snapshot.summary().get(TOTAL_RECORDS_PROP); + if (totalRecords != null && Long.parseLong(totalRecords) == 0) { + // Table has no data (empty, or wiped out). Collect NDV stats on all columns + return getStatisticsCollectionMetadata(tableMetadata, Optional.empty(), availableColumnNames -> {}); + } + + Schema schema = SchemaParser.fromJson(table.getTableSchemaJson()); + List columns = getTopLevelColumns(schema, typeManager); + Set columnIds = columns.stream() + .map(IcebergColumnHandle::getId) + .collect(toImmutableSet()); + Map ndvs = readNdvs(icebergTable, snapshotId, columnIds, true); + // Avoid collecting NDV stats on columns where we don't know the existing NDV count + Set columnsWithExtendedStatistics = columns.stream() + .filter(column -> ndvs.containsKey(column.getId())) + .map(IcebergColumnHandle::getName) .collect(toImmutableSet()); return getStatisticsCollectionMetadata(tableMetadata, Optional.of(columnsWithExtendedStatistics), availableColumnNames -> {}); } @@ -1961,7 +2688,7 @@ public ConnectorAnalyzeMetadata getStatisticsCollectionMetadata(ConnectorSession }); return new ConnectorAnalyzeMetadata( - tableHandle, + handle.forAnalyze(), getStatisticsCollectionMetadata( tableMetadata, analyzeColumnNames, @@ -2005,6 +2732,9 @@ public ConnectorTableHandle beginStatisticsCollection(ConnectorSession session, { IcebergTableHandle handle = (IcebergTableHandle) tableHandle; Table icebergTable = catalog.loadTable(session, handle.getSchemaTableName()); + if (isS3Tables(icebergTable.location())) { + throw new TrinoException(NOT_SUPPORTED, "S3 Tables does not support analyze"); + } beginTransaction(icebergTable); return handle; } @@ -2017,45 +2747,23 @@ public void finishStatisticsCollection(ConnectorSession session, ConnectorTableH if (handle.getSnapshotId().isEmpty()) { // No snapshot, table is empty verify( - computedStatistics.isEmpty(), - "Unexpected computed statistics that cannot be attached to a snapshot because none exists: %s", + computedStatistics.size() == 1, + "The computedStatistics size must be 1: %s", + computedStatistics); + ComputedStatistics statistics = getOnlyElement(computedStatistics); + verify(statistics.getGroupingColumns().isEmpty() && + statistics.getGroupingValues().isEmpty() && + statistics.getColumnStatistics().isEmpty() && + statistics.getTableStatistics().isEmpty(), + "Unexpected non-empty statistics that cannot be attached to a snapshot because none exists: %s", computedStatistics); - // TODO (https://github.com/trinodb/trino/issues/15397): remove support for Trino-specific statistics properties - // Drop all stats. Empty table needs none - UpdateProperties updateProperties = transaction.updateProperties(); - table.properties().keySet().stream() - .filter(key -> key.startsWith(TRINO_STATS_PREFIX)) - .forEach(updateProperties::remove); - updateProperties.commit(); - - transaction.commitTransaction(); + commitTransaction(transaction, "statistics collection"); transaction = null; return; } long snapshotId = handle.getSnapshotId().orElseThrow(); - Set columnIds = table.schema().columns().stream() - .map(Types.NestedField::fieldId) - .collect(toImmutableSet()); - - // TODO (https://github.com/trinodb/trino/issues/15397): remove support for Trino-specific statistics properties - // Drop stats for obsolete columns - UpdateProperties updateProperties = transaction.updateProperties(); - table.properties().keySet().stream() - .filter(key -> { - if (!key.startsWith(TRINO_STATS_PREFIX)) { - return false; - } - Matcher matcher = TRINO_STATS_COLUMN_ID_PATTERN.matcher(key); - if (!matcher.matches()) { - return false; - } - return !columnIds.contains(Integer.parseInt(matcher.group("columnId"))); - }) - .forEach(updateProperties::remove); - updateProperties.commit(); - CollectedStatistics collectedStatistics = processComputedTableStatistics(table, computedStatistics); StatisticsFile statisticsFile = tableStatisticsWriter.writeStatisticsFile( session, @@ -2067,7 +2775,7 @@ public void finishStatisticsCollection(ConnectorSession session, ConnectorTableH .setStatistics(snapshotId, statisticsFile) .commit(); - transaction.commitTransaction(); + commitTransaction(transaction, "statistics collection"); transaction = null; } @@ -2105,7 +2813,10 @@ public ColumnHandle getMergeRowIdColumnHandle(ConnectorSession session, Connecto @Override public Optional getUpdateLayout(ConnectorSession session, ConnectorTableHandle tableHandle) { - return Optional.of(IcebergUpdateHandle.INSTANCE); + return getInsertLayout(session, tableHandle) + .flatMap(ConnectorTableLayout::getPartitioning) + .map(IcebergPartitioningHandle.class::cast) + .map(IcebergPartitioningHandle::forUpdate); } @Override @@ -2116,7 +2827,6 @@ public ConnectorMergeTableHandle beginMerge(ConnectorSession session, ConnectorT Table icebergTable = catalog.loadTable(session, table.getSchemaTableName()); validateNotModifyingOldSnapshot(table, icebergTable); - validateNotPartitionedByNestedField(icebergTable.schema(), icebergTable.spec()); beginTransaction(icebergTable); @@ -2162,7 +2872,8 @@ private void finishWrite(ConnectorSession session, IcebergTableHandle table, Col Table icebergTable = transaction.table(); List commitTasks = fragments.stream() - .map(slice -> commitTaskCodec.fromJson(slice.getBytes())) + .map(Slice::getBytes) + .map(commitTaskCodec::fromJson) .collect(toImmutableList()); if (commitTasks.isEmpty()) { @@ -2176,8 +2887,15 @@ private void finishWrite(ConnectorSession session, IcebergTableHandle table, Col RowDelta rowDelta = transaction.newRowDelta(); table.getSnapshotId().map(icebergTable::snapshot).ifPresent(s -> rowDelta.validateFromSnapshot(s.snapshotId())); TupleDomain dataColumnPredicate = table.getEnforcedPredicate().filter((column, domain) -> !isMetadataColumnId(column.getId())); - if (!dataColumnPredicate.isAll()) { - rowDelta.conflictDetectionFilter(toIcebergExpression(dataColumnPredicate)); + TupleDomain effectivePredicate = dataColumnPredicate.intersect(table.getUnenforcedPredicate()); + if (isFileBasedConflictDetectionEnabled(session)) { + effectivePredicate = effectivePredicate.intersect(extractTupleDomainsFromCommitTasks(table, icebergTable, commitTasks, typeManager)); + } + + effectivePredicate = effectivePredicate.filter((ignore, domain) -> isConvertibleToIcebergExpression(domain)); + + if (!effectivePredicate.isAll()) { + rowDelta.conflictDetectionFilter(toIcebergExpression(effectivePredicate)); } IsolationLevel isolationLevel = IsolationLevel.fromName(icebergTable.properties().getOrDefault(DELETE_ISOLATION_LEVEL, DELETE_ISOLATION_LEVEL_DEFAULT)); if (isolationLevel == IsolationLevel.SERIALIZABLE) { @@ -2203,6 +2921,7 @@ private void finishWrite(ConnectorSession session, IcebergTableHandle table, Col .ofPositionDeletes() .withFileSizeInBytes(task.getFileSizeInBytes()) .withMetrics(task.getMetrics().metrics()); + task.getFileSplitOffsets().ifPresent(deleteBuilder::withSplitOffsets); if (!partitionSpec.fields().isEmpty()) { String partitionDataJson = task.getPartitionDataJson() .orElseThrow(() -> new VerifyException("No partition data for partitioned table")); @@ -2236,13 +2955,41 @@ private void finishWrite(ConnectorSession session, IcebergTableHandle table, Col } rowDelta.validateDataFilesExist(referencedDataFiles.build()); - try { - commit(rowDelta, session); - transaction.commitTransaction(); - } - catch (ValidationException e) { - throw new TrinoException(ICEBERG_COMMIT_ERROR, "Failed to commit Iceberg update to table: " + table.getSchemaTableName(), e); + commitUpdateAndTransaction(rowDelta, session, transaction, "write"); + } + + static TupleDomain extractTupleDomainsFromCommitTasks(IcebergTableHandle table, Table icebergTable, List commitTasks, TypeManager typeManager) + { + Set partitionColumns = new HashSet<>(getProjectedColumns(icebergTable.schema(), typeManager, identityPartitionColumnsInAllSpecs(icebergTable))); + PartitionSpec partitionSpec = icebergTable.spec(); + Type[] partitionColumnTypes = partitionSpec.fields().stream() + .map(field -> field.transform().getResultType(icebergTable.schema().findType(field.sourceId()))) + .toArray(Type[]::new); + Schema schema = SchemaParser.fromJson(table.getTableSchemaJson()); + Map> domainsFromTasks = new HashMap<>(); + for (CommitTaskData commitTask : commitTasks) { + PartitionSpec taskPartitionSpec = PartitionSpecParser.fromJson(schema, commitTask.getPartitionSpecJson()); + if (commitTask.getPartitionDataJson().isEmpty() || taskPartitionSpec.isUnpartitioned() || !taskPartitionSpec.equals(partitionSpec)) { + // We should not produce any specific domains if there are no partitions or current partitions does not match task partitions for any of tasks + // As each partition value narrows down conflict scope we should produce values from all commit tasks or not at all, to avoid partial information + return TupleDomain.all(); + } + + PartitionData partitionData = PartitionData.fromJson(commitTask.getPartitionDataJson().get(), partitionColumnTypes); + Map> partitionKeys = getPartitionKeys(partitionData, partitionSpec); + Map partitionValues = getPartitionValues(partitionColumns, partitionKeys); + + for (Map.Entry entry : partitionValues.entrySet()) { + IcebergColumnHandle columnHandle = (IcebergColumnHandle) entry.getKey(); + NullableValue value = entry.getValue(); + Domain newDomain = value.isNull() ? Domain.onlyNull(columnHandle.getType()) : Domain.singleValue(columnHandle.getType(), value.getValue()); + domainsFromTasks.computeIfAbsent(columnHandle, ignore -> new ArrayList<>()).add(newDomain); + } } + return withColumnDomains(domainsFromTasks.entrySet().stream() + .collect(toImmutableMap( + Map.Entry::getKey, + entry -> Domain.union(entry.getValue())))); } @Override @@ -2284,7 +3031,16 @@ public Map getViews(ConnectorSession s @Override public Optional getView(ConnectorSession session, SchemaTableName viewName) { - return catalog.getView(session, viewName); + try { + return catalog.getView(session, viewName); + } + catch (TrinoException e) { + if (e.getErrorCode().equals(ICEBERG_UNSUPPORTED_VIEW_DIALECT.toErrorCode())) { + log.debug(e, "Skip unsupported view dialect: %s", viewName); + return Optional.empty(); + } + throw e; + } } @Override @@ -2296,7 +3052,7 @@ public OptionalLong executeDelete(ConnectorSession session, ConnectorTableHandle DeleteFiles deleteFiles = icebergTable.newDelete() .deleteFromRowFilter(toIcebergExpression(handle.getEnforcedPredicate())); - commit(deleteFiles, session); + commitUpdate(deleteFiles, session, "delete"); Map summary = icebergTable.currentSnapshot().summary(); String deletedRowsStr = summary.get(DELETED_RECORDS_PROP); @@ -2343,8 +3099,11 @@ public Optional> applyLimit(Connect table.getNameMappingJson(), table.getTableLocation(), table.getStorageProperties(), + table.getTablePartitioning(), table.isRecordScannedFiles(), - table.getMaxScannedFileSize()); + table.getMaxScannedFileSize(), + table.getConstraintColumns(), + table.getForAnalyze()); return Optional.of(new LimitApplicationResult<>(table, false, false)); } @@ -2354,8 +3113,9 @@ public Optional> applyFilter(C { IcebergTableHandle table = (IcebergTableHandle) handle; ConstraintExtractor.ExtractionResult extractionResult = extractTupleDomain(constraint); - TupleDomain predicate = extractionResult.tupleDomain(); - if (predicate.isAll()) { + TupleDomain predicate = extractionResult.tupleDomain() + .transformKeys(IcebergColumnHandle.class::cast); + if (predicate.isAll() && constraint.getPredicateColumns().isEmpty()) { return Optional.empty(); } if (table.getLimit().isPresent()) { @@ -2374,12 +3134,19 @@ public Optional> applyFilter(C remainingConstraint = TupleDomain.all(); } else { - Table icebergTable = catalog.loadTable(session, table.getSchemaTableName()); + BaseTable icebergTable = catalog.loadTable(session, table.getSchemaTableName()); Set partitionSpecIds = table.getSnapshotId().map( - snapshot -> icebergTable.snapshot(snapshot).allManifests(icebergTable.io()).stream() - .map(ManifestFile::partitionSpecId) - .collect(toImmutableSet())) + snapshot -> { + try { + return icebergTable.snapshot(snapshot).allManifests(icebergTable.io()).stream() + .map(ManifestFile::partitionSpecId) + .collect(toImmutableSet()); + } + catch (NotFoundException | UncheckedIOException e) { + throw new TrinoException(ICEBERG_INVALID_METADATA, "Error accessing manifest file for table %s".formatted(icebergTable.name()), e); + } + }) // No snapshot, so no data. This case doesn't matter. .orElseGet(() -> ImmutableSet.copyOf(icebergTable.specs().keySet())); @@ -2388,17 +3155,14 @@ public Optional> applyFilter(C Map newUnenforced = new LinkedHashMap<>(); Map domains = predicate.getDomains().orElseThrow(() -> new VerifyException("No domains")); domains.forEach((columnHandle, domain) -> { - // structural types cannot be used to filter a table scan in Iceberg library. - if (isStructuralType(columnHandle.getType()) || - // Iceberg orders UUID values differently than Trino (perhaps due to https://bugs.openjdk.org/browse/JDK-7025832), so allow only IS NULL / IS NOT NULL checks - (columnHandle.getType() == UUID && !(domain.isOnlyNull() || domain.getValues().isAll()))) { + if (!isConvertibleToIcebergExpression(domain)) { unsupported.put(columnHandle, domain); } else if (canEnforceColumnConstraintInSpecs(typeManager.getTypeOperators(), icebergTable, partitionSpecIds, columnHandle, domain)) { newEnforced.put(columnHandle, domain); } else if (isMetadataColumnId(columnHandle.getId())) { - if (columnHandle.isPathColumn() || columnHandle.isFileModifiedTimeColumn()) { + if (columnHandle.isPartitionColumn() || columnHandle.isPathColumn() || columnHandle.isFileModifiedTimeColumn()) { newEnforced.put(columnHandle, domain); } else { @@ -2410,13 +3174,20 @@ else if (isMetadataColumnId(columnHandle.getId())) { } }); - newEnforcedConstraint = TupleDomain.withColumnDomains(newEnforced).intersect(table.getEnforcedPredicate()); - newUnenforcedConstraint = TupleDomain.withColumnDomains(newUnenforced).intersect(table.getUnenforcedPredicate()); - remainingConstraint = TupleDomain.withColumnDomains(newUnenforced).intersect(TupleDomain.withColumnDomains(unsupported)); + newEnforcedConstraint = withColumnDomains(newEnforced).intersect(table.getEnforcedPredicate()); + newUnenforcedConstraint = withColumnDomains(newUnenforced).intersect(table.getUnenforcedPredicate()); + remainingConstraint = withColumnDomains(newUnenforced).intersect(withColumnDomains(unsupported)); } + Set newConstraintColumns = Streams.concat( + table.getConstraintColumns().stream(), + constraint.getPredicateColumns().orElseGet(ImmutableSet::of).stream() + .map(columnHandle -> (IcebergColumnHandle) columnHandle)) + .collect(toImmutableSet()); + if (newEnforcedConstraint.equals(table.getEnforcedPredicate()) - && newUnenforcedConstraint.equals(table.getUnenforcedPredicate())) { + && newUnenforcedConstraint.equals(table.getUnenforcedPredicate()) + && newConstraintColumns.equals(table.getConstraintColumns())) { return Optional.empty(); } @@ -2437,8 +3208,11 @@ else if (isMetadataColumnId(columnHandle.getId())) { table.getNameMappingJson(), table.getTableLocation(), table.getStorageProperties(), + table.getTablePartitioning(), table.isRecordScannedFiles(), - table.getMaxScannedFileSize()), + table.getMaxScannedFileSize(), + newConstraintColumns, + table.getForAnalyze()), remainingConstraint.transformKeys(ColumnHandle.class::cast), extractionResult.remainingExpression(), false)); @@ -2552,6 +3326,7 @@ private static IcebergColumnHandle createProjectedColumnHandle(IcebergColumnHand column.getBaseType(), fullPath.build(), projectedColumnType, + true, Optional.empty()); } @@ -2568,29 +3343,42 @@ public TableStatistics getTableStatistics(ConnectorSession session, ConnectorTab checkArgument(!originalHandle.isRecordScannedFiles(), "Unexpected scanned files recording set"); checkArgument(originalHandle.getMaxScannedFileSize().isEmpty(), "Unexpected max scanned file size set"); - return tableStatisticsCache.computeIfAbsent( - new IcebergTableHandle( - originalHandle.getCatalog(), - originalHandle.getSchemaName(), - originalHandle.getTableName(), - originalHandle.getTableType(), - originalHandle.getSnapshotId(), - originalHandle.getTableSchemaJson(), - originalHandle.getPartitionSpecJson(), - originalHandle.getFormatVersion(), - originalHandle.getUnenforcedPredicate(), - originalHandle.getEnforcedPredicate(), - OptionalLong.empty(), // limit is currently not included in stats and is not enforced by the connector - ImmutableSet.of(), // projectedColumns don't affect stats - originalHandle.getNameMappingJson(), - originalHandle.getTableLocation(), - originalHandle.getStorageProperties(), - originalHandle.isRecordScannedFiles(), - originalHandle.getMaxScannedFileSize()), - handle -> { - Table icebergTable = catalog.loadTable(session, handle.getSchemaTableName()); - return TableStatisticsReader.getTableStatistics(typeManager, session, handle, icebergTable); - }); + IcebergTableHandle cacheKey = new IcebergTableHandle( + originalHandle.getCatalog(), + originalHandle.getSchemaName(), + originalHandle.getTableName(), + originalHandle.getTableType(), + originalHandle.getSnapshotId(), + originalHandle.getTableSchemaJson(), + originalHandle.getPartitionSpecJson(), + originalHandle.getFormatVersion(), + originalHandle.getUnenforcedPredicate(), + originalHandle.getEnforcedPredicate(), + OptionalLong.empty(), // limit is currently not included in stats and is not enforced by the connector + ImmutableSet.of(), // projectedColumns are used to request statistics only for the required columns, but are not part of cache key + originalHandle.getNameMappingJson(), + originalHandle.getTableLocation(), + originalHandle.getStorageProperties(), + Optional.empty(), // requiredTablePartitioning does not affect stats + false, // recordScannedFiles does not affect stats + originalHandle.getMaxScannedFileSize(), + ImmutableSet.of(), // constraintColumns do not affect stats + Optional.empty()); // forAnalyze does not affect stats + return getIncrementally( + tableStatisticsCache, + cacheKey, + currentStatistics -> currentStatistics.getColumnStatistics().keySet().containsAll(originalHandle.getProjectedColumns()), + projectedColumns -> { + Table icebergTable = catalog.loadTable(session, originalHandle.getSchemaTableName()); + return TableStatisticsReader.getTableStatistics( + typeManager, + session, + originalHandle, + projectedColumns, + icebergTable, + fileSystemFactory.create(session.getIdentity(), icebergTable.io().properties())); + }, + originalHandle.getProjectedColumns()); } @Override @@ -2612,7 +3400,7 @@ Table getIcebergTable(ConnectorSession session, SchemaTableName schemaTableName) @Override public void createMaterializedView(ConnectorSession session, SchemaTableName viewName, ConnectorMaterializedViewDefinition definition, boolean replace, boolean ignoreExisting) { - catalog.createMaterializedView(session, viewName, definition, replace, ignoreExisting); + catalog.createMaterializedView(session, viewName, definition, definition.getProperties(), replace, ignoreExisting); } @Override @@ -2630,6 +3418,7 @@ public boolean delegateMaterializedViewRefreshToConnector(ConnectorSession sessi @Override public ConnectorInsertTableHandle beginRefreshMaterializedView(ConnectorSession session, ConnectorTableHandle tableHandle, List sourceTableHandles, RetryMode retryMode) { + checkState(fromSnapshotForRefresh.isEmpty(), "From Snapshot must be empty at the start of MV refresh operation."); IcebergTableHandle table = (IcebergTableHandle) tableHandle; Table icebergTable = catalog.loadTable(session, table.getSchemaTableName()); beginTransaction(icebergTable); @@ -2649,10 +3438,17 @@ public Optional finishRefreshMaterializedView( IcebergWritableTableHandle table = (IcebergWritableTableHandle) insertHandle; Table icebergTable = transaction.table(); - // delete before insert .. simulating overwrite - transaction.newDelete() - .deleteFromRowFilter(Expressions.alwaysTrue()) - .commit(); + boolean isFullRefresh = fromSnapshotForRefresh.isEmpty(); + if (isFullRefresh) { + // delete before insert .. simulating overwrite + log.info("Performing full MV refresh for storage table: %s", table.getName()); + transaction.newDelete() + .deleteFromRowFilter(Expressions.alwaysTrue()) + .commit(); + } + else { + log.info("Performing incremental MV refresh for storage table: %s", table.getName()); + } List commitTasks = fragments.stream() .map(slice -> commitTaskCodec.fromJson(slice.getBytes())) @@ -2705,10 +3501,18 @@ public Optional finishRefreshMaterializedView( // Update the 'dependsOnTables' property that tracks tables on which the materialized view depends and the corresponding snapshot ids of the tables appendFiles.set(DEPENDS_ON_TABLES, dependencies); appendFiles.set(TRINO_QUERY_START_TIME, session.getStart().toString()); - commit(appendFiles, session); - - transaction.commitTransaction(); + commitUpdateAndTransaction(appendFiles, session, transaction, "refresh materialized view"); transaction = null; + fromSnapshotForRefresh = Optional.empty(); + + // cleanup old snapshots + try { + executeExpireSnapshots(icebergTable, session, System.currentTimeMillis()); + } + catch (Exception e) { + log.error(e, "Failed to delete old snapshot files during materialized view refresh"); + } + return Optional.of(new HiveWrittenPartitions(commitTasks.stream() .map(CommitTaskData::getPath) .collect(toImmutableList()))); @@ -2717,7 +3521,10 @@ public Optional finishRefreshMaterializedView( @Override public List listMaterializedViews(ConnectorSession session, Optional schemaName) { - return catalog.listMaterializedViews(session, schemaName); + return catalog.listTables(session, schemaName).stream() + .filter(info -> info.extendedRelationType() == TableInfo.ExtendedRelationType.TRINO_MATERIALIZED_VIEW) + .map(TableInfo::tableName) + .toList(); } @Override @@ -2766,9 +3573,12 @@ public MaterializedViewFreshness getMaterializedViewFreshness(ConnectorSession s .orElseThrow(() -> new IllegalStateException("Storage table missing in definition of materialized view " + materializedViewName)); Table icebergTable = catalog.loadTable(session, storageTableName); - String dependsOnTables = icebergTable.currentSnapshot().summary().getOrDefault(DEPENDS_ON_TABLES, ""); + String dependsOnTables = Optional.ofNullable(icebergTable.currentSnapshot()) + .map(snapshot -> snapshot.summary().getOrDefault(DEPENDS_ON_TABLES, "")) + .orElse(""); if (dependsOnTables.isEmpty()) { - // Information missing. While it's "unknown" whether storage is stale, we return "stale": under no normal circumstances dependsOnTables should be missing. + // Information missing. While it's "unknown" whether storage is stale, we return "stale". + // Normally dependsOnTables may be missing only when there was no refresh yet. return new MaterializedViewFreshness(STALE, Optional.empty()); } Instant refreshTime = Optional.ofNullable(icebergTable.currentSnapshot().summary().get(TRINO_QUERY_START_TIME)) @@ -2884,6 +3694,16 @@ public Optional redirectTable(ConnectorSession session, return catalog.redirectTable(session, tableName, targetCatalogName.get()); } + public Optional getIncrementalRefreshFromSnapshot() + { + return fromSnapshotForRefresh; + } + + public void disableIncrementalRefresh() + { + fromSnapshotForRefresh = Optional.empty(); + } + private static CollectedStatistics processComputedTableStatistics(Table table, Collection computedStatistics) { Map columnNameToId = table.schema().columns().stream() @@ -2918,6 +3738,37 @@ private void beginTransaction(Table icebergTable) transaction = icebergTable.newTransaction(); } + private void executeExpireSnapshots(Table icebergTable, ConnectorSession session, long expireTimestampMillis) + { + TrinoFileSystem fileSystem = fileSystemFactory.create(session.getIdentity(), icebergTable.io().properties()); + List pathsToDelete = new ArrayList<>(); + // deleteFunction is not accessed from multiple threads unless .executeDeleteWith() is used + Consumer deleteFunction = path -> { + pathsToDelete.add(Location.of(path)); + if (pathsToDelete.size() == DELETE_BATCH_SIZE) { + try { + fileSystem.deleteFiles(pathsToDelete); + pathsToDelete.clear(); + } + catch (IOException e) { + throw new TrinoException(ICEBERG_FILESYSTEM_ERROR, "Failed to delete files during snapshot expiration", e); + } + } + }; + + try { + icebergTable.expireSnapshots() + .expireOlderThan(expireTimestampMillis) + .deleteWith(deleteFunction) + .commit(); + + fileSystem.deleteFiles(pathsToDelete); + } + catch (IOException e) { + throw new TrinoException(ICEBERG_FILESYSTEM_ERROR, "Failed to delete files during snapshot expiration", e); + } + } + private static IcebergTableHandle checkValidTableHandle(ConnectorTableHandle tableHandle) { requireNonNull(tableHandle, "tableHandle is null"); @@ -2944,4 +3795,48 @@ private record FirstChangeSnapshot(Snapshot snapshot) private static final class UnknownTableChange implements TableChangeInfo {} + + private static TableStatistics getIncrementally( + Map> cache, + IcebergTableHandle key, + Predicate isSufficient, + Function, TableStatistics> columnStatisticsLoader, + Set projectedColumns) + { + AtomicReference valueHolder = cache.computeIfAbsent(key, ignore -> new AtomicReference<>()); + TableStatistics oldValue = valueHolder.get(); + if (oldValue != null && isSufficient.test(oldValue)) { + return oldValue; + } + + TableStatistics newValue; + if (oldValue == null) { + newValue = columnStatisticsLoader.apply(projectedColumns); + } + else { + Sets.SetView missingColumns = difference(projectedColumns, oldValue.getColumnStatistics().keySet()); + newValue = columnStatisticsLoader.apply(missingColumns); + } + + verifyNotNull(newValue, "loader returned null for %s", key); + + TableStatistics merged = mergeColumnStatistics(oldValue, newValue); + if (!valueHolder.compareAndSet(oldValue, merged)) { + // if the value changed in the valueHolder, we only add newly loaded value to be sure we have up-to-date value + valueHolder.accumulateAndGet(newValue, IcebergMetadata::mergeColumnStatistics); + } + return merged; + } + + private static TableStatistics mergeColumnStatistics(TableStatistics currentStats, TableStatistics newStats) + { + requireNonNull(newStats, "newStats is null"); + TableStatistics.Builder statisticsBuilder = TableStatistics.builder(); + if (currentStats != null) { + currentStats.getColumnStatistics().forEach(statisticsBuilder::setColumnStatistics); + } + statisticsBuilder.setRowCount(newStats.getRowCount()); + newStats.getColumnStatistics().forEach(statisticsBuilder::setColumnStatistics); + return statisticsBuilder.build(); + } } diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergMetadataColumn.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergMetadataColumn.java index 57543c76ed42..049414faf5cf 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergMetadataColumn.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergMetadataColumn.java @@ -27,6 +27,7 @@ public enum IcebergMetadataColumn { + PARTITION(MetadataColumns.PARTITION_COLUMN_ID, "$partition", VARCHAR, PRIMITIVE), // Avoid row type considering partition evolutions FILE_PATH(MetadataColumns.FILE_PATH.fieldId(), "$path", VARCHAR, PRIMITIVE), FILE_MODIFIED_TIME(Integer.MAX_VALUE - 1001, "$file_modified_time", TIMESTAMP_TZ_MILLIS, PRIMITIVE), // https://github.com/apache/iceberg/issues/5240 /**/; diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergMetadataFactory.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergMetadataFactory.java index 36163eadc864..17ddc7d1ae0a 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergMetadataFactory.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergMetadataFactory.java @@ -13,14 +13,24 @@ */ package io.trino.plugin.iceberg; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableSet; import com.google.inject.Inject; +import io.airlift.concurrent.BoundedExecutor; import io.airlift.json.JsonCodec; -import io.trino.filesystem.TrinoFileSystemFactory; +import io.trino.plugin.hive.metastore.HiveMetastoreFactory; +import io.trino.plugin.hive.metastore.RawHiveMetastoreFactory; import io.trino.plugin.iceberg.catalog.TrinoCatalogFactory; import io.trino.spi.connector.CatalogHandle; import io.trino.spi.security.ConnectorIdentity; import io.trino.spi.type.TypeManager; +import java.util.Optional; +import java.util.concurrent.Executor; +import java.util.concurrent.ExecutorService; +import java.util.function.Predicate; + +import static com.google.common.util.concurrent.MoreExecutors.directExecutor; import static java.util.Objects.requireNonNull; public class IcebergMetadataFactory @@ -29,8 +39,13 @@ public class IcebergMetadataFactory private final CatalogHandle trinoCatalogHandle; private final JsonCodec commitTaskCodec; private final TrinoCatalogFactory catalogFactory; - private final TrinoFileSystemFactory fileSystemFactory; + private final IcebergFileSystemFactory fileSystemFactory; private final TableStatisticsWriter tableStatisticsWriter; + private final Optional metastoreFactory; + private final boolean addFilesProcedureEnabled; + private final Predicate allowedExtraProperties; + private final ExecutorService icebergScanExecutor; + private final Executor metadataFetchingExecutor; @Inject public IcebergMetadataFactory( @@ -38,8 +53,12 @@ public IcebergMetadataFactory( CatalogHandle trinoCatalogHandle, JsonCodec commitTaskCodec, TrinoCatalogFactory catalogFactory, - TrinoFileSystemFactory fileSystemFactory, - TableStatisticsWriter tableStatisticsWriter) + IcebergFileSystemFactory fileSystemFactory, + TableStatisticsWriter tableStatisticsWriter, + @RawHiveMetastoreFactory Optional metastoreFactory, + @ForIcebergScanPlanning ExecutorService icebergScanExecutor, + @ForIcebergMetadata ExecutorService metadataExecutorService, + IcebergConfig config) { this.typeManager = requireNonNull(typeManager, "typeManager is null"); this.trinoCatalogHandle = requireNonNull(trinoCatalogHandle, "trinoCatalogHandle is null"); @@ -47,6 +66,22 @@ public IcebergMetadataFactory( this.catalogFactory = requireNonNull(catalogFactory, "catalogFactory is null"); this.fileSystemFactory = requireNonNull(fileSystemFactory, "fileSystemFactory is null"); this.tableStatisticsWriter = requireNonNull(tableStatisticsWriter, "tableStatisticsWriter is null"); + this.metastoreFactory = requireNonNull(metastoreFactory, "metastoreFactory is null"); + this.icebergScanExecutor = requireNonNull(icebergScanExecutor, "icebergScanExecutor is null"); + this.addFilesProcedureEnabled = config.isAddFilesProcedureEnabled(); + if (config.getAllowedExtraProperties().equals(ImmutableList.of("*"))) { + this.allowedExtraProperties = ignore -> true; + } + else { + this.allowedExtraProperties = ImmutableSet.copyOf(requireNonNull(config.getAllowedExtraProperties(), "allowedExtraProperties is null"))::contains; + } + + if (config.getMetadataParallelism() == 1) { + this.metadataFetchingExecutor = directExecutor(); + } + else { + this.metadataFetchingExecutor = new BoundedExecutor(metadataExecutorService, config.getMetadataParallelism()); + } } public IcebergMetadata create(ConnectorIdentity identity) @@ -57,6 +92,11 @@ public IcebergMetadata create(ConnectorIdentity identity) commitTaskCodec, catalogFactory.create(identity), fileSystemFactory, - tableStatisticsWriter); + tableStatisticsWriter, + metastoreFactory, + addFilesProcedureEnabled, + allowedExtraProperties, + icebergScanExecutor, + metadataFetchingExecutor); } } diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergModule.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergModule.java index c0b5bd47923f..337311a35088 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergModule.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergModule.java @@ -13,36 +13,63 @@ */ package io.trino.plugin.iceberg; +import com.google.common.util.concurrent.ListeningExecutorService; import com.google.inject.Binder; import com.google.inject.Key; import com.google.inject.Module; +import com.google.inject.Provides; import com.google.inject.Scopes; +import com.google.inject.Singleton; import com.google.inject.multibindings.Multibinder; +import io.trino.plugin.base.CatalogName; +import io.trino.plugin.base.classloader.ClassLoaderSafeConnectorPageSinkProvider; +import io.trino.plugin.base.classloader.ClassLoaderSafeConnectorPageSourceProvider; +import io.trino.plugin.base.classloader.ClassLoaderSafeConnectorSplitManager; +import io.trino.plugin.base.classloader.ClassLoaderSafeNodePartitioningProvider; +import io.trino.plugin.base.classloader.ForClassLoaderSafe; import io.trino.plugin.base.session.SessionPropertiesProvider; import io.trino.plugin.hive.FileFormatDataSourceStats; import io.trino.plugin.hive.SortingFileWriterConfig; +import io.trino.plugin.hive.metastore.HiveMetastoreFactory; +import io.trino.plugin.hive.metastore.RawHiveMetastoreFactory; import io.trino.plugin.hive.metastore.thrift.TranslateHiveViews; import io.trino.plugin.hive.orc.OrcReaderConfig; import io.trino.plugin.hive.orc.OrcWriterConfig; import io.trino.plugin.hive.parquet.ParquetReaderConfig; import io.trino.plugin.hive.parquet.ParquetWriterConfig; +import io.trino.plugin.iceberg.catalog.rest.DefaultIcebergFileSystemFactory; +import io.trino.plugin.iceberg.functions.IcebergFunctionProvider; +import io.trino.plugin.iceberg.functions.tablechanges.TableChangesFunctionProcessorProvider; +import io.trino.plugin.iceberg.functions.tablechanges.TableChangesFunctionProvider; import io.trino.plugin.iceberg.procedure.DropExtendedStatsTableProcedure; import io.trino.plugin.iceberg.procedure.ExpireSnapshotsTableProcedure; +import io.trino.plugin.iceberg.procedure.OptimizeManifestsTableProcedure; import io.trino.plugin.iceberg.procedure.OptimizeTableProcedure; import io.trino.plugin.iceberg.procedure.RegisterTableProcedure; import io.trino.plugin.iceberg.procedure.RemoveOrphanFilesTableProcedure; +import io.trino.plugin.iceberg.procedure.RollbackToSnapshotProcedure; import io.trino.plugin.iceberg.procedure.UnregisterTableProcedure; import io.trino.spi.connector.ConnectorNodePartitioningProvider; import io.trino.spi.connector.ConnectorPageSinkProvider; import io.trino.spi.connector.ConnectorPageSourceProvider; import io.trino.spi.connector.ConnectorSplitManager; import io.trino.spi.connector.TableProcedureMetadata; +import io.trino.spi.function.FunctionProvider; +import io.trino.spi.function.table.ConnectorTableFunction; import io.trino.spi.procedure.Procedure; +import java.util.concurrent.ExecutorService; + +import static com.google.common.util.concurrent.MoreExecutors.listeningDecorator; +import static com.google.common.util.concurrent.MoreExecutors.newDirectExecutorService; import static com.google.inject.multibindings.Multibinder.newSetBinder; import static com.google.inject.multibindings.OptionalBinder.newOptionalBinder; +import static io.airlift.concurrent.Threads.daemonThreadsNamed; import static io.airlift.configuration.ConfigBinder.configBinder; import static io.airlift.json.JsonCodecBinder.jsonCodecBinder; +import static io.trino.plugin.base.ClosingBinder.closingBinder; +import static java.util.concurrent.Executors.newCachedThreadPool; +import static java.util.concurrent.Executors.newFixedThreadPool; import static org.weakref.jmx.guice.ExportBinder.newExporter; public class IcebergModule @@ -57,16 +84,22 @@ public void configure(Binder binder) configBinder(binder).bindConfig(SortingFileWriterConfig.class, "iceberg"); newSetBinder(binder, SessionPropertiesProvider.class).addBinding().to(IcebergSessionProperties.class).in(Scopes.SINGLETON); + binder.bind(IcebergSchemaProperties.class).in(Scopes.SINGLETON); binder.bind(IcebergTableProperties.class).in(Scopes.SINGLETON); - binder.bind(IcebergMaterializedViewAdditionalProperties.class).in(Scopes.SINGLETON); + binder.bind(IcebergMaterializedViewProperties.class).in(Scopes.SINGLETON); binder.bind(IcebergAnalyzeProperties.class).in(Scopes.SINGLETON); - newOptionalBinder(binder, Key.get(boolean.class, AsyncIcebergSplitProducer.class)) + newOptionalBinder(binder, Key.get(boolean.class, ForIcebergSplitManager.class)) .setDefault().toInstance(true); - binder.bind(ConnectorSplitManager.class).to(IcebergSplitManager.class).in(Scopes.SINGLETON); - newOptionalBinder(binder, ConnectorPageSourceProvider.class).setDefault().to(IcebergPageSourceProvider.class).in(Scopes.SINGLETON); - binder.bind(ConnectorPageSinkProvider.class).to(IcebergPageSinkProvider.class).in(Scopes.SINGLETON); - binder.bind(ConnectorNodePartitioningProvider.class).to(IcebergNodePartitioningProvider.class).in(Scopes.SINGLETON); + binder.bind(ConnectorSplitManager.class).annotatedWith(ForClassLoaderSafe.class).to(IcebergSplitManager.class).in(Scopes.SINGLETON); + binder.bind(ConnectorSplitManager.class).to(ClassLoaderSafeConnectorSplitManager.class).in(Scopes.SINGLETON); + binder.bind(ConnectorPageSourceProvider.class).annotatedWith(ForClassLoaderSafe.class).to(IcebergPageSourceProviderFactory.class).in(Scopes.SINGLETON); + binder.bind(IcebergPageSourceProviderFactory.class).in(Scopes.SINGLETON); + binder.bind(ConnectorPageSourceProvider.class).to(ClassLoaderSafeConnectorPageSourceProvider.class).in(Scopes.SINGLETON); + binder.bind(ConnectorPageSinkProvider.class).annotatedWith(ForClassLoaderSafe.class).to(IcebergPageSinkProvider.class).in(Scopes.SINGLETON); + binder.bind(ConnectorPageSinkProvider.class).to(ClassLoaderSafeConnectorPageSinkProvider.class).in(Scopes.SINGLETON); + binder.bind(ConnectorNodePartitioningProvider.class).annotatedWith(ForClassLoaderSafe.class).to(IcebergNodePartitioningProvider.class).in(Scopes.SINGLETON); + binder.bind(ConnectorNodePartitioningProvider.class).to(ClassLoaderSafeNodePartitioningProvider.class).in(Scopes.SINGLETON); configBinder(binder).bindConfig(OrcReaderConfig.class); configBinder(binder).bindConfig(OrcWriterConfig.class); @@ -76,6 +109,7 @@ public void configure(Binder binder) binder.bind(TableStatisticsWriter.class).in(Scopes.SINGLETON); binder.bind(IcebergMetadataFactory.class).in(Scopes.SINGLETON); + newOptionalBinder(binder, Key.get(HiveMetastoreFactory.class, RawHiveMetastoreFactory.class)); jsonCodecBinder(binder).bindJsonCodec(CommitTaskData.class); @@ -85,6 +119,8 @@ public void configure(Binder binder) binder.bind(IcebergFileWriterFactory.class).in(Scopes.SINGLETON); newExporter(binder).export(IcebergFileWriterFactory.class).withGeneratedName(); + binder.bind(IcebergEnvironmentContext.class).asEagerSingleton(); + Multibinder procedures = newSetBinder(binder, Procedure.class); procedures.addBinding().toProvider(RollbackToSnapshotProcedure.class).in(Scopes.SINGLETON); procedures.addBinding().toProvider(RegisterTableProcedure.class).in(Scopes.SINGLETON); @@ -92,8 +128,50 @@ public void configure(Binder binder) Multibinder tableProcedures = newSetBinder(binder, TableProcedureMetadata.class); tableProcedures.addBinding().toProvider(OptimizeTableProcedure.class).in(Scopes.SINGLETON); + tableProcedures.addBinding().toProvider(OptimizeManifestsTableProcedure.class).in(Scopes.SINGLETON); tableProcedures.addBinding().toProvider(DropExtendedStatsTableProcedure.class).in(Scopes.SINGLETON); tableProcedures.addBinding().toProvider(ExpireSnapshotsTableProcedure.class).in(Scopes.SINGLETON); tableProcedures.addBinding().toProvider(RemoveOrphanFilesTableProcedure.class).in(Scopes.SINGLETON); + + newSetBinder(binder, ConnectorTableFunction.class).addBinding().toProvider(TableChangesFunctionProvider.class).in(Scopes.SINGLETON); + binder.bind(FunctionProvider.class).to(IcebergFunctionProvider.class).in(Scopes.SINGLETON); + binder.bind(TableChangesFunctionProcessorProvider.class).in(Scopes.SINGLETON); + + newOptionalBinder(binder, IcebergFileSystemFactory.class).setDefault().to(DefaultIcebergFileSystemFactory.class).in(Scopes.SINGLETON); + + closingBinder(binder).registerExecutor(Key.get(ExecutorService.class, ForIcebergMetadata.class)); + closingBinder(binder).registerExecutor(Key.get(ListeningExecutorService.class, ForIcebergSplitManager.class)); + closingBinder(binder).registerExecutor(Key.get(ExecutorService.class, ForIcebergScanPlanning.class)); + + binder.bind(IcebergConnector.class).in(Scopes.SINGLETON); + } + + @Singleton + @Provides + @ForIcebergMetadata + public ExecutorService createIcebergMetadataExecutor(CatalogName catalogName) + { + return newCachedThreadPool(daemonThreadsNamed("iceberg-metadata-" + catalogName + "-%s")); + } + + @Provides + @Singleton + @ForIcebergSplitManager + public ListeningExecutorService createSplitSourceExecutor(CatalogName catalogName) + { + return listeningDecorator(newCachedThreadPool(daemonThreadsNamed("iceberg-split-source-" + catalogName + "-%s"))); + } + + @Provides + @Singleton + @ForIcebergScanPlanning + public ExecutorService createSplitManagerExecutor(CatalogName catalogName, IcebergConfig config) + { + if (config.getSplitManagerThreads() == 0) { + return newDirectExecutorService(); + } + return newFixedThreadPool( + config.getSplitManagerThreads(), + daemonThreadsNamed("iceberg-split-manager-" + catalogName + "-%s")); } } diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergNodePartitioningProvider.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergNodePartitioningProvider.java index 2ad39c2f426c..9f2df36c16bd 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergNodePartitioningProvider.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergNodePartitioningProvider.java @@ -15,6 +15,7 @@ import com.google.inject.Inject; import io.trino.spi.connector.BucketFunction; +import io.trino.spi.connector.ConnectorBucketNodeMap; import io.trino.spi.connector.ConnectorNodePartitioningProvider; import io.trino.spi.connector.ConnectorPartitioningHandle; import io.trino.spi.connector.ConnectorSession; @@ -22,12 +23,12 @@ import io.trino.spi.type.Type; import io.trino.spi.type.TypeManager; import io.trino.spi.type.TypeOperators; -import org.apache.iceberg.Schema; import java.util.List; +import java.util.Optional; -import static io.trino.plugin.iceberg.IcebergUtil.schemaFromHandles; -import static io.trino.plugin.iceberg.PartitionFields.parsePartitionFields; +import static io.trino.plugin.iceberg.IcebergPartitionFunction.Transform.BUCKET; +import static io.trino.spi.connector.ConnectorBucketNodeMap.createBucketNodeMap; public class IcebergNodePartitioningProvider implements ConnectorNodePartitioningProvider @@ -40,6 +41,23 @@ public IcebergNodePartitioningProvider(TypeManager typeManager) this.typeOperators = typeManager.getTypeOperators(); } + @Override + public Optional getBucketNodeMapping( + ConnectorTransactionHandle transactionHandle, + ConnectorSession session, + ConnectorPartitioningHandle partitioningHandle) + { + IcebergPartitioningHandle handle = (IcebergPartitioningHandle) partitioningHandle; + + List partitionFunctions = handle.partitionFunctions(); + // when there is a single bucket partition function, inform the engine there is a limit on the number of buckets + // TODO: when there are multiple bucket partition functions, we could compute the product of bucket counts, but this causes the engine to create too many writers + if (partitionFunctions.size() == 1 && partitionFunctions.get(0).transform() == BUCKET) { + return Optional.of(createBucketNodeMap(partitionFunctions.get(0).size().orElseThrow()).withCacheKeyHint(handle.getCacheKeyHint())); + } + return Optional.empty(); + } + @Override public BucketFunction getBucketFunction( ConnectorTransactionHandle transactionHandle, @@ -48,16 +66,11 @@ public BucketFunction getBucketFunction( List partitionChannelTypes, int bucketCount) { - if (partitioningHandle instanceof IcebergUpdateHandle) { + IcebergPartitioningHandle handle = (IcebergPartitioningHandle) partitioningHandle; + if (handle.update()) { return new IcebergUpdateBucketFunction(bucketCount); } - IcebergPartitioningHandle handle = (IcebergPartitioningHandle) partitioningHandle; - Schema schema = schemaFromHandles(handle.getPartitioningColumns()); - return new IcebergBucketFunction( - typeOperators, - parsePartitionFields(schema, handle.getPartitioning()), - handle.getPartitioningColumns(), - bucketCount); + return new IcebergBucketFunction(handle, typeOperators, bucketCount); } } diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergOrcFileWriter.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergOrcFileWriter.java index 6eb0c19f289b..f0e04051e7b4 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergOrcFileWriter.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergOrcFileWriter.java @@ -96,9 +96,11 @@ public IcebergOrcFileWriter( } @Override - public Metrics getMetrics() + public FileMetrics getFileMetrics() { - return computeMetrics(metricsConfig, icebergSchema, orcColumns, orcWriter.getFileRowCount(), orcWriter.getFileStats()); + return new FileMetrics( + computeMetrics(metricsConfig, icebergSchema, orcColumns, orcWriter.getFileRowCount(), orcWriter.getFileStats()), + Optional.empty()); } private static Metrics computeMetrics(MetricsConfig metricsConfig, Schema icebergSchema, ColumnMetadata orcColumns, long fileRowCount, Optional> columnStatistics) diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergPageSink.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergPageSink.java index b54d8f25fa9b..dd439bc639d4 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergPageSink.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergPageSink.java @@ -42,6 +42,7 @@ import org.apache.iceberg.Schema; import org.apache.iceberg.io.LocationProvider; import org.apache.iceberg.transforms.Transform; +import org.apache.iceberg.types.TypeUtil; import org.apache.iceberg.types.Types; import java.io.Closeable; @@ -55,17 +56,18 @@ import java.util.concurrent.CompletableFuture; import java.util.function.Function; -import static com.google.common.base.Preconditions.checkArgument; +import static com.google.common.base.Preconditions.checkState; import static com.google.common.base.Verify.verify; import static com.google.common.collect.ImmutableList.toImmutableList; import static io.airlift.slice.Slices.wrappedBuffer; import static io.trino.plugin.iceberg.IcebergErrorCode.ICEBERG_INVALID_METADATA; import static io.trino.plugin.iceberg.IcebergErrorCode.ICEBERG_TOO_MANY_OPEN_PARTITIONS; import static io.trino.plugin.iceberg.IcebergSessionProperties.isSortedWritingEnabled; -import static io.trino.plugin.iceberg.IcebergUtil.getColumns; +import static io.trino.plugin.iceberg.IcebergUtil.getTopLevelColumns; import static io.trino.plugin.iceberg.PartitionTransforms.getColumnTransform; import static io.trino.plugin.iceberg.util.Timestamps.getTimestampTz; import static io.trino.plugin.iceberg.util.Timestamps.timestampTzToMicros; +import static io.trino.spi.block.RowBlock.getRowFieldsFromBlock; import static io.trino.spi.type.BigintType.BIGINT; import static io.trino.spi.type.BooleanType.BOOLEAN; import static io.trino.spi.type.DateType.DATE; @@ -105,6 +107,7 @@ public class IcebergPageSink private final MetricsConfig metricsConfig; private final PagePartitioner pagePartitioner; private final long targetMaxFileSize; + private final long idleWriterMinFileSize; private final Map storageProperties; private final List sortOrder; private final boolean sortedWritingEnabled; @@ -120,10 +123,12 @@ public class IcebergPageSink private final List writers = new ArrayList<>(); private final List closedWriterRollbackActions = new ArrayList<>(); private final Collection commitTasks = new ArrayList<>(); + private final List activeWriters = new ArrayList<>(); private long writtenBytes; private long memoryUsage; private long validationCpuNanos; + private long currentOpenWriters; public IcebergPageSink( Schema outputSchema, @@ -155,8 +160,9 @@ public IcebergPageSink( this.fileFormat = requireNonNull(fileFormat, "fileFormat is null"); this.metricsConfig = MetricsConfig.fromProperties(requireNonNull(storageProperties, "storageProperties is null")); this.maxOpenWriters = maxOpenWriters; - this.pagePartitioner = new PagePartitioner(pageIndexerFactory, toPartitionColumns(inputColumns, partitionSpec)); + this.pagePartitioner = new PagePartitioner(pageIndexerFactory, toPartitionColumns(inputColumns, partitionSpec, outputSchema)); this.targetMaxFileSize = IcebergSessionProperties.getTargetMaxFileSize(session); + this.idleWriterMinFileSize = IcebergSessionProperties.getIdleWriterMinFileSize(session); this.storageProperties = requireNonNull(storageProperties, "storageProperties is null"); this.sortOrder = requireNonNull(sortOrder, "sortOrder is null"); this.sortedWritingEnabled = isSortedWritingEnabled(session); @@ -165,7 +171,7 @@ public IcebergPageSink( this.tempDirectory = Location.of(locationProvider.newDataLocation("trino-tmp-files")); this.typeManager = requireNonNull(typeManager, "typeManager is null"); this.pageSorter = requireNonNull(pageSorter, "pageSorter is null"); - this.columnTypes = getColumns(outputSchema, typeManager).stream() + this.columnTypes = getTopLevelColumns(outputSchema, typeManager).stream() .map(IcebergColumnHandle::getType) .collect(toImmutableList()); @@ -211,6 +217,7 @@ public long getValidationCpuNanos() public CompletableFuture appendPage(Page page) { doAppend(page); + closeIdleWriters(); return NOT_BLOCKED; } @@ -300,7 +307,9 @@ private void writePage(Page page) pageForWriter = pageForWriter.getPositions(positions, 0, positions.length); } - IcebergFileWriter writer = writers.get(index).getWriter(); + WriteContext writeContext = writers.get(index); + verify(writeContext != null, "Expected writer at index %s", index); + IcebergFileWriter writer = writeContext.getWriter(); long currentWritten = writer.getWrittenBytes(); long currentMemory = writer.getMemoryUsage(); @@ -309,6 +318,8 @@ private void writePage(Page page) writtenBytes += (writer.getWrittenBytes() - currentWritten); memoryUsage += (writer.getMemoryUsage() - currentMemory); + // Mark this writer as active (i.e. not idle) + activeWriters.set(index, true); } } @@ -316,13 +327,10 @@ private int[] getWriterIndexes(Page page) { int[] writerIndexes = pagePartitioner.partitionPage(page); - if (pagePartitioner.getMaxIndex() >= maxOpenWriters) { - throw new TrinoException(ICEBERG_TOO_MANY_OPEN_PARTITIONS, format("Exceeded limit of %s open writers for partitions", maxOpenWriters)); - } - // expand writers list to new size while (writers.size() <= pagePartitioner.getMaxIndex()) { writers.add(null); + activeWriters.add(false); } // create missing writers @@ -366,17 +374,36 @@ private int[] getWriterIndexes(Page page) } writers.set(writerIndex, writer); + currentOpenWriters++; memoryUsage += writer.getWriter().getMemoryUsage(); } verify(writers.size() == pagePartitioner.getMaxIndex() + 1); - verify(!writers.contains(null)); + + if (currentOpenWriters > maxOpenWriters) { + throw new TrinoException(ICEBERG_TOO_MANY_OPEN_PARTITIONS, format("Exceeded limit of %s open writers for partitions: %s", maxOpenWriters, currentOpenWriters)); + } return writerIndexes; } + public void closeIdleWriters() + { + for (int writerIndex = 0; writerIndex < writers.size(); writerIndex++) { + WriteContext writeContext = writers.get(writerIndex); + if (activeWriters.get(writerIndex) || writeContext == null || writeContext.getWriter().getWrittenBytes() <= idleWriterMinFileSize) { + activeWriters.set(writerIndex, false); + continue; + } + closeWriter(writerIndex); + } + } + private void closeWriter(int writerIndex) { WriteContext writeContext = writers.get(writerIndex); + if (writeContext == null) { + return; + } IcebergFileWriter writer = writeContext.getWriter(); long currentWritten = writer.getWrittenBytes(); @@ -389,16 +416,18 @@ private void closeWriter(int writerIndex) memoryUsage -= currentMemory; writers.set(writerIndex, null); + currentOpenWriters--; CommitTaskData task = new CommitTaskData( writeContext.getPath(), fileFormat, writer.getWrittenBytes(), - new MetricsWrapper(writer.getMetrics()), + new MetricsWrapper(writer.getFileMetrics().metrics()), PartitionSpecParser.toJson(partitionSpec), writeContext.getPartitionData().map(PartitionData::toJson), DATA, - Optional.empty()); + Optional.empty(), + writer.getFileMetrics().splitOffsets()); commitTasks.add(wrappedBuffer(jsonCodec.toJsonBytes(task))); } @@ -426,7 +455,7 @@ private Optional getPartitionData(List columns, Object[] values = new Object[columns.size()]; for (int i = 0; i < columns.size(); i++) { PartitionColumn column = columns.get(i); - Block block = page.getBlock(column.getSourceChannel()); + Block block = PagePartitioner.getPartitionBlock(column, page); Type type = column.getSourceType(); org.apache.iceberg.types.Type icebergType = outputSchema.findType(column.getField().sourceId()); Object value = getIcebergValue(block, position, type); @@ -494,7 +523,7 @@ public static Object getIcebergValue(Block block, int position, Type type) throw new UnsupportedOperationException("Type not supported as partition column: " + type.getDisplayName()); } - private static List toPartitionColumns(List handles, PartitionSpec partitionSpec) + private static List toPartitionColumns(List handles, PartitionSpec partitionSpec, Schema schema) { Map idChannels = new HashMap<>(); for (int i = 0; i < handles.size(); i++) { @@ -502,16 +531,62 @@ private static List toPartitionColumns(List { - Integer channel = idChannels.get(field.sourceId()); - checkArgument(channel != null, "partition field not found: %s", field); - Type inputType = handles.get(channel).getType(); - ColumnTransform transform = getColumnTransform(field, inputType); - return new PartitionColumn(field, channel, inputType, transform.getType(), transform.getBlockTransform()); - }) + .map(field -> getPartitionColumn(field, handles, schema.asStruct(), idChannels)) .collect(toImmutableList()); } + private static PartitionColumn getPartitionColumn(PartitionField field, List handles, Types.StructType schema, Map idChannels) + { + List sourceChannels = getIndexPathToField(schema, getNestedFieldIds(schema, field.sourceId())); + Type sourceType = handles.get(idChannels.get(field.sourceId())).getType(); + ColumnTransform transform = getColumnTransform(field, sourceType); + return new PartitionColumn(field, sourceChannels, sourceType, transform.getType(), transform.getBlockTransform()); + } + + private static List getNestedFieldIds(Types.StructType schema, Integer sourceId) + { + Map parentIndex = TypeUtil.indexParents(schema); + Map idIndex = TypeUtil.indexById(schema); + ImmutableList.Builder parentColumnsBuilder = ImmutableList.builder(); + + parentColumnsBuilder.add(idIndex.get(sourceId).fieldId()); + Integer current = parentIndex.get(sourceId); + + while (current != null) { + parentColumnsBuilder.add(idIndex.get(current).fieldId()); + current = parentIndex.get(current); + } + return parentColumnsBuilder.build().reverse(); + } + + private static List getIndexPathToField(Types.StructType schema, List nestedFieldIds) + { + ImmutableList.Builder sourceIdsBuilder = ImmutableList.builder(); + Types.StructType current = schema; + + // Iterate over field names while finding position in schema + for (int i = 0; i < nestedFieldIds.size(); i++) { + int fieldId = nestedFieldIds.get(i); + sourceIdsBuilder.add(findFieldPosFromSchema(fieldId, current)); + + if (i + 1 < nestedFieldIds.size()) { + checkState(current.field(fieldId).type().isStructType(), "Could not find field " + nestedFieldIds + " in schema"); + current = current.field(fieldId).type().asStructType(); + } + } + return sourceIdsBuilder.build(); + } + + private static int findFieldPosFromSchema(int fieldId, Types.StructType struct) + { + for (int i = 0; i < struct.fields().size(); i++) { + if (struct.fields().get(i).fieldId() == fieldId) { + return i; + } + } + throw new IllegalArgumentException("Could not find field " + fieldId + " in schema"); + } + private static class WriteContext { private final IcebergFileWriter writer; @@ -569,7 +644,7 @@ public int[] partitionPage(Page page) Block[] blocks = new Block[columns.size()]; for (int i = 0; i < columns.size(); i++) { PartitionColumn column = columns.get(i); - Block block = page.getBlock(column.getSourceChannel()); + Block block = getPartitionBlock(column, page); blocks[i] = column.getBlockTransform().apply(block); } Page transformed = new Page(page.getPositionCount(), blocks); @@ -577,6 +652,16 @@ public int[] partitionPage(Page page) return pageIndexer.indexPage(transformed); } + private static Block getPartitionBlock(PartitionColumn column, Page page) + { + List sourceChannels = column.getSourceChannels(); + Block block = page.getBlock(sourceChannels.get(0)); + for (int i = 1; i < sourceChannels.size(); i++) { + block = getRowFieldsFromBlock(block).get(sourceChannels.get(i)); + } + return block; + } + public int getMaxIndex() { return pageIndexer.getMaxIndex(); @@ -591,15 +676,15 @@ public List getColumns() private static class PartitionColumn { private final PartitionField field; - private final int sourceChannel; + private final List sourceChannels; private final Type sourceType; private final Type resultType; private final Function blockTransform; - public PartitionColumn(PartitionField field, int sourceChannel, Type sourceType, Type resultType, Function blockTransform) + public PartitionColumn(PartitionField field, List sourceChannels, Type sourceType, Type resultType, Function blockTransform) { this.field = requireNonNull(field, "field is null"); - this.sourceChannel = sourceChannel; + this.sourceChannels = ImmutableList.copyOf(requireNonNull(sourceChannels, "sourceChannels is null")); this.sourceType = requireNonNull(sourceType, "sourceType is null"); this.resultType = requireNonNull(resultType, "resultType is null"); this.blockTransform = requireNonNull(blockTransform, "blockTransform is null"); @@ -610,9 +695,9 @@ public PartitionField getField() return field; } - public int getSourceChannel() + public List getSourceChannels() { - return sourceChannel; + return sourceChannels; } public Type getSourceType() diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergPageSinkProvider.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergPageSinkProvider.java index 1242c9dd4c19..69ec65a25fa0 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergPageSinkProvider.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergPageSinkProvider.java @@ -16,7 +16,6 @@ import com.google.inject.Inject; import io.airlift.json.JsonCodec; import io.airlift.units.DataSize; -import io.trino.filesystem.TrinoFileSystemFactory; import io.trino.plugin.hive.SortingFileWriterConfig; import io.trino.plugin.iceberg.procedure.IcebergOptimizeHandle; import io.trino.plugin.iceberg.procedure.IcebergTableExecuteHandle; @@ -48,7 +47,7 @@ public class IcebergPageSinkProvider implements ConnectorPageSinkProvider { - private final TrinoFileSystemFactory fileSystemFactory; + private final IcebergFileSystemFactory fileSystemFactory; private final JsonCodec jsonCodec; private final IcebergFileWriterFactory fileWriterFactory; private final PageIndexerFactory pageIndexerFactory; @@ -60,7 +59,7 @@ public class IcebergPageSinkProvider @Inject public IcebergPageSinkProvider( - TrinoFileSystemFactory fileSystemFactory, + IcebergFileSystemFactory fileSystemFactory, JsonCodec jsonCodec, IcebergFileWriterFactory fileWriterFactory, PageIndexerFactory pageIndexerFactory, @@ -104,7 +103,7 @@ private ConnectorPageSink createPageSink(ConnectorSession session, IcebergWritab locationProvider, fileWriterFactory, pageIndexerFactory, - fileSystemFactory.create(session), + fileSystemFactory.create(session.getIdentity(), tableHandle.getFileIoProperties()), tableHandle.getInputColumns(), jsonCodec, session, @@ -135,7 +134,7 @@ public ConnectorPageSink createPageSink(ConnectorTransactionHandle transactionHa locationProvider, fileWriterFactory, pageIndexerFactory, - fileSystemFactory.create(session), + fileSystemFactory.create(session.getIdentity(), executeHandle.getFileIoProperties()), optimizeHandle.getTableColumns(), jsonCodec, session, @@ -147,9 +146,13 @@ public ConnectorPageSink createPageSink(ConnectorTransactionHandle transactionHa sortingFileWriterMaxOpenFiles, typeManager, pageSorter); + case OPTIMIZE_MANIFESTS: case DROP_EXTENDED_STATS: + case ROLLBACK_TO_SNAPSHOT: case EXPIRE_SNAPSHOTS: case REMOVE_ORPHAN_FILES: + case ADD_FILES: + case ADD_FILES_FROM_TABLE: // handled via ConnectorMetadata.executeTableExecute } throw new IllegalArgumentException("Unknown procedure: " + executeHandle.getProcedureId()); @@ -168,7 +171,7 @@ public ConnectorMergeSink createMergeSink(ConnectorTransactionHandle transaction return new IcebergMergeSink( locationProvider, fileWriterFactory, - fileSystemFactory.create(session), + fileSystemFactory.create(session.getIdentity(), tableHandle.getFileIoProperties()), jsonCodec, session, tableHandle.getFileFormat(), @@ -176,6 +179,6 @@ public ConnectorMergeSink createMergeSink(ConnectorTransactionHandle transaction schema, partitionsSpecs, pageSink, - tableHandle.getInputColumns().size()); + schema.columns().size()); } } diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergPageSource.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergPageSource.java index 092428eac85e..ec5bd9e55d85 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergPageSource.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergPageSource.java @@ -13,7 +13,6 @@ */ package io.trino.plugin.iceberg; -import com.google.common.collect.ImmutableMap; import io.trino.plugin.hive.ReaderProjectionsAdapter; import io.trino.plugin.iceberg.delete.RowPredicate; import io.trino.spi.Page; @@ -26,16 +25,16 @@ import java.io.IOException; import java.io.UncheckedIOException; import java.util.List; -import java.util.Map; import java.util.Optional; import java.util.OptionalLong; +import java.util.concurrent.CompletableFuture; +import java.util.function.Function; import java.util.function.Supplier; import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.base.Throwables.throwIfInstanceOf; import static io.trino.plugin.base.util.Closables.closeAllSuppress; import static io.trino.plugin.iceberg.IcebergErrorCode.ICEBERG_BAD_DATA; -import static java.lang.String.format; import static java.util.Objects.requireNonNull; public class IcebergPageSource @@ -45,10 +44,9 @@ public class IcebergPageSource private final ConnectorPageSource delegate; private final Optional projectionsAdapter; private final Supplier> deletePredicate; - // An array with one element per field in the $row_id column. The value in the array points to the - // channel where the data can be read from. - private int[] rowIdChildColumnIndexes = new int[0]; + private final Function rowIdBlockFactory; // The $row_id's index in 'expectedColumns', or -1 if there isn't one + // this column with contain row position populated in the source, and must be wrapped with constant data for full row id private int rowIdColumnIndex = -1; // Maps the Iceberg field ids of unmodified columns to their indexes in updateRowIdChildColumnIndexes @@ -57,7 +55,8 @@ public IcebergPageSource( List requiredColumns, ConnectorPageSource delegate, Optional projectionsAdapter, - Supplier> deletePredicate) + Supplier> deletePredicate, + Function rowIdBlockFactory) { // expectedColumns should contain columns which should be in the final Page // requiredColumns should include all expectedColumns as well as any columns needed by the DeleteFilter @@ -69,24 +68,15 @@ public IcebergPageSource( checkArgument(expectedColumn.equals(requiredColumns.get(i)), "Expected columns must be a prefix of required columns"); expectedColumnIndexes[i] = i; - if (expectedColumn.isUpdateRowIdColumn() || expectedColumn.isMergeRowIdColumn()) { + if (expectedColumn.isMergeRowIdColumn()) { this.rowIdColumnIndex = i; - - Map fieldIdToColumnIndex = mapFieldIdsToIndex(requiredColumns); - List rowIdFields = expectedColumn.getColumnIdentity().getChildren(); - ImmutableMap.Builder fieldIdToRowIdIndex = ImmutableMap.builder(); - this.rowIdChildColumnIndexes = new int[rowIdFields.size()]; - for (int columnIndex = 0; columnIndex < rowIdFields.size(); columnIndex++) { - int fieldId = rowIdFields.get(columnIndex).getId(); - rowIdChildColumnIndexes[columnIndex] = requireNonNull(fieldIdToColumnIndex.get(fieldId), () -> format("Column %s not found in requiredColumns", fieldId)); - fieldIdToRowIdIndex.put(fieldId, columnIndex); - } } } this.delegate = requireNonNull(delegate, "delegate is null"); this.projectionsAdapter = requireNonNull(projectionsAdapter, "projectionsAdapter is null"); this.deletePredicate = requireNonNull(deletePredicate, "deletePredicate is null"); + this.rowIdBlockFactory = requireNonNull(rowIdBlockFactory, "rowIdBlockFactory is null"); } @Override @@ -113,6 +103,12 @@ public boolean isFinished() return delegate.isFinished(); } + @Override + public CompletableFuture isBlocked() + { + return delegate.isBlocked(); + } + @Override public Page getNextPage() { @@ -156,21 +152,11 @@ private Page withRowIdBlock(Page page) return page; } - Block[] rowIdFields = new Block[rowIdChildColumnIndexes.length]; - for (int childIndex = 0; childIndex < rowIdChildColumnIndexes.length; childIndex++) { - rowIdFields[childIndex] = page.getBlock(rowIdChildColumnIndexes[childIndex]); - } - + RowBlock rowIdBlock = rowIdBlockFactory.apply(page.getBlock(rowIdColumnIndex)); Block[] fullPage = new Block[page.getChannelCount()]; for (int channel = 0; channel < page.getChannelCount(); channel++) { - if (channel == rowIdColumnIndex) { - fullPage[channel] = RowBlock.fromFieldBlocks(page.getPositionCount(), Optional.empty(), rowIdFields); - continue; - } - - fullPage[channel] = page.getBlock(channel); + fullPage[channel] = channel == rowIdColumnIndex ? rowIdBlock : page.getBlock(channel); } - return new Page(page.getPositionCount(), fullPage); } @@ -207,13 +193,4 @@ protected void closeWithSuppression(Throwable throwable) { closeAllSuppress(throwable, this); } - - private static Map mapFieldIdsToIndex(List columns) - { - ImmutableMap.Builder fieldIdsToIndex = ImmutableMap.builder(); - for (int i = 0; i < columns.size(); i++) { - fieldIdsToIndex.put(columns.get(i).getId(), i); - } - return fieldIdsToIndex.buildOrThrow(); - } } diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergPageSourceProvider.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergPageSourceProvider.java index 259ea8d5db6d..ebc8a3461b9a 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergPageSourceProvider.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergPageSourceProvider.java @@ -13,20 +13,16 @@ */ package io.trino.plugin.iceberg; -import com.google.common.base.Suppliers; -import com.google.common.base.VerifyException; import com.google.common.collect.AbstractIterator; import com.google.common.collect.BiMap; import com.google.common.collect.ImmutableBiMap; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; -import com.google.common.graph.Traverser; import com.google.inject.Inject; import io.airlift.slice.Slice; import io.trino.filesystem.Location; import io.trino.filesystem.TrinoFileSystem; -import io.trino.filesystem.TrinoFileSystemFactory; import io.trino.filesystem.TrinoInputFile; import io.trino.memory.context.AggregatedMemoryContext; import io.trino.orc.OrcColumn; @@ -38,33 +34,34 @@ import io.trino.orc.OrcRecordReader; import io.trino.orc.TupleDomainOrcPredicate; import io.trino.orc.TupleDomainOrcPredicate.TupleDomainOrcPredicateBuilder; -import io.trino.orc.metadata.OrcType; -import io.trino.parquet.BloomFilterStore; +import io.trino.parquet.Column; import io.trino.parquet.Field; import io.trino.parquet.ParquetCorruptionException; import io.trino.parquet.ParquetDataSource; import io.trino.parquet.ParquetDataSourceId; import io.trino.parquet.ParquetReaderOptions; +import io.trino.parquet.metadata.FileMetadata; +import io.trino.parquet.metadata.ParquetMetadata; import io.trino.parquet.predicate.TupleDomainParquetPredicate; import io.trino.parquet.reader.MetadataReader; import io.trino.parquet.reader.ParquetReader; +import io.trino.parquet.reader.RowGroupInfo; import io.trino.plugin.hive.FileFormatDataSourceStats; -import io.trino.plugin.hive.ReaderColumns; -import io.trino.plugin.hive.ReaderPageSource; -import io.trino.plugin.hive.ReaderProjectionsAdapter; +import io.trino.plugin.hive.TransformConnectorPageSource; import io.trino.plugin.hive.orc.OrcPageSource; -import io.trino.plugin.hive.orc.OrcPageSource.ColumnAdaptation; -import io.trino.plugin.hive.orc.OrcReaderConfig; import io.trino.plugin.hive.parquet.ParquetPageSource; -import io.trino.plugin.hive.parquet.ParquetReaderConfig; -import io.trino.plugin.hive.parquet.TrinoParquetDataSource; import io.trino.plugin.iceberg.IcebergParquetColumnIOConverter.FieldContext; import io.trino.plugin.iceberg.delete.DeleteFile; -import io.trino.plugin.iceberg.delete.DeleteFilter; -import io.trino.plugin.iceberg.delete.PositionDeleteFilter; +import io.trino.plugin.iceberg.delete.DeleteManager; import io.trino.plugin.iceberg.delete.RowPredicate; import io.trino.plugin.iceberg.fileio.ForwardingInputFile; +import io.trino.spi.Page; import io.trino.spi.TrinoException; +import io.trino.spi.block.Block; +import io.trino.spi.block.IntArrayBlock; +import io.trino.spi.block.RowBlock; +import io.trino.spi.block.RunLengthEncodedBlock; +import io.trino.spi.block.VariableWidthBlock; import io.trino.spi.connector.ColumnHandle; import io.trino.spi.connector.ConnectorPageSource; import io.trino.spi.connector.ConnectorPageSourceProvider; @@ -74,20 +71,17 @@ import io.trino.spi.connector.ConnectorTransactionHandle; import io.trino.spi.connector.DynamicFilter; import io.trino.spi.connector.EmptyPageSource; +import io.trino.spi.connector.FixedPageSource; import io.trino.spi.predicate.Domain; import io.trino.spi.predicate.NullableValue; -import io.trino.spi.predicate.Range; import io.trino.spi.predicate.TupleDomain; -import io.trino.spi.predicate.ValueSet; import io.trino.spi.type.ArrayType; import io.trino.spi.type.MapType; import io.trino.spi.type.RowType; -import io.trino.spi.type.StandardTypes; import io.trino.spi.type.Type; import io.trino.spi.type.TypeManager; import org.apache.avro.file.DataFileStream; import org.apache.avro.generic.GenericDatumReader; -import org.apache.iceberg.MetadataColumns; import org.apache.iceberg.PartitionSpec; import org.apache.iceberg.PartitionSpecParser; import org.apache.iceberg.Schema; @@ -99,22 +93,16 @@ import org.apache.iceberg.mapping.NameMapping; import org.apache.iceberg.mapping.NameMappingParser; import org.apache.iceberg.parquet.ParquetSchemaUtil; -import org.apache.iceberg.types.Conversions; +import org.apache.iceberg.types.Types; +import org.apache.iceberg.util.StructLikeWrapper; import org.apache.parquet.column.ColumnDescriptor; -import org.apache.parquet.hadoop.metadata.BlockMetaData; -import org.apache.parquet.hadoop.metadata.FileMetaData; -import org.apache.parquet.hadoop.metadata.ParquetMetadata; -import org.apache.parquet.io.ColumnIO; import org.apache.parquet.io.MessageColumnIO; import org.apache.parquet.schema.GroupType; import org.apache.parquet.schema.MessageType; import org.apache.parquet.schema.PrimitiveType; -import org.roaringbitmap.longlong.LongBitmapDataProvider; -import org.roaringbitmap.longlong.Roaring64Bitmap; import java.io.IOException; import java.io.UncheckedIOException; -import java.nio.ByteBuffer; import java.util.ArrayList; import java.util.HashMap; import java.util.List; @@ -123,31 +111,38 @@ import java.util.Optional; import java.util.OptionalLong; import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; +import java.util.function.Function; import java.util.function.Supplier; +import java.util.stream.IntStream; +import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.base.Preconditions.checkState; +import static com.google.common.base.Suppliers.memoize; +import static com.google.common.base.Throwables.throwIfInstanceOf; import static com.google.common.base.Verify.verify; import static com.google.common.collect.ImmutableList.toImmutableList; import static com.google.common.collect.ImmutableMap.toImmutableMap; import static com.google.common.collect.ImmutableSet.toImmutableSet; import static com.google.common.collect.Maps.uniqueIndex; +import static io.airlift.slice.SizeOf.SIZE_OF_LONG; import static io.airlift.slice.Slices.utf8Slice; import static io.trino.memory.context.AggregatedMemoryContext.newSimpleAggregatedMemoryContext; import static io.trino.orc.OrcReader.INITIAL_BATCH_SIZE; import static io.trino.orc.OrcReader.ProjectedLayout; import static io.trino.orc.OrcReader.fullyProjectedLayout; -import static io.trino.parquet.BloomFilterStore.getBloomFilterStore; import static io.trino.parquet.ParquetTypeUtils.getColumnIO; import static io.trino.parquet.ParquetTypeUtils.getDescriptors; import static io.trino.parquet.predicate.PredicateUtils.buildPredicate; -import static io.trino.parquet.predicate.PredicateUtils.predicateMatches; -import static io.trino.plugin.iceberg.IcebergColumnHandle.TRINO_MERGE_PARTITION_DATA; -import static io.trino.plugin.iceberg.IcebergColumnHandle.TRINO_MERGE_PARTITION_SPEC_ID; +import static io.trino.parquet.predicate.PredicateUtils.getFilteredRowGroups; +import static io.trino.plugin.hive.parquet.ParquetPageSourceFactory.createDataSource; +import static io.trino.plugin.iceberg.ColumnIdentity.TypeCategory.PRIMITIVE; import static io.trino.plugin.iceberg.IcebergErrorCode.ICEBERG_BAD_DATA; import static io.trino.plugin.iceberg.IcebergErrorCode.ICEBERG_CANNOT_OPEN_SPLIT; import static io.trino.plugin.iceberg.IcebergErrorCode.ICEBERG_CURSOR_ERROR; import static io.trino.plugin.iceberg.IcebergMetadataColumn.FILE_MODIFIED_TIME; import static io.trino.plugin.iceberg.IcebergMetadataColumn.FILE_PATH; +import static io.trino.plugin.iceberg.IcebergMetadataColumn.PARTITION; import static io.trino.plugin.iceberg.IcebergSessionProperties.getOrcLazyReadSmallRanges; import static io.trino.plugin.iceberg.IcebergSessionProperties.getOrcMaxBufferSize; import static io.trino.plugin.iceberg.IcebergSessionProperties.getOrcMaxMergeDistance; @@ -156,42 +151,39 @@ import static io.trino.plugin.iceberg.IcebergSessionProperties.getOrcTinyStripeThreshold; import static io.trino.plugin.iceberg.IcebergSessionProperties.getParquetMaxReadBlockRowCount; import static io.trino.plugin.iceberg.IcebergSessionProperties.getParquetMaxReadBlockSize; +import static io.trino.plugin.iceberg.IcebergSessionProperties.getParquetSmallFileThreshold; import static io.trino.plugin.iceberg.IcebergSessionProperties.isOrcBloomFiltersEnabled; import static io.trino.plugin.iceberg.IcebergSessionProperties.isOrcNestedLazy; -import static io.trino.plugin.iceberg.IcebergSessionProperties.isParquetOptimizedNestedReaderEnabled; -import static io.trino.plugin.iceberg.IcebergSessionProperties.isParquetOptimizedReaderEnabled; +import static io.trino.plugin.iceberg.IcebergSessionProperties.isParquetIgnoreStatistics; import static io.trino.plugin.iceberg.IcebergSessionProperties.isUseFileSizeFromMetadata; import static io.trino.plugin.iceberg.IcebergSessionProperties.useParquetBloomFilter; import static io.trino.plugin.iceberg.IcebergSplitManager.ICEBERG_DOMAIN_COMPACTION_THRESHOLD; +import static io.trino.plugin.iceberg.IcebergSplitSource.partitionMatchesPredicate; import static io.trino.plugin.iceberg.IcebergUtil.deserializePartitionValue; import static io.trino.plugin.iceberg.IcebergUtil.getColumnHandle; import static io.trino.plugin.iceberg.IcebergUtil.getPartitionKeys; +import static io.trino.plugin.iceberg.IcebergUtil.getPartitionValues; import static io.trino.plugin.iceberg.IcebergUtil.schemaFromHandles; -import static io.trino.plugin.iceberg.TypeConverter.ICEBERG_BINARY_TYPE; import static io.trino.plugin.iceberg.TypeConverter.ORC_ICEBERG_ID_KEY; -import static io.trino.plugin.iceberg.delete.EqualityDeleteFilter.readEqualityDeletes; -import static io.trino.plugin.iceberg.delete.PositionDeleteFilter.readPositionDeletes; -import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED; +import static io.trino.plugin.iceberg.util.OrcIcebergIds.fileColumnsByIcebergId; +import static io.trino.spi.block.PageBuilderStatus.DEFAULT_MAX_PAGE_SIZE_IN_BYTES; import static io.trino.spi.predicate.Utils.nativeValueToBlock; -import static io.trino.spi.type.BigintType.BIGINT; import static io.trino.spi.type.BooleanType.BOOLEAN; import static io.trino.spi.type.DateTimeEncoding.packDateTimeWithZone; -import static io.trino.spi.type.IntegerType.INTEGER; import static io.trino.spi.type.TimeZoneKey.UTC_KEY; -import static io.trino.spi.type.UuidType.UUID; -import static io.trino.spi.type.VarcharType.VARCHAR; +import static java.lang.Math.min; +import static java.lang.Math.toIntExact; import static java.lang.String.format; import static java.util.Locale.ENGLISH; +import static java.util.Objects.checkIndex; import static java.util.Objects.requireNonNull; +import static java.util.function.Function.identity; import static java.util.function.Predicate.not; import static java.util.stream.Collectors.groupingBy; import static java.util.stream.Collectors.mapping; -import static java.util.stream.Collectors.toList; import static java.util.stream.Collectors.toUnmodifiableList; import static org.apache.iceberg.FileContent.EQUALITY_DELETES; import static org.apache.iceberg.FileContent.POSITION_DELETES; -import static org.apache.iceberg.MetadataColumns.DELETE_FILE_PATH; -import static org.apache.iceberg.MetadataColumns.DELETE_FILE_POS; import static org.apache.iceberg.MetadataColumns.ROW_POSITION; import static org.joda.time.DateTimeZone.UTC; @@ -200,25 +192,35 @@ public class IcebergPageSourceProvider { private static final String AVRO_FIELD_ID = "field-id"; - private final TrinoFileSystemFactory fileSystemFactory; + // This is used whenever a query doesn't reference any data columns. + // We need to limit the number of rows per page in case there are projections + // in the query that can cause page sizes to explode. For example: SELECT rand() FROM some_table + // TODO (https://github.com/trinodb/trino/issues/16824) allow connector to return pages of arbitrary row count and handle this gracefully in engine + private static final int MAX_RLE_PAGE_SIZE = DEFAULT_MAX_PAGE_SIZE_IN_BYTES / SIZE_OF_LONG; + + private final IcebergFileSystemFactory fileSystemFactory; private final FileFormatDataSourceStats fileFormatDataSourceStats; private final OrcReaderOptions orcReaderOptions; private final ParquetReaderOptions parquetReaderOptions; private final TypeManager typeManager; + private final DeleteManager unpartitionedTableDeleteManager; + private final Map> partitionKeyFactories = new ConcurrentHashMap<>(); + private final Map partitionedDeleteManagers = new ConcurrentHashMap<>(); @Inject public IcebergPageSourceProvider( - TrinoFileSystemFactory fileSystemFactory, + IcebergFileSystemFactory fileSystemFactory, FileFormatDataSourceStats fileFormatDataSourceStats, - OrcReaderConfig orcReaderConfig, - ParquetReaderConfig parquetReaderConfig, + OrcReaderOptions orcReaderOptions, + ParquetReaderOptions parquetReaderOptions, TypeManager typeManager) { this.fileSystemFactory = requireNonNull(fileSystemFactory, "fileSystemFactory is null"); this.fileFormatDataSourceStats = requireNonNull(fileFormatDataSourceStats, "fileFormatDataSourceStats is null"); - this.orcReaderOptions = orcReaderConfig.toOrcReaderOptions(); - this.parquetReaderOptions = parquetReaderConfig.toParquetReaderOptions(); + this.orcReaderOptions = requireNonNull(orcReaderOptions, "orcReaderOptions is null"); + this.parquetReaderOptions = requireNonNull(parquetReaderOptions, "parquetReaderOptions is null"); this.typeManager = requireNonNull(typeManager, "typeManager is null"); + this.unpartitionedTableDeleteManager = new DeleteManager(typeManager); } @Override @@ -231,219 +233,246 @@ public ConnectorPageSource createPageSource( DynamicFilter dynamicFilter) { IcebergSplit split = (IcebergSplit) connectorSplit; - IcebergTableHandle table = (IcebergTableHandle) connectorTable; - List icebergColumns = columns.stream() .map(IcebergColumnHandle.class::cast) .collect(toImmutableList()); - - Schema tableSchema = SchemaParser.fromJson(table.getTableSchemaJson()); - - Set deleteFilterRequiredColumns = requiredColumnsForDeletes(tableSchema, split.getDeletes()); - - PartitionSpec partitionSpec = PartitionSpecParser.fromJson(tableSchema, split.getPartitionSpecJson()); + IcebergTableHandle tableHandle = (IcebergTableHandle) connectorTable; + Schema schema = SchemaParser.fromJson(tableHandle.getTableSchemaJson()); + PartitionSpec partitionSpec = PartitionSpecParser.fromJson(schema, split.getPartitionSpecJson()); org.apache.iceberg.types.Type[] partitionColumnTypes = partitionSpec.fields().stream() - .map(field -> field.transform().getResultType(tableSchema.findType(field.sourceId()))) + .map(field -> field.transform().getResultType(schema.findType(field.sourceId()))) .toArray(org.apache.iceberg.types.Type[]::new); - PartitionData partitionData = PartitionData.fromJson(split.getPartitionDataJson(), partitionColumnTypes); + + return createPageSource( + session, + icebergColumns, + schema, + partitionSpec, + PartitionData.fromJson(split.getPartitionDataJson(), partitionColumnTypes), + split.getDeletes(), + dynamicFilter, + tableHandle.getUnenforcedPredicate(), + split.getFileStatisticsDomain(), + split.getPath(), + split.getStart(), + split.getLength(), + split.getFileSize(), + split.getFileRecordCount(), + split.getPartitionDataJson(), + split.getFileFormat(), + split.getFileIoProperties(), + split.getDataSequenceNumber(), + tableHandle.getNameMappingJson().map(NameMappingParser::fromJson)); + } + + public ConnectorPageSource createPageSource( + ConnectorSession session, + List icebergColumns, + Schema tableSchema, + PartitionSpec partitionSpec, + PartitionData partitionData, + List deletes, + DynamicFilter dynamicFilter, + TupleDomain unenforcedPredicate, + TupleDomain fileStatisticsDomain, + String path, + long start, + long length, + long fileSize, + long fileRecordCount, + String partitionDataJson, + IcebergFileFormat fileFormat, + Map fileIoProperties, + long dataSequenceNumber, + Optional nameMapping) + { Map> partitionKeys = getPartitionKeys(partitionData, partitionSpec); + TupleDomain effectivePredicate = getUnenforcedPredicate( + tableSchema, + partitionKeys, + dynamicFilter, + unenforcedPredicate, + fileStatisticsDomain); + if (effectivePredicate.isNone()) { + return new EmptyPageSource(); + } + + // exit early when only reading partition keys from a simple split + String partition = partitionSpec.partitionToPath(partitionData); + TrinoFileSystem fileSystem = fileSystemFactory.create(session.getIdentity(), fileIoProperties); + TrinoInputFile inputFile = isUseFileSizeFromMetadata(session) + ? fileSystem.newInputFile(Location.of(path), fileSize) + : fileSystem.newInputFile(Location.of(path)); + try { + if (effectivePredicate.isAll() && + start == 0 && length == inputFile.length() && + deletes.isEmpty() && + icebergColumns.stream().allMatch(column -> partitionKeys.containsKey(column.getId()))) { + return generatePages( + fileRecordCount, + icebergColumns, + partitionKeys); + } + } + catch (IOException e) { + throw new UncheckedIOException(e); + } List requiredColumns = new ArrayList<>(icebergColumns); + Set deleteFilterRequiredColumns = requiredColumnsForDeletes(tableSchema, deletes); deleteFilterRequiredColumns.stream() .filter(not(icebergColumns::contains)) .forEach(requiredColumns::add); - icebergColumns.stream() - .filter(column -> column.isUpdateRowIdColumn() || column.isMergeRowIdColumn()) - .findFirst().ifPresent(rowIdColumn -> { - Set alreadyRequiredColumnIds = requiredColumns.stream() - .map(IcebergColumnHandle::getId) - .collect(toImmutableSet()); - for (ColumnIdentity identity : rowIdColumn.getColumnIdentity().getChildren()) { - if (alreadyRequiredColumnIds.contains(identity.getId())) { - // ignore - } - else if (identity.getId() == MetadataColumns.FILE_PATH.fieldId()) { - requiredColumns.add(new IcebergColumnHandle(identity, VARCHAR, ImmutableList.of(), VARCHAR, Optional.empty())); - } - else if (identity.getId() == ROW_POSITION.fieldId()) { - requiredColumns.add(new IcebergColumnHandle(identity, BIGINT, ImmutableList.of(), BIGINT, Optional.empty())); - } - else if (identity.getId() == TRINO_MERGE_PARTITION_SPEC_ID) { - requiredColumns.add(new IcebergColumnHandle(identity, INTEGER, ImmutableList.of(), INTEGER, Optional.empty())); - } - else if (identity.getId() == TRINO_MERGE_PARTITION_DATA) { - requiredColumns.add(new IcebergColumnHandle(identity, VARCHAR, ImmutableList.of(), VARCHAR, Optional.empty())); - } - else { - requiredColumns.add(getColumnHandle(tableSchema.findField(identity.getId()), typeManager)); - } - } - }); - - TupleDomain effectivePredicate = table.getUnenforcedPredicate() - .intersect(dynamicFilter.getCurrentPredicate().transformKeys(IcebergColumnHandle.class::cast)) - .simplify(ICEBERG_DOMAIN_COMPACTION_THRESHOLD); - if (effectivePredicate.isNone()) { - return new EmptyPageSource(); - } - - TrinoFileSystem fileSystem = fileSystemFactory.create(session); - TrinoInputFile inputfile = isUseFileSizeFromMetadata(session) - ? fileSystem.newInputFile(Location.of(split.getPath()), split.getFileSize()) - : fileSystem.newInputFile(Location.of(split.getPath())); - ReaderPageSourceWithRowPositions readerPageSourceWithRowPositions = createDataPageSource( session, - inputfile, - split.getStart(), - split.getLength(), + inputFile, + start, + length, + fileSize, partitionSpec.specId(), - split.getPartitionDataJson(), - split.getFileFormat(), - SchemaParser.fromJson(table.getTableSchemaJson()), + partitionDataJson, + fileFormat, + tableSchema, requiredColumns, effectivePredicate, - table.getNameMappingJson().map(NameMappingParser::fromJson), + nameMapping, + partition, partitionKeys); - ReaderPageSource dataPageSource = readerPageSourceWithRowPositions.getReaderPageSource(); - - Optional projectionsAdapter = dataPageSource.getReaderColumns().map(readerColumns -> - new ReaderProjectionsAdapter( - requiredColumns, - readerColumns, - column -> ((IcebergColumnHandle) column).getType(), - IcebergPageSourceProvider::applyProjection)); - - List readColumns = dataPageSource.getReaderColumns() - .map(readerColumns -> readerColumns.get().stream().map(IcebergColumnHandle.class::cast).collect(toList())) - .orElse(requiredColumns); - - Supplier> deletePredicate = Suppliers.memoize(() -> { - List deleteFilters = readDeletes( - session, - tableSchema, - split.getPath(), - split.getDeletes(), - readerPageSourceWithRowPositions.getStartRowPosition(), - readerPageSourceWithRowPositions.getEndRowPosition()); - return deleteFilters.stream() - .map(filter -> filter.createPredicate(readColumns)) - .reduce(RowPredicate::and); - }); - return new IcebergPageSource( - icebergColumns, - requiredColumns, - dataPageSource.get(), - projectionsAdapter, - deletePredicate); + ConnectorPageSource pageSource = readerPageSourceWithRowPositions.pageSource(); + + // filter out deleted rows + if (!deletes.isEmpty()) { + Supplier> deletePredicate = memoize(() -> getDeleteManager(partitionSpec, partitionData) + .getDeletePredicate( + path, + dataSequenceNumber, + deletes, + requiredColumns, + tableSchema, + readerPageSourceWithRowPositions, + (deleteFile, deleteColumns, tupleDomain) -> openDeletes(session, fileSystem, deleteFile, deleteColumns, tupleDomain))); + pageSource = TransformConnectorPageSource.create(pageSource, page -> { + try { + Optional rowPredicate = deletePredicate.get(); + if (rowPredicate.isPresent()) { + page = rowPredicate.get().filterPage(page); + } + if (icebergColumns.size() == page.getChannelCount()) { + return page; + } + return new PrefixColumnsPage(page, icebergColumns.size()).getPage(); + } + catch (RuntimeException e) { + throwIfInstanceOf(e, TrinoException.class); + throw new TrinoException(ICEBERG_BAD_DATA, e); + } + }); + } + return pageSource; } - private Set requiredColumnsForDeletes(Schema schema, List deletes) + private DeleteManager getDeleteManager(PartitionSpec partitionSpec, PartitionData partitionData) { - ImmutableSet.Builder requiredColumns = ImmutableSet.builder(); - for (DeleteFile deleteFile : deletes) { - if (deleteFile.content() == POSITION_DELETES) { - requiredColumns.add(getColumnHandle(ROW_POSITION, typeManager)); - } - else if (deleteFile.content() == EQUALITY_DELETES) { - deleteFile.equalityFieldIds().stream() - .map(id -> getColumnHandle(schema.findField(id), typeManager)) - .forEach(requiredColumns::add); - } + if (partitionSpec.isUnpartitioned()) { + return unpartitionedTableDeleteManager; } - return requiredColumns.build(); + Types.StructType structType = partitionSpec.partitionType(); + PartitionKey partitionKey = partitionKeyFactories.computeIfAbsent( + partitionSpec.specId(), + key -> { + // creating the template wrapper is expensive, reuse it for all partitions of the same spec + // reuse is only safe because we only use the copyFor method which is thread safe + StructLikeWrapper templateWrapper = StructLikeWrapper.forType(structType); + return data -> new PartitionKey(key, templateWrapper.copyFor(data)); + }) + .apply(partitionData); + + return partitionedDeleteManagers.computeIfAbsent(partitionKey, ignored -> new DeleteManager(typeManager)); } - private List readDeletes( - ConnectorSession session, - Schema schema, - String dataFilePath, - List deleteFiles, - Optional startRowPosition, - Optional endRowPosition) + private record PartitionKey(int specId, StructLikeWrapper partitionData) {} + + private TupleDomain getUnenforcedPredicate( + Schema tableSchema, + Map> partitionKeys, + DynamicFilter dynamicFilter, + TupleDomain unenforcedPredicate, + TupleDomain fileStatisticsDomain) { - verify(startRowPosition.isPresent() == endRowPosition.isPresent(), "startRowPosition and endRowPosition must be specified together"); - - Slice targetPath = utf8Slice(dataFilePath); - List filters = new ArrayList<>(); - LongBitmapDataProvider deletedRows = new Roaring64Bitmap(); - - IcebergColumnHandle deleteFilePath = getColumnHandle(DELETE_FILE_PATH, typeManager); - IcebergColumnHandle deleteFilePos = getColumnHandle(DELETE_FILE_POS, typeManager); - List deleteColumns = ImmutableList.of(deleteFilePath, deleteFilePos); - TupleDomain deleteDomain = TupleDomain.fromFixedValues(ImmutableMap.of(deleteFilePath, NullableValue.of(VARCHAR, targetPath))); - if (startRowPosition.isPresent()) { - Range positionRange = Range.range(deleteFilePos.getType(), startRowPosition.get(), true, endRowPosition.get(), true); - TupleDomain positionDomain = TupleDomain.withColumnDomains(ImmutableMap.of(deleteFilePos, Domain.create(ValueSet.ofRanges(positionRange), false))); - deleteDomain = deleteDomain.intersect(positionDomain); - } + return prunePredicate( + tableSchema, + partitionKeys, + // We reach here when we could not prune the split using file level stats, table predicate + // and the dynamic filter in the coordinator during split generation. The file level stats + // in IcebergSplit#fileStatisticsDomain could help to prune this split when a more selective dynamic filter + // is available now, without having to access parquet/orc file footer for row-group/stripe stats. + TupleDomain.intersect(ImmutableList.of( + unenforcedPredicate, + fileStatisticsDomain, + dynamicFilter.getCurrentPredicate().transformKeys(IcebergColumnHandle.class::cast))), + fileStatisticsDomain) + .simplify(ICEBERG_DOMAIN_COMPACTION_THRESHOLD); + } - for (DeleteFile delete : deleteFiles) { - if (delete.content() == POSITION_DELETES) { - if (startRowPosition.isPresent()) { - byte[] lowerBoundBytes = delete.getLowerBounds().get(DELETE_FILE_POS.fieldId()); - Optional positionLowerBound = Optional.ofNullable(lowerBoundBytes) - .map(bytes -> Conversions.fromByteBuffer(DELETE_FILE_POS.type(), ByteBuffer.wrap(bytes))); + private TupleDomain prunePredicate( + Schema tableSchema, + Map> partitionKeys, + TupleDomain unenforcedPredicate, + TupleDomain fileStatisticsDomain) + { + if (unenforcedPredicate.isAll() || unenforcedPredicate.isNone()) { + return unenforcedPredicate; + } - byte[] upperBoundBytes = delete.getUpperBounds().get(DELETE_FILE_POS.fieldId()); - Optional positionUpperBound = Optional.ofNullable(upperBoundBytes) - .map(bytes -> Conversions.fromByteBuffer(DELETE_FILE_POS.type(), ByteBuffer.wrap(bytes))); + Set partitionColumns = partitionKeys.keySet().stream() + .map(fieldId -> getColumnHandle(tableSchema.findField(fieldId), typeManager)) + .collect(toImmutableSet()); + Supplier> partitionValues = memoize(() -> getPartitionValues(partitionColumns, partitionKeys)); + if (!partitionMatchesPredicate(partitionColumns, partitionValues, unenforcedPredicate)) { + return TupleDomain.none(); + } - if ((positionLowerBound.isPresent() && positionLowerBound.get() > endRowPosition.get()) || - (positionUpperBound.isPresent() && positionUpperBound.get() < startRowPosition.get())) { - continue; - } - } + return unenforcedPredicate + // Filter out partition columns domains from the dynamic filter because they should be irrelevant at data file level + .filter((columnHandle, ignore) -> !partitionKeys.containsKey(columnHandle.getId())) + // remove domains from predicate that fully contain split data because they are irrelevant for filtering + .filter((handle, domain) -> !domain.contains(fileStatisticsDomain.getDomain(handle, domain.getType()))); + } - try (ConnectorPageSource pageSource = openDeletes(session, delete, deleteColumns, deleteDomain)) { - readPositionDeletes(pageSource, targetPath, deletedRows); - } - catch (IOException e) { - throw new UncheckedIOException(e); - } + private Set requiredColumnsForDeletes(Schema schema, List deletes) + { + ImmutableSet.Builder requiredColumns = ImmutableSet.builder(); + for (DeleteFile deleteFile : deletes) { + if (deleteFile.content() == POSITION_DELETES) { + requiredColumns.add(getColumnHandle(ROW_POSITION, typeManager)); } - else if (delete.content() == EQUALITY_DELETES) { - List fieldIds = delete.equalityFieldIds(); - verify(!fieldIds.isEmpty(), "equality field IDs are missing"); - List columns = fieldIds.stream() + else if (deleteFile.content() == EQUALITY_DELETES) { + deleteFile.equalityFieldIds().stream() .map(id -> getColumnHandle(schema.findField(id), typeManager)) - .collect(toImmutableList()); - - try (ConnectorPageSource pageSource = openDeletes(session, delete, columns, TupleDomain.all())) { - filters.add(readEqualityDeletes(pageSource, columns, schema)); - } - catch (IOException e) { - throw new UncheckedIOException(e); - } - } - else { - throw new VerifyException("Unknown delete content: " + delete.content()); + .forEach(requiredColumns::add); } } - if (!deletedRows.isEmpty()) { - filters.add(new PositionDeleteFilter(deletedRows)); - } - - return filters; + return requiredColumns.build(); } private ConnectorPageSource openDeletes( ConnectorSession session, + TrinoFileSystem fileSystem, DeleteFile delete, List columns, TupleDomain tupleDomain) { - TrinoFileSystem fileSystem = fileSystemFactory.create(session); return createDataPageSource( session, fileSystem.newInputFile(Location.of(delete.path()), delete.fileSizeInBytes()), 0, delete.fileSizeInBytes(), + delete.fileSizeInBytes(), 0, "", IcebergFileFormat.fromIceberg(delete.format()), @@ -451,16 +480,17 @@ private ConnectorPageSource openDeletes( columns, tupleDomain, Optional.empty(), + "", ImmutableMap.of()) - .getReaderPageSource() - .get(); + .pageSource(); } - public ReaderPageSourceWithRowPositions createDataPageSource( + private ReaderPageSourceWithRowPositions createDataPageSource( ConnectorSession session, TrinoInputFile inputFile, long start, long length, + long fileSize, int partitionSpecId, String partitionData, IcebergFileFormat fileFormat, @@ -468,62 +498,99 @@ public ReaderPageSourceWithRowPositions createDataPageSource( List dataColumns, TupleDomain predicate, Optional nameMapping, + String partition, Map> partitionKeys) { - switch (fileFormat) { - case ORC: - return createOrcPageSource( - inputFile, - start, - length, - partitionSpecId, - partitionData, - dataColumns, - predicate, - orcReaderOptions - .withMaxMergeDistance(getOrcMaxMergeDistance(session)) - .withMaxBufferSize(getOrcMaxBufferSize(session)) - .withStreamBufferSize(getOrcStreamBufferSize(session)) - .withTinyStripeThreshold(getOrcTinyStripeThreshold(session)) - .withMaxReadBlockSize(getOrcMaxReadBlockSize(session)) - .withLazyReadSmallRanges(getOrcLazyReadSmallRanges(session)) - .withNestedLazy(isOrcNestedLazy(session)) - .withBloomFiltersEnabled(isOrcBloomFiltersEnabled(session)), - fileFormatDataSourceStats, - typeManager, - nameMapping, - partitionKeys); - case PARQUET: - return createParquetPageSource( - inputFile, - start, - length, - partitionSpecId, - partitionData, - dataColumns, - parquetReaderOptions - .withMaxReadBlockSize(getParquetMaxReadBlockSize(session)) - .withMaxReadBlockRowCount(getParquetMaxReadBlockRowCount(session)) - .withBatchColumnReaders(isParquetOptimizedReaderEnabled(session)) - .withBloomFilter(useParquetBloomFilter(session)) - .withBatchNestedColumnReaders(isParquetOptimizedNestedReaderEnabled(session)), - predicate, - fileFormatDataSourceStats, - nameMapping, - partitionKeys); - case AVRO: - return createAvroPageSource( - inputFile, - start, - length, - partitionSpecId, - partitionData, - fileSchema, - nameMapping, - dataColumns); - default: - throw new TrinoException(NOT_SUPPORTED, "File format not supported for Iceberg: " + fileFormat); + return switch (fileFormat) { + case ORC -> createOrcPageSource( + inputFile, + start, + length, + partitionSpecId, + partitionData, + dataColumns, + predicate, + orcReaderOptions + .withMaxMergeDistance(getOrcMaxMergeDistance(session)) + .withMaxBufferSize(getOrcMaxBufferSize(session)) + .withStreamBufferSize(getOrcStreamBufferSize(session)) + .withTinyStripeThreshold(getOrcTinyStripeThreshold(session)) + .withMaxReadBlockSize(getOrcMaxReadBlockSize(session)) + .withLazyReadSmallRanges(getOrcLazyReadSmallRanges(session)) + .withNestedLazy(isOrcNestedLazy(session)) + .withBloomFiltersEnabled(isOrcBloomFiltersEnabled(session)), + fileFormatDataSourceStats, + typeManager, + nameMapping, + partition, + partitionKeys); + case PARQUET -> createParquetPageSource( + inputFile, + start, + length, + fileSize, + partitionSpecId, + partitionData, + dataColumns, + parquetReaderOptions + .withMaxReadBlockSize(getParquetMaxReadBlockSize(session)) + .withMaxReadBlockRowCount(getParquetMaxReadBlockRowCount(session)) + .withSmallFileThreshold(getParquetSmallFileThreshold(session)) + .withIgnoreStatistics(isParquetIgnoreStatistics(session)) + .withBloomFilter(useParquetBloomFilter(session)) + // TODO https://github.com/trinodb/trino/issues/11000 + .withUseColumnIndex(false), + predicate, + fileFormatDataSourceStats, + nameMapping, + partition, + partitionKeys); + case AVRO -> createAvroPageSource( + inputFile, + start, + length, + partitionSpecId, + partitionData, + fileSchema, + nameMapping, + partition, + dataColumns); + }; + } + + private static ConnectorPageSource generatePages( + long totalRowCount, + List icebergColumns, + Map> partitionKeys) + { + int maxPageSize = MAX_RLE_PAGE_SIZE; + Block[] pageBlocks = new Block[icebergColumns.size()]; + for (int i = 0; i < icebergColumns.size(); i++) { + IcebergColumnHandle column = icebergColumns.get(i); + Type trinoType = column.getType(); + Object partitionValue = deserializePartitionValue(trinoType, partitionKeys.get(column.getId()).orElse(null), column.getName()); + pageBlocks[i] = RunLengthEncodedBlock.create(nativeValueToBlock(trinoType, partitionValue), maxPageSize); } + Page maxPage = new Page(maxPageSize, pageBlocks); + + return new FixedPageSource( + new AbstractIterator<>() + { + private long rowIndex; + + @Override + protected Page computeNext() + { + if (rowIndex == totalRowCount) { + return endOfData(); + } + int pageSize = toIntExact(min(maxPageSize, totalRowCount - rowIndex)); + Page page = maxPage.getRegion(0, pageSize); + rowIndex += pageSize; + return page; + } + }, + maxPage.getRetainedSizeInBytes()); } private static ReaderPageSourceWithRowPositions createOrcPageSource( @@ -538,6 +605,7 @@ private static ReaderPageSourceWithRowPositions createOrcPageSource( FileFormatDataSourceStats stats, TypeManager typeManager, Optional nameMapping, + String partition, Map> partitionKeys) { OrcDataSource orcDataSource = null; @@ -547,92 +615,90 @@ private static ReaderPageSourceWithRowPositions createOrcPageSource( OrcReader reader = OrcReader.createOrcReader(orcDataSource, options) .orElseThrow(() -> new TrinoException(ICEBERG_BAD_DATA, "ORC file is zero length")); - List fileColumns = reader.getRootColumn().getNestedColumns(); - if (nameMapping.isPresent() && !hasIds(reader.getRootColumn())) { - fileColumns = fileColumns.stream() - .map(orcColumn -> setMissingFieldIds(orcColumn, nameMapping.get(), ImmutableList.of(orcColumn.getColumnName()))) - .collect(toImmutableList()); - } - - Map fileColumnsByIcebergId = mapIdsToOrcFileColumns(fileColumns); + Map fileColumnsByIcebergId = fileColumnsByIcebergId(reader, nameMapping); TupleDomainOrcPredicateBuilder predicateBuilder = TupleDomainOrcPredicate.builder() .setBloomFiltersEnabled(options.isBloomFiltersEnabled()); Map effectivePredicateDomains = effectivePredicate.getDomains() .orElseThrow(() -> new IllegalArgumentException("Effective predicate is none")); + for (IcebergColumnHandle column : columns) { + for (Map.Entry domainEntry : effectivePredicateDomains.entrySet()) { + IcebergColumnHandle predicateColumn = domainEntry.getKey(); + OrcColumn predicateOrcColumn = fileColumnsByIcebergId.get(predicateColumn.getId()); + if (predicateOrcColumn != null && column.getBaseColumnIdentity().equals(predicateColumn.getBaseColumnIdentity())) { + predicateBuilder.addColumn(predicateOrcColumn.getColumnId(), domainEntry.getValue()); + } + } + } - Optional baseColumnProjections = projectBaseColumns(columns); Map>> projectionsByFieldId = columns.stream() .collect(groupingBy( column -> column.getBaseColumnIdentity().getId(), mapping(IcebergColumnHandle::getPath, toUnmodifiableList()))); - List readBaseColumns = baseColumnProjections - .map(readerColumns -> (List) readerColumns.get().stream().map(IcebergColumnHandle.class::cast).collect(toImmutableList())) - .orElse(columns); - List fileReadColumns = new ArrayList<>(readBaseColumns.size()); - List fileReadTypes = new ArrayList<>(readBaseColumns.size()); - List projectedLayouts = new ArrayList<>(readBaseColumns.size()); - List columnAdaptations = new ArrayList<>(readBaseColumns.size()); - - for (IcebergColumnHandle column : readBaseColumns) { - verify(column.isBaseColumn(), "Column projections must be based from a root column"); - OrcColumn orcColumn = fileColumnsByIcebergId.get(column.getId()); + List baseColumns = new ArrayList<>(columns.size()); + Map baseColumnIdToOrdinal = new HashMap<>(); + List fileReadColumns = new ArrayList<>(columns.size()); + List fileReadTypes = new ArrayList<>(columns.size()); + List projectedLayouts = new ArrayList<>(columns.size()); + TransformConnectorPageSource.Builder transforms = TransformConnectorPageSource.builder(); + boolean appendRowNumberColumn = false; + for (IcebergColumnHandle column : columns) { if (column.isIsDeletedColumn()) { - columnAdaptations.add(ColumnAdaptation.constantColumn(nativeValueToBlock(BOOLEAN, false))); + transforms.constantValue(nativeValueToBlock(BOOLEAN, false)); } else if (partitionKeys.containsKey(column.getId())) { Type trinoType = column.getType(); - columnAdaptations.add(ColumnAdaptation.constantColumn(nativeValueToBlock( + transforms.constantValue(nativeValueToBlock( trinoType, - deserializePartitionValue(trinoType, partitionKeys.get(column.getId()).orElse(null), column.getName())))); + deserializePartitionValue(trinoType, partitionKeys.get(column.getId()).orElse(null), column.getName()))); + } + else if (column.isPartitionColumn()) { + transforms.constantValue(nativeValueToBlock(PARTITION.getType(), utf8Slice(partition))); } else if (column.isPathColumn()) { - columnAdaptations.add(ColumnAdaptation.constantColumn(nativeValueToBlock(FILE_PATH.getType(), utf8Slice(inputFile.location().toString())))); + transforms.constantValue(nativeValueToBlock(FILE_PATH.getType(), utf8Slice(inputFile.location().toString()))); } else if (column.isFileModifiedTimeColumn()) { - columnAdaptations.add(ColumnAdaptation.constantColumn(nativeValueToBlock(FILE_MODIFIED_TIME.getType(), packDateTimeWithZone(inputFile.lastModified().toEpochMilli(), UTC_KEY)))); + transforms.constantValue(nativeValueToBlock(FILE_MODIFIED_TIME.getType(), packDateTimeWithZone(inputFile.lastModified().toEpochMilli(), UTC_KEY))); } - else if (column.isUpdateRowIdColumn() || column.isMergeRowIdColumn()) { - // $row_id is a composite of multiple physical columns. It is assembled by the IcebergPageSource - columnAdaptations.add(ColumnAdaptation.nullColumn(column.getType())); + else if (column.isMergeRowIdColumn()) { + appendRowNumberColumn = true; + transforms.transform(MergeRowIdTransform.create(utf8Slice(inputFile.location().toString()), partitionSpecId, utf8Slice(partitionData))); } else if (column.isRowPositionColumn()) { - columnAdaptations.add(ColumnAdaptation.positionColumn()); + appendRowNumberColumn = true; + transforms.transform(new GetRowPositionFromSource()); } - else if (column.getId() == TRINO_MERGE_PARTITION_SPEC_ID) { - columnAdaptations.add(ColumnAdaptation.constantColumn(nativeValueToBlock(column.getType(), (long) partitionSpecId))); + else if (!fileColumnsByIcebergId.containsKey(column.getBaseColumnIdentity().getId())) { + transforms.constantValue(column.getType().createNullBlock()); } - else if (column.getId() == TRINO_MERGE_PARTITION_DATA) { - columnAdaptations.add(ColumnAdaptation.constantColumn(nativeValueToBlock(column.getType(), utf8Slice(partitionData)))); - } - else if (orcColumn != null) { - Type readType = getOrcReadType(column.getType(), typeManager); - - if (column.getType() == UUID && !"UUID".equals(orcColumn.getAttributes().get(ICEBERG_BINARY_TYPE))) { - throw new TrinoException(ICEBERG_BAD_DATA, format("Expected ORC column for UUID data to be annotated with %s=UUID: %s", ICEBERG_BINARY_TYPE, orcColumn)); + else { + IcebergColumnHandle baseColumn = column.getBaseColumn(); + Integer ordinal = baseColumnIdToOrdinal.get(baseColumn.getId()); + if (ordinal == null) { + ordinal = baseColumns.size(); + baseColumns.add(baseColumn); + baseColumnIdToOrdinal.put(baseColumn.getId(), ordinal); + + OrcColumn orcBaseColumn = requireNonNull(fileColumnsByIcebergId.get(baseColumn.getId())); + fileReadColumns.add(orcBaseColumn); + fileReadTypes.add(getOrcReadType(baseColumn.getType(), typeManager)); + projectedLayouts.add(IcebergOrcProjectedLayout.createProjectedLayout( + orcBaseColumn, + projectionsByFieldId.get(baseColumn.getId()))); } - List> fieldIdProjections = projectionsByFieldId.get(column.getId()); - ProjectedLayout projectedLayout = IcebergOrcProjectedLayout.createProjectedLayout(orcColumn, fieldIdProjections); - - int sourceIndex = fileReadColumns.size(); - columnAdaptations.add(ColumnAdaptation.sourceColumn(sourceIndex)); - fileReadColumns.add(orcColumn); - fileReadTypes.add(readType); - projectedLayouts.add(projectedLayout); - - for (Map.Entry domainEntry : effectivePredicateDomains.entrySet()) { - IcebergColumnHandle predicateColumn = domainEntry.getKey(); - OrcColumn predicateOrcColumn = fileColumnsByIcebergId.get(predicateColumn.getId()); - if (predicateOrcColumn != null && column.getColumnIdentity().equals(predicateColumn.getBaseColumnIdentity())) { - predicateBuilder.addColumn(predicateOrcColumn.getColumnId(), domainEntry.getValue()); - } + if (column.isBaseColumn()) { + transforms.column(ordinal); + } + else { + transforms.dereferenceField(ImmutableList.builder() + .add(ordinal) + .addAll(applyProjection(column, baseColumn)) + .build()); } - } - else { - columnAdaptations.add(ColumnAdaptation.nullColumn(column.getType())); } } @@ -649,20 +715,21 @@ else if (orcColumn != null) { memoryUsage, INITIAL_BATCH_SIZE, exception -> handleException(orcDataSourceId, exception), - new IdBasedFieldMapperFactory(readBaseColumns)); + new IdBasedFieldMapperFactory(baseColumns)); + + ConnectorPageSource pageSource = new OrcPageSource( + recordReader, + orcDataSource, + Optional.empty(), + Optional.empty(), + memoryUsage, + stats, + reader.getCompressionKind()); + + pageSource = transforms.build(pageSource); return new ReaderPageSourceWithRowPositions( - new ReaderPageSource( - new OrcPageSource( - recordReader, - columnAdaptations, - orcDataSource, - Optional.empty(), - Optional.empty(), - memoryUsage, - stats, - reader.getCompressionKind()), - baseColumnProjections), + pageSource, recordReader.getStartRowPosition(), recordReader.getEndRowPosition()); } @@ -677,8 +744,8 @@ else if (orcColumn != null) { } } } - if (e instanceof TrinoException) { - throw (TrinoException) e; + if (e instanceof TrinoException trinoException) { + throw trinoException; } if (e instanceof OrcCorruptionException) { throw new TrinoException(ICEBERG_BAD_DATA, e); @@ -688,48 +755,6 @@ else if (orcColumn != null) { } } - private static boolean hasIds(OrcColumn column) - { - if (column.getAttributes().containsKey(ORC_ICEBERG_ID_KEY)) { - return true; - } - - return column.getNestedColumns().stream().anyMatch(IcebergPageSourceProvider::hasIds); - } - - private static OrcColumn setMissingFieldIds(OrcColumn column, NameMapping nameMapping, List qualifiedPath) - { - MappedField mappedField = nameMapping.find(qualifiedPath); - - ImmutableMap.Builder attributes = ImmutableMap.builder() - .putAll(column.getAttributes()); - if (mappedField != null && mappedField.id() != null) { - attributes.put(ORC_ICEBERG_ID_KEY, String.valueOf(mappedField.id())); - } - - return new OrcColumn( - column.getPath(), - column.getColumnId(), - column.getColumnName(), - column.getColumnType(), - column.getOrcDataSourceId(), - column.getNestedColumns().stream() - .map(nestedColumn -> { - ImmutableList.Builder nextQualifiedPath = ImmutableList.builder() - .addAll(qualifiedPath); - if (column.getColumnType() == OrcType.OrcTypeKind.LIST) { - // The Trino ORC reader uses "item" for list element names, but the NameMapper expects "element" - nextQualifiedPath.add("element"); - } - else { - nextQualifiedPath.add(nestedColumn.getColumnName()); - } - return setMissingFieldIds(nestedColumn, nameMapping, nextQualifiedPath.build()); - }) - .collect(toImmutableList()), - attributes.buildOrThrow()); - } - /** * Gets the index based dereference chain to get from the readColumnHandle to the expectedColumnHandle */ @@ -750,20 +775,6 @@ private static List applyProjection(ColumnHandle expectedColumnHandle, return dereferenceChain.build(); } - private static Map mapIdsToOrcFileColumns(List columns) - { - ImmutableMap.Builder columnsById = ImmutableMap.builder(); - Traverser.forTree(OrcColumn::getNestedColumns) - .depthFirstPreOrder(columns) - .forEach(column -> { - String fieldId = column.getAttributes().get(ORC_ICEBERG_ID_KEY); - if (fieldId != null) { - columnsById.put(Integer.parseInt(fieldId), column); - } - }); - return columnsById.buildOrThrow(); - } - private static Integer getIcebergFieldId(OrcColumn column) { String icebergId = column.getAttributes().get(ORC_ICEBERG_ID_KEY); @@ -802,7 +813,7 @@ public IdBasedFieldMapperFactory(List columns) ImmutableMap.Builder> mapping = ImmutableMap.builder(); for (IcebergColumnHandle column : columns) { - if (column.isUpdateRowIdColumn() || column.isMergeRowIdColumn()) { + if (column.isMergeRowIdColumn()) { // The update $row_id column contains fields which should not be accounted for in the mapping. continue; } @@ -868,22 +879,24 @@ private static ReaderPageSourceWithRowPositions createParquetPageSource( TrinoInputFile inputFile, long start, long length, + long fileSize, int partitionSpecId, String partitionData, - List regularColumns, + List columns, ParquetReaderOptions options, TupleDomain effectivePredicate, FileFormatDataSourceStats fileFormatDataSourceStats, Optional nameMapping, + String partition, Map> partitionKeys) { AggregatedMemoryContext memoryContext = newSimpleAggregatedMemoryContext(); ParquetDataSource dataSource = null; try { - dataSource = new TrinoParquetDataSource(inputFile, options, fileFormatDataSourceStats); + dataSource = createDataSource(inputFile, OptionalLong.of(fileSize), options, memoryContext, fileFormatDataSourceStats); ParquetMetadata parquetMetadata = MetadataReader.readFooter(dataSource, Optional.empty()); - FileMetaData fileMetaData = parquetMetadata.getFileMetaData(); + FileMetadata fileMetaData = parquetMetadata.getFileMetaData(); MessageType fileSchema = fileMetaData.getSchema(); if (nameMapping.isPresent() && !ParquetSchemaUtil.hasIds(fileSchema)) { // NameMapping conversion is necessary because MetadataReader converts all column names to lowercase and NameMapping is case sensitive @@ -891,114 +904,123 @@ private static ReaderPageSourceWithRowPositions createParquetPageSource( } // Mapping from Iceberg field ID to Parquet fields. - Map parquetIdToField = createParquetIdToFieldMapping(fileSchema); - - Optional baseColumnProjections = projectBaseColumns(regularColumns); - List readBaseColumns = baseColumnProjections - .map(readerColumns -> (List) readerColumns.get().stream().map(IcebergColumnHandle.class::cast).collect(toImmutableList())) - .orElse(regularColumns); - - List parquetFields = readBaseColumns.stream() - .map(column -> parquetIdToField.get(column.getId())) - .collect(toList()); + Map parquetIdToFieldName = createParquetIdToFieldMapping(fileSchema); - MessageType requestedSchema = getMessageType(regularColumns, fileSchema.getName(), parquetIdToField); + MessageType requestedSchema = getMessageType(columns, fileSchema.getName(), parquetIdToFieldName); Map, ColumnDescriptor> descriptorsByPath = getDescriptors(fileSchema, requestedSchema); - TupleDomain parquetTupleDomain = getParquetTupleDomain(descriptorsByPath, effectivePredicate); + TupleDomain parquetTupleDomain = options.isIgnoreStatistics() ? TupleDomain.all() : getParquetTupleDomain(descriptorsByPath, effectivePredicate); TupleDomainParquetPredicate parquetPredicate = buildPredicate(requestedSchema, parquetTupleDomain, descriptorsByPath, UTC); - long nextStart = 0; - Optional startRowPosition = Optional.empty(); - Optional endRowPosition = Optional.empty(); - ImmutableList.Builder blockStarts = ImmutableList.builder(); - List blocks = new ArrayList<>(); - for (BlockMetaData block : parquetMetadata.getBlocks()) { - long firstDataPage = block.getColumns().get(0).getFirstDataPageOffset(); - Optional bloomFilterStore = getBloomFilterStore(dataSource, block, parquetTupleDomain, options); - - if (start <= firstDataPage && firstDataPage < start + length && - predicateMatches(parquetPredicate, block, dataSource, descriptorsByPath, parquetTupleDomain, Optional.empty(), bloomFilterStore, UTC, ICEBERG_DOMAIN_COMPACTION_THRESHOLD)) { - blocks.add(block); - blockStarts.add(nextStart); - if (startRowPosition.isEmpty()) { - startRowPosition = Optional.of(nextStart); - } - endRowPosition = Optional.of(nextStart + block.getRowCount()); - } - nextStart += block.getRowCount(); - } - MessageColumnIO messageColumnIO = getColumnIO(fileSchema, requestedSchema); - ParquetPageSource.Builder pageSourceBuilder = ParquetPageSource.builder(); - int parquetSourceChannel = 0; - - ImmutableList.Builder parquetColumnFieldsBuilder = ImmutableList.builder(); - for (int columnIndex = 0; columnIndex < readBaseColumns.size(); columnIndex++) { - IcebergColumnHandle column = readBaseColumns.get(columnIndex); + Map baseColumnIdToOrdinal = new HashMap<>(); + TransformConnectorPageSource.Builder transforms = TransformConnectorPageSource.builder(); + boolean appendRowNumberColumn = false; + int nextOrdinal = 0; + ImmutableList.Builder parquetColumnFieldsBuilder = ImmutableList.builder(); + for (IcebergColumnHandle column : columns) { if (column.isIsDeletedColumn()) { - pageSourceBuilder.addConstantColumn(nativeValueToBlock(BOOLEAN, false)); + transforms.constantValue(nativeValueToBlock(BOOLEAN, false)); } else if (partitionKeys.containsKey(column.getId())) { Type trinoType = column.getType(); - pageSourceBuilder.addConstantColumn(nativeValueToBlock( + transforms.constantValue(nativeValueToBlock( trinoType, deserializePartitionValue(trinoType, partitionKeys.get(column.getId()).orElse(null), column.getName()))); } + else if (column.isPartitionColumn()) { + transforms.constantValue(nativeValueToBlock(PARTITION.getType(), utf8Slice(partition))); + } else if (column.isPathColumn()) { - pageSourceBuilder.addConstantColumn(nativeValueToBlock(FILE_PATH.getType(), utf8Slice(inputFile.location().toString()))); + transforms.constantValue(nativeValueToBlock(FILE_PATH.getType(), utf8Slice(inputFile.location().toString()))); } else if (column.isFileModifiedTimeColumn()) { - pageSourceBuilder.addConstantColumn(nativeValueToBlock(FILE_MODIFIED_TIME.getType(), packDateTimeWithZone(inputFile.lastModified().toEpochMilli(), UTC_KEY))); + transforms.constantValue(nativeValueToBlock(FILE_MODIFIED_TIME.getType(), packDateTimeWithZone(inputFile.lastModified().toEpochMilli(), UTC_KEY))); } - else if (column.isUpdateRowIdColumn() || column.isMergeRowIdColumn()) { - // $row_id is a composite of multiple physical columns, it is assembled by the IcebergPageSource - pageSourceBuilder.addNullColumn(column.getType()); + else if (column.isMergeRowIdColumn()) { + appendRowNumberColumn = true; + transforms.transform(MergeRowIdTransform.create(utf8Slice(inputFile.location().toString()), partitionSpecId, utf8Slice(partitionData))); } else if (column.isRowPositionColumn()) { - pageSourceBuilder.addRowIndexColumn(); + appendRowNumberColumn = true; + transforms.transform(new GetRowPositionFromSource()); } - else if (column.getId() == TRINO_MERGE_PARTITION_SPEC_ID) { - pageSourceBuilder.addConstantColumn(nativeValueToBlock(column.getType(), (long) partitionSpecId)); - } - else if (column.getId() == TRINO_MERGE_PARTITION_DATA) { - pageSourceBuilder.addConstantColumn(nativeValueToBlock(column.getType(), utf8Slice(partitionData))); + else if (!parquetIdToFieldName.containsKey(column.getBaseColumn().getId())) { + transforms.constantValue(column.getType().createNullBlock()); } else { - org.apache.parquet.schema.Type parquetField = parquetFields.get(columnIndex); - Type trinoType = column.getBaseType(); - if (parquetField == null) { - pageSourceBuilder.addNullColumn(trinoType); - continue; + IcebergColumnHandle baseColumn = column.getBaseColumn(); + Integer ordinal = baseColumnIdToOrdinal.get(baseColumn.getId()); + if (ordinal == null) { + String parquetFieldName = requireNonNull(parquetIdToFieldName.get(baseColumn.getId())).getName(); + + // The top level columns are already mapped by name/id appropriately. + Optional field = IcebergParquetColumnIOConverter.constructField( + new FieldContext(baseColumn.getType(), baseColumn.getColumnIdentity()), + messageColumnIO.getChild(parquetFieldName)); + if (field.isEmpty()) { + // base column is missing so return a null + transforms.constantValue(column.getType().createNullBlock()); + continue; + } + + ordinal = nextOrdinal; + nextOrdinal++; + baseColumnIdToOrdinal.put(baseColumn.getId(), ordinal); + + parquetColumnFieldsBuilder.add(new Column(parquetFieldName, field.get())); } - // The top level columns are already mapped by name/id appropriately. - ColumnIO columnIO = messageColumnIO.getChild(parquetField.getName()); - Optional field = IcebergParquetColumnIOConverter.constructField(new FieldContext(trinoType, column.getColumnIdentity()), columnIO); - if (field.isEmpty()) { - pageSourceBuilder.addNullColumn(trinoType); - continue; + if (column.isBaseColumn()) { + transforms.column(ordinal); + } + else { + transforms.dereferenceField(ImmutableList.builder() + .add(ordinal) + .addAll(applyProjection(column, baseColumn)) + .build()); } - parquetColumnFieldsBuilder.add(field.get()); - pageSourceBuilder.addSourceColumn(parquetSourceChannel); - parquetSourceChannel++; } } + List rowGroups = getFilteredRowGroups( + start, + length, + dataSource, + parquetMetadata, + ImmutableList.of(parquetTupleDomain), + ImmutableList.of(parquetPredicate), + descriptorsByPath, + UTC, + ICEBERG_DOMAIN_COMPACTION_THRESHOLD, + options); + ParquetDataSourceId dataSourceId = dataSource.getId(); ParquetReader parquetReader = new ParquetReader( Optional.ofNullable(fileMetaData.getCreatedBy()), parquetColumnFieldsBuilder.build(), - blocks, - blockStarts.build(), + appendRowNumberColumn, + rowGroups, dataSource, UTC, memoryContext, options, - exception -> handleException(dataSourceId, exception)); + exception -> handleException(dataSourceId, exception), + Optional.empty(), + Optional.empty()); + + ConnectorPageSource pageSource = new ParquetPageSource(parquetReader); + pageSource = transforms.build(pageSource); + + Optional startRowPosition = Optional.empty(); + Optional endRowPosition = Optional.empty(); + if (!rowGroups.isEmpty()) { + startRowPosition = Optional.of(rowGroups.get(0).fileRowOffset()); + RowGroupInfo lastRowGroup = rowGroups.get(rowGroups.size() - 1); + endRowPosition = Optional.of(lastRowGroup.fileRowOffset() + lastRowGroup.prunedBlockMetadata().getRowCount()); + } + return new ReaderPageSourceWithRowPositions( - new ReaderPageSource( - pageSourceBuilder.build(parquetReader), - baseColumnProjections), + pageSource, startRowPosition, endRowPosition); } @@ -1013,8 +1035,8 @@ else if (column.getId() == TRINO_MERGE_PARTITION_DATA) { e.addSuppressed(ex); } } - if (e instanceof TrinoException) { - throw (TrinoException) e; + if (e instanceof TrinoException trinoException) { + throw trinoException; } if (e instanceof ParquetCorruptionException) { throw new TrinoException(ICEBERG_BAD_DATA, e); @@ -1051,10 +1073,7 @@ else if (type instanceof GroupType groupType) { private static MessageType getMessageType(List regularColumns, String fileSchemaName, Map parquetIdToField) { - return projectSufficientColumns(regularColumns) - .map(readerColumns -> readerColumns.get().stream().map(IcebergColumnHandle.class::cast).collect(toUnmodifiableList())) - .orElse(regularColumns) - .stream() + return projectSufficientColumns(regularColumns).stream() .map(column -> getColumnType(column, parquetIdToField)) .filter(Optional::isPresent) .map(Optional::get) @@ -1071,21 +1090,13 @@ private static ReaderPageSourceWithRowPositions createAvroPageSource( String partitionData, Schema fileSchema, Optional nameMapping, + String partition, List columns) { - ConstantPopulatingPageSource.Builder constantPopulatingPageSourceBuilder = ConstantPopulatingPageSource.builder(); - int avroSourceChannel = 0; - - Optional baseColumnProjections = projectBaseColumns(columns); - - List readBaseColumns = baseColumnProjections - .map(readerColumns -> (List) readerColumns.get().stream().map(IcebergColumnHandle.class::cast).collect(toImmutableList())) - .orElse(columns); - InputFile file = new ForwardingInputFile(inputFile); OptionalLong fileModifiedTime = OptionalLong.empty(); try { - if (readBaseColumns.stream().anyMatch(IcebergColumnHandle::isFileModifiedTimeColumn)) { + if (columns.stream().anyMatch(IcebergColumnHandle::isFileModifiedTimeColumn)) { fileModifiedTime = OptionalLong.of(inputFile.lastModified().toEpochMilli()); } } @@ -1107,57 +1118,70 @@ private static ReaderPageSourceWithRowPositions createAvroPageSource( ImmutableList.Builder columnNames = ImmutableList.builder(); ImmutableList.Builder columnTypes = ImmutableList.builder(); - ImmutableList.Builder rowIndexChannels = ImmutableList.builder(); + TransformConnectorPageSource.Builder transforms = TransformConnectorPageSource.builder(); + boolean appendRowNumberColumn = false; + Map baseColumnIdToOrdinal = new HashMap<>(); - for (IcebergColumnHandle column : readBaseColumns) { - verify(column.isBaseColumn(), "Column projections must be based from a root column"); - org.apache.avro.Schema.Field field = fileColumnsByIcebergId.get(column.getId()); - - if (column.isPathColumn()) { - constantPopulatingPageSourceBuilder.addConstantColumn(nativeValueToBlock(FILE_PATH.getType(), utf8Slice(file.location()))); + int nextOrdinal = 0; + for (IcebergColumnHandle column : columns) { + if (column.isPartitionColumn()) { + transforms.constantValue(nativeValueToBlock(PARTITION.getType(), utf8Slice(partition))); } - else if (column.isFileModifiedTimeColumn()) { - constantPopulatingPageSourceBuilder.addConstantColumn(nativeValueToBlock(FILE_MODIFIED_TIME.getType(), packDateTimeWithZone(fileModifiedTime.orElseThrow(), UTC_KEY))); + else if (column.isPathColumn()) { + transforms.constantValue(nativeValueToBlock(FILE_PATH.getType(), utf8Slice(file.location()))); } - // For delete - else if (column.isRowPositionColumn()) { - rowIndexChannels.add(true); - columnNames.add(ROW_POSITION.name()); - columnTypes.add(BIGINT); - constantPopulatingPageSourceBuilder.addDelegateColumn(avroSourceChannel); - avroSourceChannel++; + else if (column.isFileModifiedTimeColumn()) { + transforms.constantValue(nativeValueToBlock(FILE_MODIFIED_TIME.getType(), packDateTimeWithZone(fileModifiedTime.orElseThrow(), UTC_KEY))); } - else if (column.getId() == TRINO_MERGE_PARTITION_SPEC_ID) { - constantPopulatingPageSourceBuilder.addConstantColumn(nativeValueToBlock(column.getType(), (long) partitionSpecId)); + else if (column.isMergeRowIdColumn()) { + appendRowNumberColumn = true; + transforms.transform(MergeRowIdTransform.create(utf8Slice(file.location()), partitionSpecId, utf8Slice(partitionData))); } - else if (column.getId() == TRINO_MERGE_PARTITION_DATA) { - constantPopulatingPageSourceBuilder.addConstantColumn(nativeValueToBlock(column.getType(), utf8Slice(partitionData))); + else if (column.isRowPositionColumn()) { + appendRowNumberColumn = true; + transforms.transform(new GetRowPositionFromSource()); } - else if (field == null) { - constantPopulatingPageSourceBuilder.addConstantColumn(nativeValueToBlock(column.getType(), null)); + else if (!fileColumnsByIcebergId.containsKey(column.getBaseColumn().getId())) { + transforms.constantValue(nativeValueToBlock(column.getType(), null)); } else { - rowIndexChannels.add(false); - columnNames.add(column.getName()); - columnTypes.add(column.getType()); - constantPopulatingPageSourceBuilder.addDelegateColumn(avroSourceChannel); - avroSourceChannel++; + IcebergColumnHandle baseColumn = column.getBaseColumn(); + Integer ordinal = baseColumnIdToOrdinal.get(baseColumn.getId()); + if (ordinal == null) { + ordinal = nextOrdinal; + nextOrdinal++; + baseColumnIdToOrdinal.put(baseColumn.getId(), ordinal); + + columnNames.add(baseColumn.getName()); + columnTypes.add(baseColumn.getType()); + } + + if (column.isBaseColumn()) { + transforms.column(ordinal); + } + else { + transforms.dereferenceField(ImmutableList.builder() + .add(ordinal) + .addAll(applyProjection(column, baseColumn)) + .build()); + } } } + ConnectorPageSource pageSource = new IcebergAvroPageSource( + file, + start, + length, + fileSchema, + nameMapping, + columnNames.build(), + columnTypes.build(), + appendRowNumberColumn, + newSimpleAggregatedMemoryContext()); + pageSource = transforms.build(pageSource); + return new ReaderPageSourceWithRowPositions( - new ReaderPageSource( - constantPopulatingPageSourceBuilder.build(new IcebergAvroPageSource( - file, - start, - length, - fileSchema, - nameMapping, - columnNames.build(), - columnTypes.build(), - rowIndexChannels.build(), - newSimpleAggregatedMemoryContext())), - baseColumnProjections), + pageSource, Optional.empty(), Optional.empty()); } @@ -1262,52 +1286,17 @@ public ProjectedLayout getFieldLayout(OrcColumn orcColumn) } } - /** - * Creates a mapping between the input {@code columns} and base columns if required. - */ - public static Optional projectBaseColumns(List columns) - { - requireNonNull(columns, "columns is null"); - - // No projection is required if all columns are base columns - if (columns.stream().allMatch(IcebergColumnHandle::isBaseColumn)) { - return Optional.empty(); - } - - ImmutableList.Builder projectedColumns = ImmutableList.builder(); - ImmutableList.Builder outputColumnMapping = ImmutableList.builder(); - Map mappedFieldIds = new HashMap<>(); - int projectedColumnCount = 0; - - for (IcebergColumnHandle column : columns) { - int baseColumnId = column.getBaseColumnIdentity().getId(); - Integer mapped = mappedFieldIds.get(baseColumnId); - - if (mapped == null) { - projectedColumns.add(column.getBaseColumn()); - mappedFieldIds.put(baseColumnId, projectedColumnCount); - outputColumnMapping.add(projectedColumnCount); - projectedColumnCount++; - } - else { - outputColumnMapping.add(mapped); - } - } - - return Optional.of(new ReaderColumns(projectedColumns.build(), outputColumnMapping.build())); - } - /** * Creates a set of sufficient columns for the input projected columns and prepares a mapping between the two. - * For example, if input {@param columns} include columns "a.b" and "a.b.c", then they will be projected + * For example, if input columns include columns "a.b" and "a.b.c", then they will be projected * from a single column "a.b". */ - private static Optional projectSufficientColumns(List columns) + private static List projectSufficientColumns(List columns) { requireNonNull(columns, "columns is null"); if (columns.stream().allMatch(IcebergColumnHandle::isBaseColumn)) { - return Optional.empty(); + return columns; } ImmutableBiMap.Builder dereferenceChainsBuilder = ImmutableBiMap.builder(); @@ -1319,14 +1308,13 @@ private static Optional projectSufficientColumns(List dereferenceChains = dereferenceChainsBuilder.build(); - List sufficientColumns = new ArrayList<>(); - ImmutableList.Builder outputColumnMapping = ImmutableList.builder(); + List sufficientColumns = new ArrayList<>(); Map pickedColumns = new HashMap<>(); // Pick a covering column for every column for (IcebergColumnHandle columnHandle : columns) { - DereferenceChain dereferenceChain = dereferenceChains.inverse().get(columnHandle); + DereferenceChain dereferenceChain = requireNonNull(dereferenceChains.inverse().get(columnHandle)); DereferenceChain chosenColumn = null; // Shortest existing prefix is chosen as the input. @@ -1338,23 +1326,15 @@ private static Optional projectSufficientColumns(List getColumnType(IcebergColumnHandle column, Map parquetIdToField) @@ -1390,12 +1370,15 @@ private static TupleDomain getParquetTupleDomain(Map descriptorsById = descriptorsByPath.values().stream() + .filter(descriptor -> descriptor.getPrimitiveType().getId() != null) + .collect(toImmutableMap(descriptor -> descriptor.getPrimitiveType().getId().intValue(), identity())); ImmutableMap.Builder predicate = ImmutableMap.builder(); effectivePredicate.getDomains().orElseThrow().forEach((columnHandle, domain) -> { - String baseType = columnHandle.getType().getTypeSignature().getBase(); + ColumnIdentity columnIdentity = columnHandle.getColumnIdentity(); // skip looking up predicates for complex types as Parquet only stores stats for primitives - if (columnHandle.isBaseColumn() && (!baseType.equals(StandardTypes.MAP) && !baseType.equals(StandardTypes.ARRAY) && !baseType.equals(StandardTypes.ROW))) { - ColumnDescriptor descriptor = descriptorsByPath.get(ImmutableList.of(columnHandle.getName())); + if (PRIMITIVE == columnIdentity.getTypeCategory()) { + ColumnDescriptor descriptor = descriptorsById.get(columnHandle.getId()); if (descriptor != null) { predicate.put(descriptor, domain); } @@ -1417,8 +1400,8 @@ private static TrinoException handleException(OrcDataSourceId dataSourceId, Exce private static TrinoException handleException(ParquetDataSourceId dataSourceId, Exception exception) { - if (exception instanceof TrinoException) { - return (TrinoException) exception; + if (exception instanceof TrinoException trinoException) { + return trinoException; } if (exception instanceof ParquetCorruptionException) { return new TrinoException(ICEBERG_BAD_DATA, exception); @@ -1426,35 +1409,13 @@ private static TrinoException handleException(ParquetDataSourceId dataSourceId, return new TrinoException(ICEBERG_CURSOR_ERROR, format("Failed to read Parquet file: %s", dataSourceId), exception); } - public static final class ReaderPageSourceWithRowPositions + public record ReaderPageSourceWithRowPositions(ConnectorPageSource pageSource, Optional startRowPosition, Optional endRowPosition) { - private final ReaderPageSource readerPageSource; - private final Optional startRowPosition; - private final Optional endRowPosition; - - public ReaderPageSourceWithRowPositions( - ReaderPageSource readerPageSource, - Optional startRowPosition, - Optional endRowPosition) - { - this.readerPageSource = requireNonNull(readerPageSource, "readerPageSource is null"); - this.startRowPosition = requireNonNull(startRowPosition, "startRowPosition is null"); - this.endRowPosition = requireNonNull(endRowPosition, "endRowPosition is null"); - } - - public ReaderPageSource getReaderPageSource() - { - return readerPageSource; - } - - public Optional getStartRowPosition() + public ReaderPageSourceWithRowPositions { - return startRowPosition; - } - - public Optional getEndRowPosition() - { - return endRowPosition; + requireNonNull(pageSource, "pageSource is null"); + requireNonNull(startRowPosition, "startRowPosition is null"); + requireNonNull(endRowPosition, "endRowPosition is null"); } } @@ -1510,4 +1471,70 @@ public int hashCode() return Objects.hash(baseColumnIdentity, path); } } + + private record MergeRowIdTransform(VariableWidthBlock filePath, IntArrayBlock partitionSpecId, VariableWidthBlock partitionData) + implements Function + { + private static Function create(Slice filePath, int partitionSpecId, Slice partitionData) + { + return new MergeRowIdTransform( + new VariableWidthBlock(1, filePath, new int[] {0, filePath.length()}, Optional.empty()), + new IntArrayBlock(1, Optional.empty(), new int[] {partitionSpecId}), + new VariableWidthBlock(1, partitionData, new int[] {0, partitionData.length()}, Optional.empty())); + } + + @Override + public Block apply(Page page) + { + Block rowPosition = page.getBlock(page.getChannelCount() - 1); + Block[] fields = new Block[] { + RunLengthEncodedBlock.create(filePath, rowPosition.getPositionCount()), + rowPosition, + RunLengthEncodedBlock.create(partitionSpecId, rowPosition.getPositionCount()), + RunLengthEncodedBlock.create(partitionData, rowPosition.getPositionCount()) + }; + return RowBlock.fromFieldBlocks(rowPosition.getPositionCount(), fields); + } + } + + private record GetRowPositionFromSource() + implements Function + { + @Override + public Block apply(Page page) + { + return page.getBlock(page.getChannelCount() - 1); + } + } + + private record PrefixColumnsPage(Page page, int channelCount, int[] channels) + { + private PrefixColumnsPage + { + requireNonNull(page, "page is null"); + checkArgument(channelCount >= 0, "channelCount is negative"); + checkArgument(channelCount < page.getChannelCount(), "channelCount is greater than or equal to Page channel count"); + checkArgument(channels.length == channelCount, "channels length does not match channelCount"); + } + + private PrefixColumnsPage(Page page, int channelCount) + { + this(page, channelCount, IntStream.range(0, channelCount).toArray()); + } + + public Block getBlock(int channel) + { + checkIndex(channel, channelCount); + return page.getBlock(channel); + } + + public Page getPage() + { + Block[] blocks = new Block[channels.length]; + for (int i = 0; i < channels.length; i++) { + blocks[i] = getBlock(channels[i]); + } + return new Page(page.getPositionCount(), blocks); + } + } } diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergPageSourceProviderFactory.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergPageSourceProviderFactory.java new file mode 100644 index 000000000000..41a6444e4546 --- /dev/null +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergPageSourceProviderFactory.java @@ -0,0 +1,70 @@ +/* + * 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.iceberg; + +import com.google.inject.Inject; +import io.trino.orc.OrcReaderOptions; +import io.trino.parquet.ParquetReaderOptions; +import io.trino.plugin.hive.FileFormatDataSourceStats; +import io.trino.plugin.hive.orc.OrcReaderConfig; +import io.trino.plugin.hive.parquet.ParquetReaderConfig; +import io.trino.spi.connector.ColumnHandle; +import io.trino.spi.connector.ConnectorPageSource; +import io.trino.spi.connector.ConnectorPageSourceProvider; +import io.trino.spi.connector.ConnectorSession; +import io.trino.spi.connector.ConnectorSplit; +import io.trino.spi.connector.ConnectorTableHandle; +import io.trino.spi.connector.ConnectorTransactionHandle; +import io.trino.spi.connector.DynamicFilter; +import io.trino.spi.type.TypeManager; + +import java.util.List; + +import static java.util.Objects.requireNonNull; + +public class IcebergPageSourceProviderFactory + implements ConnectorPageSourceProvider +{ + private final IcebergFileSystemFactory fileSystemFactory; + private final FileFormatDataSourceStats fileFormatDataSourceStats; + private final OrcReaderOptions orcReaderOptions; + private final ParquetReaderOptions parquetReaderOptions; + private final TypeManager typeManager; + + @Inject + public IcebergPageSourceProviderFactory( + IcebergFileSystemFactory fileSystemFactory, + FileFormatDataSourceStats fileFormatDataSourceStats, + OrcReaderConfig orcReaderConfig, + ParquetReaderConfig parquetReaderConfig, + TypeManager typeManager) + { + this.fileSystemFactory = requireNonNull(fileSystemFactory, "fileSystemFactory is null"); + this.fileFormatDataSourceStats = requireNonNull(fileFormatDataSourceStats, "fileFormatDataSourceStats is null"); + this.orcReaderOptions = orcReaderConfig.toOrcReaderOptions(); + this.parquetReaderOptions = parquetReaderConfig.toParquetReaderOptions(); + this.typeManager = requireNonNull(typeManager, "typeManager is null"); + } + + public ConnectorPageSourceProvider createPageSourceProvider() + { + return new IcebergPageSourceProvider(fileSystemFactory, fileFormatDataSourceStats, orcReaderOptions, parquetReaderOptions, typeManager); + } + + @Override + public ConnectorPageSource createPageSource(ConnectorTransactionHandle transaction, ConnectorSession session, ConnectorSplit split, ConnectorTableHandle table, List columns, DynamicFilter dynamicFilter) + { + return createPageSourceProvider().createPageSource(transaction, session, split, table, columns, dynamicFilter); + } +} diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergParquetFileWriter.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergParquetFileWriter.java index 104598c8dd0a..06cccccebe3c 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergParquetFileWriter.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergParquetFileWriter.java @@ -13,15 +13,16 @@ */ package io.trino.plugin.iceberg; -import io.trino.filesystem.TrinoFileSystem; +import io.trino.filesystem.Location; import io.trino.filesystem.TrinoOutputFile; +import io.trino.parquet.ParquetDataSourceId; +import io.trino.parquet.metadata.ParquetMetadata; import io.trino.parquet.writer.ParquetWriterOptions; import io.trino.plugin.hive.parquet.ParquetFileWriter; -import io.trino.plugin.iceberg.fileio.ForwardingInputFile; +import io.trino.spi.Page; +import io.trino.spi.TrinoException; import io.trino.spi.type.Type; -import org.apache.iceberg.Metrics; import org.apache.iceberg.MetricsConfig; -import org.apache.iceberg.io.InputFile; import org.apache.parquet.format.CompressionCodec; import org.apache.parquet.schema.MessageType; @@ -30,16 +31,20 @@ import java.util.List; import java.util.Map; import java.util.Optional; +import java.util.stream.Stream; +import static io.trino.plugin.iceberg.util.ParquetUtil.footerMetrics; +import static io.trino.plugin.iceberg.util.ParquetUtil.getSplitOffsets; +import static io.trino.spi.StandardErrorCode.GENERIC_INTERNAL_ERROR; +import static java.lang.String.format; import static java.util.Objects.requireNonNull; -import static org.apache.iceberg.parquet.ParquetUtil.fileMetrics; -public class IcebergParquetFileWriter - extends ParquetFileWriter +public final class IcebergParquetFileWriter implements IcebergFileWriter { private final MetricsConfig metricsConfig; - private final InputFile inputFile; + private final ParquetFileWriter parquetFileWriter; + private final Location location; public IcebergParquetFileWriter( MetricsConfig metricsConfig, @@ -52,11 +57,11 @@ public IcebergParquetFileWriter( ParquetWriterOptions parquetWriterOptions, int[] fileInputColumnIndexes, CompressionCodec compressionCodec, - String trinoVersion, - TrinoFileSystem fileSystem) + String trinoVersion) throws IOException { - super(outputFile, + this.parquetFileWriter = new ParquetFileWriter( + outputFile, rollbackAction, fileColumnTypes, fileColumnNames, @@ -66,16 +71,58 @@ public IcebergParquetFileWriter( fileInputColumnIndexes, compressionCodec, trinoVersion, - false, Optional.empty(), Optional.empty()); + this.location = outputFile.location(); this.metricsConfig = requireNonNull(metricsConfig, "metricsConfig is null"); - this.inputFile = new ForwardingInputFile(fileSystem.newInputFile(outputFile.location())); } @Override - public Metrics getMetrics() + public FileMetrics getFileMetrics() { - return fileMetrics(inputFile, metricsConfig); + ParquetMetadata parquetMetadata; + try { + parquetMetadata = new ParquetMetadata(parquetFileWriter.getFileMetadata(), new ParquetDataSourceId(location.toString())); + return new FileMetrics(footerMetrics(parquetMetadata, Stream.empty(), metricsConfig), Optional.of(getSplitOffsets(parquetMetadata))); + } + catch (IOException e) { + throw new TrinoException(GENERIC_INTERNAL_ERROR, format("Error creating metadata for Parquet file %s", location), e); + } + } + + @Override + public long getWrittenBytes() + { + return parquetFileWriter.getWrittenBytes(); + } + + @Override + public long getMemoryUsage() + { + return parquetFileWriter.getMemoryUsage(); + } + + @Override + public void appendRows(Page dataPage) + { + parquetFileWriter.appendRows(dataPage); + } + + @Override + public Closeable commit() + { + return parquetFileWriter.commit(); + } + + @Override + public void rollback() + { + parquetFileWriter.rollback(); + } + + @Override + public long getValidationCpuNanos() + { + return parquetFileWriter.getValidationCpuNanos(); } } diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergPartitionColumn.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergPartitionColumn.java new file mode 100644 index 000000000000..804fcaa6311c --- /dev/null +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergPartitionColumn.java @@ -0,0 +1,30 @@ +/* + * 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.iceberg; + +import com.google.common.collect.ImmutableList; +import io.trino.spi.type.RowType; + +import java.util.List; + +import static java.util.Objects.requireNonNull; + +public record IcebergPartitionColumn(RowType rowType, List fieldIds) +{ + public IcebergPartitionColumn + { + requireNonNull(rowType, "rowType is null"); + fieldIds = ImmutableList.copyOf(requireNonNull(fieldIds, "fieldIds is null")); + } +} diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergPartitionFunction.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergPartitionFunction.java new file mode 100644 index 000000000000..a558c2363546 --- /dev/null +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergPartitionFunction.java @@ -0,0 +1,98 @@ +/* + * 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.iceberg; + +import com.google.common.collect.ImmutableList; +import io.trino.spi.type.Type; + +import java.util.List; +import java.util.OptionalInt; +import java.util.regex.Matcher; + +import static com.google.common.base.Preconditions.checkArgument; +import static io.trino.plugin.iceberg.PartitionFields.ICEBERG_BUCKET_PATTERN; +import static io.trino.plugin.iceberg.PartitionFields.ICEBERG_TRUNCATE_PATTERN; +import static java.util.Objects.requireNonNull; + +// NOTE: the partitioning function must contain the data path for nested fields because the partitioning columns +// reference the top level column, not the nested column. This means that even thought the partitioning functions +// with a different path should be compatible, the system does not consider them compatible. Fortunately, partitioning +// on nested columns is not common. +public record IcebergPartitionFunction(Transform transform, List dataPath, Type type, OptionalInt size) +{ + public enum Transform + { + IDENTITY, + YEAR, + MONTH, + DAY, + HOUR, + VOID, + BUCKET, + TRUNCATE + } + + public IcebergPartitionFunction(Transform transform, List dataPath, Type type) + { + this(transform, dataPath, type, OptionalInt.empty()); + } + + public IcebergPartitionFunction + { + requireNonNull(transform, "transform is null"); + requireNonNull(dataPath, "dataPath is null"); + checkArgument(!dataPath.isEmpty(), "dataPath is empty"); + requireNonNull(type, "type is null"); + requireNonNull(size, "size is null"); + checkArgument(size.orElse(0) >= 0, "size must be greater than or equal to zero"); + checkArgument(size.isEmpty() || transform == Transform.BUCKET || transform == Transform.TRUNCATE, "size is only valid for BUCKET and TRUNCATE transforms"); + } + + public IcebergPartitionFunction withTopLevelColumnIndex(int newColumnIndex) + { + return new IcebergPartitionFunction( + transform, + ImmutableList.builder() + .add(newColumnIndex) + .addAll(dataPath().subList(1, dataPath().size())) + .build(), + type, + size); + } + + public static IcebergPartitionFunction create(String transform, List dataPath, Type type) + { + return switch (transform) { + case "identity" -> new IcebergPartitionFunction(Transform.IDENTITY, dataPath, type); + case "year" -> new IcebergPartitionFunction(Transform.YEAR, dataPath, type); + case "month" -> new IcebergPartitionFunction(Transform.MONTH, dataPath, type); + case "day" -> new IcebergPartitionFunction(Transform.DAY, dataPath, type); + case "hour" -> new IcebergPartitionFunction(Transform.HOUR, dataPath, type); + case "void" -> new IcebergPartitionFunction(Transform.VOID, dataPath, type); + default -> { + Matcher matcher = ICEBERG_BUCKET_PATTERN.matcher(transform); + if (matcher.matches()) { + yield new IcebergPartitionFunction(Transform.BUCKET, dataPath, type, OptionalInt.of(Integer.parseInt(matcher.group(1)))); + } + + matcher = ICEBERG_TRUNCATE_PATTERN.matcher(transform); + if (matcher.matches()) { + yield new IcebergPartitionFunction(Transform.TRUNCATE, dataPath, type, OptionalInt.of(Integer.parseInt(matcher.group(1)))); + } + + throw new UnsupportedOperationException("Unsupported partition transform: " + transform); + } + }; + } +} diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergPartitioningHandle.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergPartitioningHandle.java index 884ead60950e..a42d629db219 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergPartitioningHandle.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergPartitioningHandle.java @@ -13,48 +13,141 @@ */ package io.trino.plugin.iceberg; -import com.fasterxml.jackson.annotation.JsonCreator; -import com.fasterxml.jackson.annotation.JsonProperty; import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.common.hash.Hasher; +import com.google.common.hash.Hashing; import io.trino.spi.connector.ConnectorPartitioningHandle; +import io.trino.spi.type.TypeManager; +import org.apache.iceberg.PartitionField; +import org.apache.iceberg.PartitionSpec; +import org.apache.iceberg.types.Types; +import java.util.ArrayDeque; +import java.util.HashMap; import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.concurrent.atomic.AtomicInteger; -import static com.google.common.base.MoreObjects.toStringHelper; +import static com.google.common.collect.ImmutableList.toImmutableList; +import static com.google.common.collect.ImmutableSet.toImmutableSet; +import static io.trino.plugin.iceberg.TypeConverter.toTrinoType; +import static java.nio.charset.StandardCharsets.UTF_8; import static java.util.Objects.requireNonNull; -public class IcebergPartitioningHandle +public record IcebergPartitioningHandle(boolean update, List partitionFunctions) implements ConnectorPartitioningHandle { - private final List partitioning; - private final List partitioningColumns; + public IcebergPartitioningHandle + { + partitionFunctions = ImmutableList.copyOf(requireNonNull(partitionFunctions, "partitioning is null")); + } + + public IcebergPartitioningHandle forUpdate() + { + return new IcebergPartitioningHandle(true, partitionFunctions); + } - @JsonCreator - public IcebergPartitioningHandle( - @JsonProperty("partitioning") List partitioning, - @JsonProperty("partitioningColumns") List partitioningColumns) + public static IcebergPartitioningHandle create(PartitionSpec spec, TypeManager typeManager, List partitioningColumns) { - this.partitioning = ImmutableList.copyOf(requireNonNull(partitioning, "partitioning is null")); - this.partitioningColumns = ImmutableList.copyOf(requireNonNull(partitioningColumns, "partitioningColumns is null")); + Map> dataPaths = buildDataPaths(spec); + List partitionFields = spec.fields().stream() + .map(field -> IcebergPartitionFunction.create( + field.transform().toString(), + dataPaths.get(field.sourceId()), + toTrinoType(spec.schema().findType(field.sourceId()), typeManager))) + .collect(toImmutableList()); + + return new IcebergPartitioningHandle(false, partitionFields); } - @JsonProperty - public List getPartitioning() + /** + * Constructs a map of field IDs to data paths. + * The data path for root field is the ordinal position of the partition field under this root field, defined by {@link IcebergMetadata#getWriteLayout} + * The data path for non-root nested fields is the ordinal position in its parent's nested field. + * e.g. for a schema {f1: {f3, f4}, f2, f5} + * when partitioned by f1.f3 and f2, the data paths are {3 : [1,0], 2 : [0]} + * when partitioned by f1.f4 and f5, the data paths are {4 : [0, 1], 5 : [1]} + */ + private static Map> buildDataPaths(PartitionSpec spec) { - return partitioning; + Set partitionFieldIds = spec.fields().stream().map(PartitionField::sourceId).collect(toImmutableSet()); + + /* + * In this loop, the field ID acts as a placeholder in the first position + * Later, these placeholders will be replaced with the actual channel IDs by the order of its partitioned sub-field ID. + */ + Map> fieldInfo = new HashMap<>(); + for (Types.NestedField field : spec.schema().asStruct().fields()) { + // Partition fields can only be nested in a struct + if (field.type() instanceof Types.StructType nestedStruct) { + buildDataPaths(partitionFieldIds, nestedStruct, new ArrayDeque<>(ImmutableList.of(field.fieldId())), fieldInfo); + } + else if (field.type().isPrimitiveType() && partitionFieldIds.contains(field.fieldId())) { + fieldInfo.put(field.fieldId(), ImmutableList.of(field.fieldId())); + } + } + + /* + * Replace the root field ID with the actual channel ID. + * Transformation: {fieldId : rootFieldId.structOrdinalX.structOrdinalY} -> {fieldId : channel.structOrdinalX.structOrdinalY}. + * Root field's channelId is assigned sequentially based on the key fieldId. + */ + List sortedFieldIds = fieldInfo.keySet().stream() + .sorted() + .collect(toImmutableList()); + + ImmutableMap.Builder> builder = ImmutableMap + .builderWithExpectedSize(sortedFieldIds.size()); + + Map fieldChannels = new HashMap<>(); + AtomicInteger channel = new AtomicInteger(); + for (int sortedFieldId : sortedFieldIds) { + List dataPath = fieldInfo.get(sortedFieldId); + int fieldChannel = fieldChannels.computeIfAbsent(dataPath.get(0), ignore -> channel.getAndIncrement()); + List channelDataPath = ImmutableList.builder() + .add(fieldChannel) + .addAll(dataPath.stream() + .skip(1) + .iterator()) + .build(); + builder.put(sortedFieldId, channelDataPath); + } + + return builder.buildOrThrow(); } - @JsonProperty - public List getPartitioningColumns() + private static void buildDataPaths(Set partitionFieldIds, Types.StructType struct, ArrayDeque currentPaths, Map> dataPaths) { - return partitioningColumns; + List fields = struct.fields(); + for (int fieldOrdinal = 0; fieldOrdinal < fields.size(); fieldOrdinal++) { + Types.NestedField field = fields.get(fieldOrdinal); + int fieldId = field.fieldId(); + + currentPaths.addLast(fieldOrdinal); + org.apache.iceberg.types.Type type = field.type(); + if (type instanceof Types.StructType nestedStruct) { + buildDataPaths(partitionFieldIds, nestedStruct, currentPaths, dataPaths); + } + // Map and List types are not supported in partitioning + else if (type.isPrimitiveType() && partitionFieldIds.contains(fieldId)) { + dataPaths.put(fieldId, ImmutableList.copyOf(currentPaths)); + } + currentPaths.removeLast(); + } } - @Override - public String toString() + public long getCacheKeyHint() { - return toStringHelper(this) - .add("partitioning", partitioning) - .toString(); + Hasher hasher = Hashing.goodFastHash(64).newHasher(); + hasher.putBoolean(update); + for (IcebergPartitionFunction function : partitionFunctions) { + hasher.putInt(function.transform().ordinal()); + function.dataPath().forEach(hasher::putInt); + hasher.putString(function.type().getTypeSignature().toString(), UTF_8); + function.size().ifPresent(hasher::putInt); + } + return hasher.hash().asLong(); } } diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergSchemaProperties.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergSchemaProperties.java index a87c48a5eaf3..bdc43ab1b9e3 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergSchemaProperties.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergSchemaProperties.java @@ -14,11 +14,10 @@ package io.trino.plugin.iceberg; import com.google.common.collect.ImmutableList; +import com.google.inject.Inject; import io.trino.spi.session.PropertyMetadata; import java.util.List; -import java.util.Map; -import java.util.Optional; import static io.trino.spi.session.PropertyMetadata.stringProperty; @@ -26,18 +25,22 @@ public final class IcebergSchemaProperties { public static final String LOCATION_PROPERTY = "location"; - public static final List> SCHEMA_PROPERTIES = ImmutableList.>builder() - .add(stringProperty( - LOCATION_PROPERTY, - "Base file system location URI", - null, - false)) - .build(); + public final List> schemaProperties; - private IcebergSchemaProperties() {} + @Inject + public IcebergSchemaProperties() + { + this.schemaProperties = ImmutableList.>builder() + .add(stringProperty( + LOCATION_PROPERTY, + "Base file system location URI", + null, + false)) + .build(); + } - public static Optional getSchemaLocation(Map schemaProperties) + public List> getSchemaProperties() { - return Optional.ofNullable((String) schemaProperties.get(LOCATION_PROPERTY)); + return schemaProperties; } } diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergSecurityModule.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergSecurityModule.java index 6aa9f419b388..d3c9624059dd 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergSecurityModule.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergSecurityModule.java @@ -19,9 +19,11 @@ import io.trino.plugin.base.security.ConnectorAccessControlModule; import io.trino.plugin.base.security.FileBasedAccessControlModule; import io.trino.plugin.base.security.ReadOnlySecurityModule; +import io.trino.plugin.hive.security.UsingSystemSecurity; import io.trino.plugin.iceberg.IcebergSecurityConfig.IcebergSecurity; import static io.airlift.configuration.ConditionalModule.conditionalModule; +import static io.airlift.configuration.ConfigurationAwareModule.combine; import static io.trino.plugin.iceberg.IcebergSecurityConfig.IcebergSecurity.ALLOW_ALL; import static io.trino.plugin.iceberg.IcebergSecurityConfig.IcebergSecurity.FILE; import static io.trino.plugin.iceberg.IcebergSecurityConfig.IcebergSecurity.READ_ONLY; @@ -33,12 +35,17 @@ public class IcebergSecurityModule protected void setup(Binder binder) { install(new ConnectorAccessControlModule()); - bindSecurityModule(ALLOW_ALL, new AllowAllSecurityModule()); - bindSecurityModule(READ_ONLY, new ReadOnlySecurityModule()); - bindSecurityModule(FILE, new FileBasedAccessControlModule()); + bindSecurityModule(ALLOW_ALL, combine(new AllowAllSecurityModule(), usingSystemSecurity(false))); + bindSecurityModule(READ_ONLY, combine(new ReadOnlySecurityModule(), usingSystemSecurity(false))); + bindSecurityModule(FILE, combine(new FileBasedAccessControlModule(), usingSystemSecurity(false))); // SYSTEM: do not bind an ConnectorAccessControl so the engine will use system security with system roles } + private static Module usingSystemSecurity(boolean system) + { + return binder -> binder.bind(boolean.class).annotatedWith(UsingSystemSecurity.class).toInstance(system); + } + private void bindSecurityModule(IcebergSecurity icebergSecurity, Module module) { install(conditionalModule( diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergSessionProperties.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergSessionProperties.java index a607c52b37cd..87c8e06536ae 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergSessionProperties.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergSessionProperties.java @@ -19,7 +19,7 @@ import io.airlift.units.Duration; import io.trino.orc.OrcWriteValidation.OrcWriteValidationMode; import io.trino.plugin.base.session.SessionPropertiesProvider; -import io.trino.plugin.hive.HiveCompressionCodec; +import io.trino.plugin.hive.HiveCompressionOption; import io.trino.plugin.hive.orc.OrcReaderConfig; import io.trino.plugin.hive.orc.OrcWriterConfig; import io.trino.plugin.hive.parquet.ParquetReaderConfig; @@ -27,19 +27,27 @@ import io.trino.spi.TrinoException; import io.trino.spi.connector.ConnectorSession; import io.trino.spi.session.PropertyMetadata; +import io.trino.spi.type.ArrayType; +import java.util.Collection; import java.util.List; import java.util.Optional; +import java.util.Set; import java.util.concurrent.ThreadLocalRandom; import static com.google.common.base.Preconditions.checkArgument; +import static com.google.common.base.Strings.isNullOrEmpty; +import static com.google.common.collect.ImmutableSet.toImmutableSet; import static io.trino.plugin.base.session.PropertyMetadataUtil.dataSizeProperty; import static io.trino.plugin.base.session.PropertyMetadataUtil.durationProperty; import static io.trino.plugin.base.session.PropertyMetadataUtil.validateMaxDataSize; import static io.trino.plugin.base.session.PropertyMetadataUtil.validateMinDataSize; +import static io.trino.plugin.hive.parquet.ParquetReaderConfig.PARQUET_READER_MAX_SMALL_FILE_THRESHOLD; import static io.trino.plugin.hive.parquet.ParquetWriterConfig.PARQUET_WRITER_MAX_BLOCK_SIZE; import static io.trino.plugin.hive.parquet.ParquetWriterConfig.PARQUET_WRITER_MAX_PAGE_SIZE; +import static io.trino.plugin.hive.parquet.ParquetWriterConfig.PARQUET_WRITER_MAX_PAGE_VALUE_COUNT; import static io.trino.plugin.hive.parquet.ParquetWriterConfig.PARQUET_WRITER_MIN_PAGE_SIZE; +import static io.trino.plugin.hive.parquet.ParquetWriterConfig.PARQUET_WRITER_MIN_PAGE_VALUE_COUNT; import static io.trino.plugin.iceberg.IcebergConfig.COLLECT_EXTENDED_STATISTICS_ON_WRITE_DESCRIPTION; import static io.trino.plugin.iceberg.IcebergConfig.EXTENDED_STATISTICS_DESCRIPTION; import static io.trino.spi.StandardErrorCode.INVALID_SESSION_PROPERTY; @@ -48,11 +56,15 @@ import static io.trino.spi.session.PropertyMetadata.enumProperty; import static io.trino.spi.session.PropertyMetadata.integerProperty; import static io.trino.spi.session.PropertyMetadata.stringProperty; +import static io.trino.spi.type.VarcharType.VARCHAR; import static java.lang.String.format; +import static java.util.Locale.ENGLISH; +import static java.util.Objects.requireNonNull; public final class IcebergSessionProperties implements SessionPropertiesProvider { + public static final String SPLIT_SIZE = "experimental_split_size"; private static final String COMPRESSION_CODEC = "compression_codec"; private static final String USE_FILE_SIZE_FROM_METADATA = "use_file_size_from_metadata"; private static final String ORC_BLOOM_FILTERS_ENABLED = "orc_bloom_filters_enabled"; @@ -69,20 +81,25 @@ public final class IcebergSessionProperties private static final String ORC_WRITER_MIN_STRIPE_SIZE = "orc_writer_min_stripe_size"; private static final String ORC_WRITER_MAX_STRIPE_SIZE = "orc_writer_max_stripe_size"; private static final String ORC_WRITER_MAX_STRIPE_ROWS = "orc_writer_max_stripe_rows"; + private static final String ORC_WRITER_MAX_ROW_GROUP_ROWS = "orc_writer_max_row_group_rows"; private static final String ORC_WRITER_MAX_DICTIONARY_MEMORY = "orc_writer_max_dictionary_memory"; private static final String PARQUET_MAX_READ_BLOCK_SIZE = "parquet_max_read_block_size"; private static final String PARQUET_USE_BLOOM_FILTER = "parquet_use_bloom_filter"; private static final String PARQUET_MAX_READ_BLOCK_ROW_COUNT = "parquet_max_read_block_row_count"; private static final String PARQUET_OPTIMIZED_READER_ENABLED = "parquet_optimized_reader_enabled"; private static final String PARQUET_OPTIMIZED_NESTED_READER_ENABLED = "parquet_optimized_nested_reader_enabled"; + private static final String PARQUET_SMALL_FILE_THRESHOLD = "parquet_small_file_threshold"; + private static final String PARQUET_IGNORE_STATISTICS = "parquet_ignore_statistics"; private static final String PARQUET_WRITER_BLOCK_SIZE = "parquet_writer_block_size"; private static final String PARQUET_WRITER_PAGE_SIZE = "parquet_writer_page_size"; + private static final String PARQUET_WRITER_PAGE_VALUE_COUNT = "parquet_writer_page_value_count"; private static final String PARQUET_WRITER_BATCH_SIZE = "parquet_writer_batch_size"; private static final String DYNAMIC_FILTERING_WAIT_TIMEOUT = "dynamic_filtering_wait_timeout"; private static final String STATISTICS_ENABLED = "statistics_enabled"; public static final String EXTENDED_STATISTICS_ENABLED = "extended_statistics_enabled"; private static final String PROJECTION_PUSHDOWN_ENABLED = "projection_pushdown_enabled"; private static final String TARGET_MAX_FILE_SIZE = "target_max_file_size"; + private static final String IDLE_WRITER_MIN_FILE_SIZE = "idle_writer_min_file_size"; public static final String COLLECT_EXTENDED_STATISTICS_ON_WRITE = "collect_extended_statistics_on_write"; private static final String HIVE_CATALOG_NAME = "hive_catalog_name"; private static final String MINIMUM_ASSIGNED_SPLIT_WEIGHT = "minimum_assigned_split_weight"; @@ -90,6 +107,11 @@ public final class IcebergSessionProperties public static final String REMOVE_ORPHAN_FILES_MIN_RETENTION = "remove_orphan_files_min_retention"; private static final String MERGE_MANIFESTS_ON_WRITE = "merge_manifests_on_write"; private static final String SORTED_WRITING_ENABLED = "sorted_writing_enabled"; + private static final String QUERY_PARTITION_FILTER_REQUIRED = "query_partition_filter_required"; + private static final String QUERY_PARTITION_FILTER_REQUIRED_SCHEMAS = "query_partition_filter_required_schemas"; + private static final String INCREMENTAL_REFRESH_ENABLED = "incremental_refresh_enabled"; + public static final String BUCKET_EXECUTION_ENABLED = "bucket_execution_enabled"; + public static final String FILE_BASED_CONFLICT_DETECTION_ENABLED = "file_based_conflict_detection_enabled"; private final List> sessionProperties; @@ -102,10 +124,17 @@ public IcebergSessionProperties( ParquetWriterConfig parquetWriterConfig) { sessionProperties = ImmutableList.>builder() + .add(dataSizeProperty( + SPLIT_SIZE, + "Target split size", + // Note: this is null by default & hidden, currently mainly for tests. + // See https://github.com/trinodb/trino/issues/9018#issuecomment-1752929193 for further discussion. + null, + true)) .add(enumProperty( COMPRESSION_CODEC, "Compression codec to use when writing files", - HiveCompressionCodec.class, + HiveCompressionOption.class, icebergConfig.getCompressionCodec(), false)) .add(booleanProperty( @@ -192,6 +221,11 @@ public IcebergSessionProperties( "ORC: Max stripe row count", orcWriterConfig.getStripeMaxRowCount(), false)) + .add(integerProperty( + ORC_WRITER_MAX_ROW_GROUP_ROWS, + "ORC: Max number of rows in a row group", + orcWriterConfig.getRowGroupMaxRowCount(), + false)) .add(dataSizeProperty( ORC_WRITER_MAX_DICTIONARY_MEMORY, "ORC: Max dictionary memory", @@ -219,15 +253,16 @@ public IcebergSessionProperties( } }, false)) - .add(booleanProperty( - PARQUET_OPTIMIZED_READER_ENABLED, - "Use optimized Parquet reader", - parquetReaderConfig.isOptimizedReaderEnabled(), + .add(dataSizeProperty( + PARQUET_SMALL_FILE_THRESHOLD, + "Parquet: Size below which a parquet file will be read entirely", + parquetReaderConfig.getSmallFileThreshold(), + value -> validateMaxDataSize(PARQUET_SMALL_FILE_THRESHOLD, value, DataSize.valueOf(PARQUET_READER_MAX_SMALL_FILE_THRESHOLD)), false)) .add(booleanProperty( - PARQUET_OPTIMIZED_NESTED_READER_ENABLED, - "Use optimized Parquet reader for nested columns", - parquetReaderConfig.isOptimizedNestedReaderEnabled(), + PARQUET_IGNORE_STATISTICS, + "Ignore statistics from Parquet to allow querying files with corrupted or incorrect statistics", + parquetReaderConfig.isIgnoreStatistics(), false)) .add(dataSizeProperty( PARQUET_WRITER_BLOCK_SIZE, @@ -244,6 +279,18 @@ public IcebergSessionProperties( validateMaxDataSize(PARQUET_WRITER_PAGE_SIZE, value, DataSize.valueOf(PARQUET_WRITER_MAX_PAGE_SIZE)); }, false)) + .add(integerProperty( + PARQUET_WRITER_PAGE_VALUE_COUNT, + "Parquet: Writer page row count", + parquetWriterConfig.getPageValueCount(), + value -> { + if (value < PARQUET_WRITER_MIN_PAGE_VALUE_COUNT || value > PARQUET_WRITER_MAX_PAGE_VALUE_COUNT) { + throw new TrinoException( + INVALID_SESSION_PROPERTY, + format("%s must be between %s and %s: %s", PARQUET_WRITER_PAGE_VALUE_COUNT, PARQUET_WRITER_MIN_PAGE_VALUE_COUNT, PARQUET_WRITER_MAX_PAGE_VALUE_COUNT, value)); + } + }, + false)) .add(integerProperty( PARQUET_WRITER_BATCH_SIZE, "Parquet: Maximum number of rows passed to the writer in each batch", @@ -274,6 +321,11 @@ public IcebergSessionProperties( "Target maximum size of written files; the actual size may be larger", icebergConfig.getTargetMaxFileSize(), false)) + .add(dataSizeProperty( + IDLE_WRITER_MIN_FILE_SIZE, + "Minimum data written by a single partition writer before it can be consider as 'idle' and could be closed by the engine", + icebergConfig.getIdleWriterMinFileSize(), + false)) .add(booleanProperty( COLLECT_EXTENDED_STATISTICS_ON_WRITE, COLLECT_EXTENDED_STATISTICS_ON_WRITE_DESCRIPTION, @@ -311,6 +363,43 @@ public IcebergSessionProperties( "Enable sorted writing to tables with a specified sort order", icebergConfig.isSortedWritingEnabled(), false)) + .add(booleanProperty( + QUERY_PARTITION_FILTER_REQUIRED, + "Require filter on partition column", + icebergConfig.isQueryPartitionFilterRequired(), + false)) + .add(new PropertyMetadata<>( + QUERY_PARTITION_FILTER_REQUIRED_SCHEMAS, + "List of schemas for which filter on partition column is enforced.", + new ArrayType(VARCHAR), + Set.class, + icebergConfig.getQueryPartitionFilterRequiredSchemas(), + false, + object -> ((Collection) object).stream() + .map(String.class::cast) + .peek(property -> { + if (isNullOrEmpty(property)) { + throw new TrinoException(INVALID_SESSION_PROPERTY, format("Invalid null or empty value in %s property", QUERY_PARTITION_FILTER_REQUIRED_SCHEMAS)); + } + }) + .map(schema -> schema.toLowerCase(ENGLISH)) + .collect(toImmutableSet()), + value -> value)) + .add(booleanProperty( + INCREMENTAL_REFRESH_ENABLED, + "Enable Incremental refresh for MVs backed by Iceberg tables, when possible.", + icebergConfig.isIncrementalRefreshEnabled(), + false)) + .add(booleanProperty( + BUCKET_EXECUTION_ENABLED, + "Enable bucket-aware execution: use physical bucketing information to optimize queries", + icebergConfig.isBucketExecutionEnabled(), + false)) + .add(booleanProperty( + FILE_BASED_CONFLICT_DETECTION_ENABLED, + "Enable file-based conflict detection: take partition information from the actual written files as a source for the conflict detection system", + icebergConfig.isFileBasedConflictDetectionEnabled(), + false)) .build(); } @@ -397,14 +486,24 @@ public static int getOrcWriterMaxStripeRows(ConnectorSession session) return session.getProperty(ORC_WRITER_MAX_STRIPE_ROWS, Integer.class); } + public static int getOrcWriterMaxRowGroupRows(ConnectorSession session) + { + return session.getProperty(ORC_WRITER_MAX_ROW_GROUP_ROWS, Integer.class); + } + public static DataSize getOrcWriterMaxDictionaryMemory(ConnectorSession session) { return session.getProperty(ORC_WRITER_MAX_DICTIONARY_MEMORY, DataSize.class); } - public static HiveCompressionCodec getCompressionCodec(ConnectorSession session) + public static Optional getSplitSize(ConnectorSession session) + { + return Optional.ofNullable(session.getProperty(SPLIT_SIZE, DataSize.class)); + } + + public static HiveCompressionOption getCompressionCodec(ConnectorSession session) { - return session.getProperty(COMPRESSION_CODEC, HiveCompressionCodec.class); + return session.getProperty(COMPRESSION_CODEC, HiveCompressionOption.class); } public static boolean isUseFileSizeFromMetadata(ConnectorSession session) @@ -432,11 +531,26 @@ public static boolean isParquetOptimizedNestedReaderEnabled(ConnectorSession ses return session.getProperty(PARQUET_OPTIMIZED_NESTED_READER_ENABLED, Boolean.class); } + public static DataSize getParquetSmallFileThreshold(ConnectorSession session) + { + return session.getProperty(PARQUET_SMALL_FILE_THRESHOLD, DataSize.class); + } + + public static boolean isParquetIgnoreStatistics(ConnectorSession session) + { + return session.getProperty(PARQUET_IGNORE_STATISTICS, Boolean.class); + } + public static DataSize getParquetWriterPageSize(ConnectorSession session) { return session.getProperty(PARQUET_WRITER_PAGE_SIZE, DataSize.class); } + public static int getParquetWriterPageValueCount(ConnectorSession session) + { + return session.getProperty(PARQUET_WRITER_PAGE_VALUE_COUNT, Integer.class); + } + public static DataSize getParquetWriterBlockSize(ConnectorSession session) { return session.getProperty(PARQUET_WRITER_BLOCK_SIZE, DataSize.class); @@ -482,6 +596,11 @@ public static long getTargetMaxFileSize(ConnectorSession session) return session.getProperty(TARGET_MAX_FILE_SIZE, DataSize.class).toBytes(); } + public static long getIdleWriterMinFileSize(ConnectorSession session) + { + return session.getProperty(IDLE_WRITER_MIN_FILE_SIZE, DataSize.class).toBytes(); + } + public static Optional getHiveCatalogName(ConnectorSession session) { return Optional.ofNullable(session.getProperty(HIVE_CATALOG_NAME, String.class)); @@ -511,4 +630,32 @@ public static boolean isSortedWritingEnabled(ConnectorSession session) { return session.getProperty(SORTED_WRITING_ENABLED, Boolean.class); } + + public static boolean isQueryPartitionFilterRequired(ConnectorSession session) + { + return session.getProperty(QUERY_PARTITION_FILTER_REQUIRED, Boolean.class); + } + + @SuppressWarnings("unchecked cast") + public static Set getQueryPartitionFilterRequiredSchemas(ConnectorSession session) + { + Set queryPartitionFilterRequiredSchemas = (Set) session.getProperty(QUERY_PARTITION_FILTER_REQUIRED_SCHEMAS, Set.class); + requireNonNull(queryPartitionFilterRequiredSchemas, "queryPartitionFilterRequiredSchemas is null"); + return queryPartitionFilterRequiredSchemas; + } + + public static boolean isIncrementalRefreshEnabled(ConnectorSession session) + { + return session.getProperty(INCREMENTAL_REFRESH_ENABLED, Boolean.class); + } + + public static boolean isBucketExecutionEnabled(ConnectorSession session) + { + return session.getProperty(BUCKET_EXECUTION_ENABLED, Boolean.class); + } + + public static boolean isFileBasedConflictDetectionEnabled(ConnectorSession session) + { + return session.getProperty(FILE_BASED_CONFLICT_DETECTION_ENABLED, Boolean.class); + } } diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergSortingFileWriter.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergSortingFileWriter.java index ac5f875b1d80..1915d2d4422c 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergSortingFileWriter.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergSortingFileWriter.java @@ -23,7 +23,6 @@ import io.trino.spi.connector.SortOrder; import io.trino.spi.type.Type; import io.trino.spi.type.TypeOperators; -import org.apache.iceberg.Metrics; import java.io.Closeable; import java.util.List; @@ -65,9 +64,9 @@ public IcebergSortingFileWriter( } @Override - public Metrics getMetrics() + public FileMetrics getFileMetrics() { - return outputWriter.getMetrics(); + return outputWriter.getFileMetrics(); } @Override diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergSplit.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergSplit.java index 20719a9d8dfb..d4371fad0804 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergSplit.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergSplit.java @@ -19,12 +19,16 @@ import com.google.common.base.MoreObjects.ToStringHelper; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; +import io.airlift.slice.SizeOf; import io.trino.plugin.iceberg.delete.DeleteFile; import io.trino.spi.HostAddress; import io.trino.spi.SplitWeight; import io.trino.spi.connector.ConnectorSplit; +import io.trino.spi.predicate.TupleDomain; import java.util.List; +import java.util.Map; +import java.util.Optional; import static com.google.common.base.MoreObjects.toStringHelper; import static io.airlift.slice.SizeOf.estimatedSizeOf; @@ -40,11 +44,17 @@ public class IcebergSplit private final long start; private final long length; private final long fileSize; + private final long fileRecordCount; private final IcebergFileFormat fileFormat; + private final Optional> partitionValues; private final String partitionSpecJson; private final String partitionDataJson; private final List deletes; private final SplitWeight splitWeight; + private final TupleDomain fileStatisticsDomain; + private final Map fileIoProperties; + private final long dataSequenceNumber; + private final List addresses; @JsonCreator public IcebergSplit( @@ -52,21 +62,66 @@ public IcebergSplit( @JsonProperty("start") long start, @JsonProperty("length") long length, @JsonProperty("fileSize") long fileSize, + @JsonProperty("fileRecordCount") long fileRecordCount, @JsonProperty("fileFormat") IcebergFileFormat fileFormat, @JsonProperty("partitionSpecJson") String partitionSpecJson, @JsonProperty("partitionDataJson") String partitionDataJson, @JsonProperty("deletes") List deletes, - @JsonProperty("splitWeight") SplitWeight splitWeight) + @JsonProperty("splitWeight") SplitWeight splitWeight, + @JsonProperty("fileStatisticsDomain") TupleDomain fileStatisticsDomain, + @JsonProperty("fileIoProperties") Map fileIoProperties, + @JsonProperty("dataSequenceNumber") long dataSequenceNumber) + { + this( + path, + start, + length, + fileSize, + fileRecordCount, + fileFormat, + Optional.empty(), + partitionSpecJson, + partitionDataJson, + deletes, + splitWeight, + fileStatisticsDomain, + fileIoProperties, + ImmutableList.of(), + dataSequenceNumber); + } + + public IcebergSplit( + String path, + long start, + long length, + long fileSize, + long fileRecordCount, + IcebergFileFormat fileFormat, + Optional> partitionValues, + String partitionSpecJson, + String partitionDataJson, + List deletes, + SplitWeight splitWeight, + TupleDomain fileStatisticsDomain, + Map fileIoProperties, + List addresses, + long dataSequenceNumber) { this.path = requireNonNull(path, "path is null"); this.start = start; this.length = length; this.fileSize = fileSize; + this.fileRecordCount = fileRecordCount; this.fileFormat = requireNonNull(fileFormat, "fileFormat is null"); + this.partitionValues = requireNonNull(partitionValues, "partitionValues is null"); this.partitionSpecJson = requireNonNull(partitionSpecJson, "partitionSpecJson is null"); this.partitionDataJson = requireNonNull(partitionDataJson, "partitionDataJson is null"); this.deletes = ImmutableList.copyOf(requireNonNull(deletes, "deletes is null")); this.splitWeight = requireNonNull(splitWeight, "splitWeight is null"); + this.fileStatisticsDomain = requireNonNull(fileStatisticsDomain, "fileStatisticsDomain is null"); + this.fileIoProperties = ImmutableMap.copyOf(requireNonNull(fileIoProperties, "fileIoProperties is null")); + this.addresses = requireNonNull(addresses, "addresses is null"); + this.dataSequenceNumber = dataSequenceNumber; } @Override @@ -106,6 +161,12 @@ public long getFileSize() return fileSize; } + @JsonProperty + public long getFileRecordCount() + { + return fileRecordCount; + } + @JsonProperty public IcebergFileFormat getFileFormat() { @@ -118,6 +179,16 @@ public String getPartitionSpecJson() return partitionSpecJson; } + /** + * Trino (stack) values of the partition columns. The values are the result of evaluating + * the partition expressions on the partition data. + */ + @JsonIgnore + public Optional> getPartitionValues() + { + return partitionValues; + } + @JsonProperty public String getPartitionDataJson() { @@ -137,6 +208,24 @@ public SplitWeight getSplitWeight() return splitWeight; } + @JsonProperty + public TupleDomain getFileStatisticsDomain() + { + return fileStatisticsDomain; + } + + @JsonProperty + public Map getFileIoProperties() + { + return fileIoProperties; + } + + @JsonProperty + public long getDataSequenceNumber() + { + return dataSequenceNumber; + } + @Override public Object getInfo() { @@ -154,8 +243,11 @@ public long getRetainedSizeInBytes() + estimatedSizeOf(path) + estimatedSizeOf(partitionSpecJson) + estimatedSizeOf(partitionDataJson) - + estimatedSizeOf(deletes, DeleteFile::getRetainedSizeInBytes) - + splitWeight.getRetainedSizeInBytes(); + + estimatedSizeOf(deletes, DeleteFile::retainedSizeInBytes) + + splitWeight.getRetainedSizeInBytes() + + fileStatisticsDomain.getRetainedSizeInBytes(IcebergColumnHandle::getRetainedSizeInBytes) + + estimatedSizeOf(fileIoProperties, SizeOf::estimatedSizeOf, SizeOf::estimatedSizeOf) + + estimatedSizeOf(addresses, HostAddress::getRetainedSizeInBytes); } @Override @@ -164,7 +256,8 @@ public String toString() ToStringHelper helper = toStringHelper(this) .addValue(path) .add("start", start) - .add("length", length); + .add("length", length) + .add("fileStatisticsDomain", fileStatisticsDomain); if (!deletes.isEmpty()) { helper.add("deleteFiles", deletes.size()); helper.add("deleteRecords", deletes.stream() diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergSplitManager.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergSplitManager.java index 262c8c3c7583..382a49fa64c6 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergSplitManager.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergSplitManager.java @@ -14,10 +14,13 @@ package io.trino.plugin.iceberg; import com.google.common.collect.ImmutableList; +import com.google.common.util.concurrent.ListeningExecutorService; import com.google.inject.Inject; import io.airlift.units.Duration; -import io.trino.filesystem.TrinoFileSystemFactory; +import io.trino.filesystem.cache.CachingHostAddressProvider; import io.trino.plugin.base.classloader.ClassLoaderSafeConnectorSplitSource; +import io.trino.plugin.iceberg.functions.tablechanges.TableChangesFunctionHandle; +import io.trino.plugin.iceberg.functions.tablechanges.TableChangesSplitSource; import io.trino.spi.connector.ConnectorSession; import io.trino.spi.connector.ConnectorSplitManager; import io.trino.spi.connector.ConnectorSplitSource; @@ -26,11 +29,19 @@ import io.trino.spi.connector.Constraint; import io.trino.spi.connector.DynamicFilter; import io.trino.spi.connector.FixedSplitSource; +import io.trino.spi.function.table.ConnectorTableFunctionHandle; import io.trino.spi.type.TypeManager; +import org.apache.iceberg.CombinedScanTask; +import org.apache.iceberg.DataOperations; +import org.apache.iceberg.FileScanTask; +import org.apache.iceberg.Scan; +import org.apache.iceberg.Snapshot; import org.apache.iceberg.Table; -import org.apache.iceberg.TableScan; +import org.apache.iceberg.metrics.InMemoryMetricsReporter; +import org.apache.iceberg.util.SnapshotUtil; + +import java.util.concurrent.ExecutorService; -import static com.google.common.util.concurrent.MoreExecutors.newDirectExecutorService; import static io.trino.plugin.iceberg.IcebergSessionProperties.getDynamicFilteringWaitTimeout; import static io.trino.plugin.iceberg.IcebergSessionProperties.getMinimumAssignedSplitWeight; import static io.trino.spi.connector.FixedSplitSource.emptySplitSource; @@ -43,20 +54,26 @@ public class IcebergSplitManager private final IcebergTransactionManager transactionManager; private final TypeManager typeManager; - private final TrinoFileSystemFactory fileSystemFactory; - private final boolean asyncIcebergSplitProducer; + private final IcebergFileSystemFactory fileSystemFactory; + private final ListeningExecutorService splitSourceExecutor; + private final ExecutorService icebergPlanningExecutor; + private final CachingHostAddressProvider cachingHostAddressProvider; @Inject public IcebergSplitManager( IcebergTransactionManager transactionManager, TypeManager typeManager, - TrinoFileSystemFactory fileSystemFactory, - @AsyncIcebergSplitProducer boolean asyncIcebergSplitProducer) + IcebergFileSystemFactory fileSystemFactory, + @ForIcebergSplitManager ListeningExecutorService splitSourceExecutor, + @ForIcebergScanPlanning ExecutorService icebergPlanningExecutor, + CachingHostAddressProvider cachingHostAddressProvider) { this.transactionManager = requireNonNull(transactionManager, "transactionManager is null"); this.typeManager = requireNonNull(typeManager, "typeManager is null"); this.fileSystemFactory = requireNonNull(fileSystemFactory, "fileSystemFactory is null"); - this.asyncIcebergSplitProducer = asyncIcebergSplitProducer; + this.splitSourceExecutor = requireNonNull(splitSourceExecutor, "splitSourceExecutor is null"); + this.icebergPlanningExecutor = requireNonNull(icebergPlanningExecutor, "icebergPlanningExecutor is null"); + this.cachingHostAddressProvider = requireNonNull(cachingHostAddressProvider, "cachingHostAddressProvider is null"); } @Override @@ -76,27 +93,72 @@ public ConnectorSplitSource getSplits( return emptySplitSource(); } - Table icebergTable = transactionManager.get(transaction, session.getIdentity()).getIcebergTable(session, table.getSchemaTableName()); + IcebergMetadata icebergMetadata = transactionManager.get(transaction, session.getIdentity()); + Table icebergTable = icebergMetadata.getIcebergTable(session, table.getSchemaTableName()); Duration dynamicFilteringWaitTimeout = getDynamicFilteringWaitTimeout(session); + InMemoryMetricsReporter metricsReporter = new InMemoryMetricsReporter(); + Scan scan = getScan(icebergMetadata, icebergTable, table, icebergPlanningExecutor); - TableScan tableScan = icebergTable.newScan() - .useSnapshot(table.getSnapshotId().get()); - if (!asyncIcebergSplitProducer) { - tableScan = tableScan.planWith(newDirectExecutorService()); - } IcebergSplitSource splitSource = new IcebergSplitSource( fileSystemFactory, session, table, - tableScan, + icebergTable, + scan, table.getMaxScannedFileSize(), dynamicFilter, dynamicFilteringWaitTimeout, constraint, typeManager, table.isRecordScannedFiles(), - getMinimumAssignedSplitWeight(session)); + getMinimumAssignedSplitWeight(session), + cachingHostAddressProvider, + splitSourceExecutor); return new ClassLoaderSafeConnectorSplitSource(splitSource, IcebergSplitManager.class.getClassLoader()); } + + private Scan getScan(IcebergMetadata icebergMetadata, Table icebergTable, IcebergTableHandle table, ExecutorService executor) + { + Long fromSnapshot = icebergMetadata.getIncrementalRefreshFromSnapshot().orElse(null); + if (fromSnapshot != null) { + // check if fromSnapshot is still part of the table's snapshot history + if (SnapshotUtil.isAncestorOf(icebergTable, fromSnapshot)) { + boolean containsModifiedRows = false; + for (Snapshot snapshot : SnapshotUtil.ancestorsBetween(icebergTable, icebergTable.currentSnapshot().snapshotId(), fromSnapshot)) { + if (snapshot.operation().equals(DataOperations.OVERWRITE) || snapshot.operation().equals(DataOperations.DELETE)) { + containsModifiedRows = true; + break; + } + } + if (!containsModifiedRows) { + return icebergTable.newIncrementalAppendScan().fromSnapshotExclusive(fromSnapshot).planWith(executor); + } + } + // fromSnapshot is missing (could be due to snapshot expiration or rollback), or snapshot range contains modifications + // (deletes or overwrites), so we cannot perform incremental refresh. Falling back to full refresh. + icebergMetadata.disableIncrementalRefresh(); + } + return icebergTable.newScan().useSnapshot(table.getSnapshotId().get()).planWith(executor); + } + + @Override + public ConnectorSplitSource getSplits( + ConnectorTransactionHandle transaction, + ConnectorSession session, + ConnectorTableFunctionHandle function) + { + if (function instanceof TableChangesFunctionHandle functionHandle) { + Table icebergTable = transactionManager.get(transaction, session.getIdentity()).getIcebergTable(session, functionHandle.schemaTableName()); + + TableChangesSplitSource tableChangesSplitSource = new TableChangesSplitSource( + icebergTable, + icebergTable.newIncrementalChangelogScan() + .fromSnapshotExclusive(functionHandle.startSnapshotId()) + .toSnapshot(functionHandle.endSnapshotId())); + return new ClassLoaderSafeConnectorSplitSource(tableChangesSplitSource, IcebergSplitManager.class.getClassLoader()); + } + + throw new IllegalStateException("Unknown table function: " + function); + } } diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergSplitSource.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergSplitSource.java index 6e4ad08f6eea..cf1ce850c573 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergSplitSource.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergSplitSource.java @@ -15,18 +15,24 @@ import com.google.common.annotations.VisibleForTesting; import com.google.common.base.Stopwatch; +import com.google.common.cache.CacheBuilder; import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; +import com.google.common.collect.Iterators; +import com.google.common.collect.Maps; import com.google.common.io.Closer; +import com.google.common.util.concurrent.ListenableFuture; +import com.google.common.util.concurrent.ListeningExecutorService; +import com.google.errorprone.annotations.concurrent.GuardedBy; +import io.airlift.log.Logger; import io.airlift.units.DataSize; import io.airlift.units.Duration; -import io.trino.filesystem.Location; -import io.trino.filesystem.TrinoFileSystemFactory; -import io.trino.filesystem.TrinoInputFile; +import io.trino.cache.NonEvictableCache; +import io.trino.filesystem.cache.CachingHostAddressProvider; import io.trino.plugin.iceberg.delete.DeleteFile; import io.trino.plugin.iceberg.util.DataFileWithDeleteFiles; import io.trino.spi.SplitWeight; -import io.trino.spi.TrinoException; import io.trino.spi.connector.ColumnHandle; import io.trino.spi.connector.ConnectorSession; import io.trino.spi.connector.ConnectorSplit; @@ -40,10 +46,17 @@ import io.trino.spi.predicate.ValueSet; import io.trino.spi.type.TypeManager; import jakarta.annotation.Nullable; +import org.apache.iceberg.BaseFileScanTask; +import org.apache.iceberg.CombinedScanTask; +import org.apache.iceberg.ContentFile; import org.apache.iceberg.FileScanTask; +import org.apache.iceberg.PartitionField; +import org.apache.iceberg.PartitionSpec; import org.apache.iceberg.PartitionSpecParser; +import org.apache.iceberg.Scan; import org.apache.iceberg.Schema; -import org.apache.iceberg.TableScan; +import org.apache.iceberg.StructLike; +import org.apache.iceberg.Table; import org.apache.iceberg.expressions.Expression; import org.apache.iceberg.io.CloseableIterable; import org.apache.iceberg.io.CloseableIterator; @@ -57,11 +70,14 @@ import java.util.Iterator; import java.util.List; import java.util.Map; +import java.util.Objects; import java.util.Optional; import java.util.OptionalLong; import java.util.Set; import java.util.concurrent.CompletableFuture; +import java.util.function.Predicate; import java.util.function.Supplier; +import java.util.stream.Stream; import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.base.Preconditions.checkState; @@ -71,89 +87,129 @@ import static com.google.common.collect.ImmutableSet.toImmutableSet; import static com.google.common.collect.Sets.intersection; import static com.google.common.math.LongMath.saturatedAdd; +import static io.airlift.concurrent.MoreFutures.toCompletableFuture; import static io.airlift.slice.Slices.utf8Slice; +import static io.trino.cache.CacheUtils.uncheckedCacheGet; +import static io.trino.cache.SafeCaches.buildNonEvictableCache; +import static io.trino.plugin.iceberg.ExpressionConverter.isConvertibleToIcebergExpression; import static io.trino.plugin.iceberg.ExpressionConverter.toIcebergExpression; -import static io.trino.plugin.iceberg.IcebergColumnHandle.fileModifiedTimeColumnHandle; -import static io.trino.plugin.iceberg.IcebergColumnHandle.pathColumnHandle; -import static io.trino.plugin.iceberg.IcebergErrorCode.ICEBERG_FILESYSTEM_ERROR; +import static io.trino.plugin.iceberg.IcebergExceptions.translateMetadataException; import static io.trino.plugin.iceberg.IcebergMetadataColumn.isMetadataColumnId; -import static io.trino.plugin.iceberg.IcebergSplitManager.ICEBERG_DOMAIN_COMPACTION_THRESHOLD; +import static io.trino.plugin.iceberg.IcebergSessionProperties.getSplitSize; import static io.trino.plugin.iceberg.IcebergTypes.convertIcebergValueToTrino; -import static io.trino.plugin.iceberg.IcebergUtil.deserializePartitionValue; import static io.trino.plugin.iceberg.IcebergUtil.getColumnHandle; +import static io.trino.plugin.iceberg.IcebergUtil.getFileModifiedTimePathDomain; +import static io.trino.plugin.iceberg.IcebergUtil.getModificationTime; +import static io.trino.plugin.iceberg.IcebergUtil.getPartitionDomain; import static io.trino.plugin.iceberg.IcebergUtil.getPartitionKeys; +import static io.trino.plugin.iceberg.IcebergUtil.getPartitionValues; +import static io.trino.plugin.iceberg.IcebergUtil.getPathDomain; import static io.trino.plugin.iceberg.IcebergUtil.primitiveFieldTypes; +import static io.trino.plugin.iceberg.StructLikeWrapperWithFieldIdToIndex.createStructLikeWrapper; import static io.trino.plugin.iceberg.TypeConverter.toIcebergType; import static io.trino.spi.type.DateTimeEncoding.packDateTimeWithZone; import static io.trino.spi.type.TimeZoneKey.UTC_KEY; import static java.util.Collections.emptyIterator; import static java.util.Objects.requireNonNull; -import static java.util.concurrent.CompletableFuture.completedFuture; import static java.util.concurrent.TimeUnit.MILLISECONDS; +import static org.apache.iceberg.FileContent.EQUALITY_DELETES; +import static org.apache.iceberg.FileContent.POSITION_DELETES; import static org.apache.iceberg.types.Conversions.fromByteBuffer; public class IcebergSplitSource implements ConnectorSplitSource { + private static final Logger log = Logger.get(IcebergSplitSource.class); private static final ConnectorSplitBatch EMPTY_BATCH = new ConnectorSplitBatch(ImmutableList.of(), false); private static final ConnectorSplitBatch NO_MORE_SPLITS_BATCH = new ConnectorSplitBatch(ImmutableList.of(), true); - private final TrinoFileSystemFactory fileSystemFactory; + private final IcebergFileSystemFactory fileSystemFactory; private final ConnectorSession session; private final IcebergTableHandle tableHandle; - private final TableScan tableScan; + private final Map fileIoProperties; + private final Scan tableScan; private final Optional maxScannedFileSizeInBytes; private final Map fieldIdToType; private final DynamicFilter dynamicFilter; private final long dynamicFilteringWaitTimeoutMillis; private final Stopwatch dynamicFilterWaitStopwatch; - private final Constraint constraint; + private final PartitionConstraintMatcher partitionConstraintMatcher; private final TypeManager typeManager; + @GuardedBy("closer") private final Closer closer = Closer.create(); + @GuardedBy("closer") + private boolean closed; + @GuardedBy("closer") + private ListenableFuture currentBatchFuture; private final double minimumAssignedSplitWeight; + private final Set projectedBaseColumns; private final TupleDomain dataColumnPredicate; + private final Domain partitionDomain; private final Domain pathDomain; private final Domain fileModifiedTimeDomain; private final OptionalLong limit; + private final Set predicatedColumnIds; + private final ListeningExecutorService executor; + @GuardedBy("this") private TupleDomain pushedDownDynamicFilterPredicate; + @GuardedBy("this") private CloseableIterable fileScanIterable; + @GuardedBy("this") private long targetSplitSize; + @GuardedBy("this") private CloseableIterator fileScanIterator; - private Iterator fileTasksIterator = emptyIterator(); - private boolean fileHasAnyDeletions; + @GuardedBy("this") + private Iterator fileTasksIterator = emptyIterator(); private final boolean recordScannedFiles; + private final int currentSpecId; + @GuardedBy("this") private final ImmutableSet.Builder scannedFiles = ImmutableSet.builder(); + @GuardedBy("this") + @Nullable + private Map> scannedFilesByPartition = new HashMap<>(); + @GuardedBy("this") private long outputRowsLowerBound; + private final CachingHostAddressProvider cachingHostAddressProvider; + private volatile boolean finished; public IcebergSplitSource( - TrinoFileSystemFactory fileSystemFactory, + IcebergFileSystemFactory fileSystemFactory, ConnectorSession session, IcebergTableHandle tableHandle, - TableScan tableScan, + Table icebergTable, + Scan tableScan, Optional maxScannedFileSize, DynamicFilter dynamicFilter, Duration dynamicFilteringWaitTimeout, Constraint constraint, TypeManager typeManager, boolean recordScannedFiles, - double minimumAssignedSplitWeight) + double minimumAssignedSplitWeight, + CachingHostAddressProvider cachingHostAddressProvider, + ListeningExecutorService executor) { this.fileSystemFactory = requireNonNull(fileSystemFactory, "fileSystemFactory is null"); this.session = requireNonNull(session, "session is null"); this.tableHandle = requireNonNull(tableHandle, "tableHandle is null"); + this.fileIoProperties = requireNonNull(icebergTable.io().properties(), "fileIoProperties is null"); this.tableScan = requireNonNull(tableScan, "tableScan is null"); this.maxScannedFileSizeInBytes = maxScannedFileSize.map(DataSize::toBytes); this.fieldIdToType = primitiveFieldTypes(tableScan.schema()); this.dynamicFilter = requireNonNull(dynamicFilter, "dynamicFilter is null"); this.dynamicFilteringWaitTimeoutMillis = dynamicFilteringWaitTimeout.toMillis(); this.dynamicFilterWaitStopwatch = Stopwatch.createStarted(); - this.constraint = requireNonNull(constraint, "constraint is null"); + this.partitionConstraintMatcher = new PartitionConstraintMatcher(constraint); this.typeManager = requireNonNull(typeManager, "typeManager is null"); this.recordScannedFiles = recordScannedFiles; + this.currentSpecId = icebergTable.spec().specId(); this.minimumAssignedSplitWeight = minimumAssignedSplitWeight; + this.projectedBaseColumns = tableHandle.getProjectedColumns().stream() + .map(column -> column.getBaseColumnIdentity().getId()) + .collect(toImmutableSet()); this.dataColumnPredicate = tableHandle.getEnforcedPredicate().filter((column, domain) -> !isMetadataColumnId(column.getId())); + this.partitionDomain = getPartitionDomain(tableHandle.getEnforcedPredicate()); this.pathDomain = getPathDomain(tableHandle.getEnforcedPredicate()); checkArgument( tableHandle.getUnenforcedPredicate().isAll() || tableHandle.getLimit().isEmpty(), @@ -161,7 +217,15 @@ public IcebergSplitSource( tableHandle.getLimit(), tableHandle.getUnenforcedPredicate()); this.limit = tableHandle.getLimit(); + this.predicatedColumnIds = Stream.concat( + tableHandle.getUnenforcedPredicate().getDomains().orElse(ImmutableMap.of()).keySet().stream(), + dynamicFilter.getColumnsCovered().stream() + .map(IcebergColumnHandle.class::cast)) + .map(IcebergColumnHandle::getId) + .collect(toImmutableSet()); this.fileModifiedTimeDomain = getFileModifiedTimePathDomain(tableHandle.getEnforcedPredicate()); + this.cachingHostAddressProvider = requireNonNull(cachingHostAddressProvider, "cachingHostAddressProvider is null"); + this.executor = requireNonNull(executor, "executor is null"); } @Override @@ -170,154 +234,238 @@ public CompletableFuture getNextBatch(int maxSize) long timeLeft = dynamicFilteringWaitTimeoutMillis - dynamicFilterWaitStopwatch.elapsed(MILLISECONDS); if (dynamicFilter.isAwaitable() && timeLeft > 0) { return dynamicFilter.isBlocked() - .thenApply(ignored -> EMPTY_BATCH) + .thenApply(ignore -> EMPTY_BATCH) .completeOnTimeout(EMPTY_BATCH, timeLeft, MILLISECONDS); } + ListenableFuture nextBatchFuture; + synchronized (closer) { + checkState(!closed, "already closed"); + checkState(currentBatchFuture == null || currentBatchFuture.isDone(), "previous batch future is not done"); + + // Avoids blocking the calling (scheduler) thread when producing splits, allowing other split sources to + // start loading splits in parallel to each other + nextBatchFuture = executor.submit(() -> getNextBatchInternal(maxSize)); + currentBatchFuture = nextBatchFuture; + } + + return toCompletableFuture(nextBatchFuture).exceptionally(t -> { + throw translateMetadataException(t, tableHandle.getSchemaTableName().toString()); + }); + } + + private synchronized ConnectorSplitBatch getNextBatchInternal(int maxSize) + { if (fileScanIterable == null) { - // Used to avoid duplicating work if the Dynamic Filter was already pushed down to the Iceberg API - boolean dynamicFilterIsComplete = dynamicFilter.isComplete(); - this.pushedDownDynamicFilterPredicate = dynamicFilter.getCurrentPredicate().transformKeys(IcebergColumnHandle.class::cast); - TupleDomain fullPredicate = tableHandle.getUnenforcedPredicate() - .intersect(pushedDownDynamicFilterPredicate); - // TODO: (https://github.com/trinodb/trino/issues/9743): Consider removing TupleDomain#simplify - TupleDomain simplifiedPredicate = fullPredicate.simplify(ICEBERG_DOMAIN_COMPACTION_THRESHOLD); - boolean usedSimplifiedPredicate = !simplifiedPredicate.equals(fullPredicate); - if (usedSimplifiedPredicate) { - // Pushed down predicate was simplified, always evaluate it against individual splits - this.pushedDownDynamicFilterPredicate = TupleDomain.all(); - } + this.pushedDownDynamicFilterPredicate = dynamicFilter.getCurrentPredicate() + .transformKeys(IcebergColumnHandle.class::cast) + .filter((columnHandle, domain) -> isConvertibleToIcebergExpression(domain)); - TupleDomain effectivePredicate = dataColumnPredicate - .intersect(simplifiedPredicate); + TupleDomain effectivePredicate = TupleDomain.intersect( + ImmutableList.of(dataColumnPredicate, tableHandle.getUnenforcedPredicate(), pushedDownDynamicFilterPredicate)); if (effectivePredicate.isNone()) { finish(); - return completedFuture(NO_MORE_SPLITS_BATCH); + return NO_MORE_SPLITS_BATCH; } Expression filterExpression = toIcebergExpression(effectivePredicate); - // If the Dynamic Filter will be evaluated against each file, stats are required. Otherwise, skip them. - boolean requiresColumnStats = usedSimplifiedPredicate || !dynamicFilterIsComplete; - TableScan scan = tableScan.filter(filterExpression); - if (requiresColumnStats) { - scan = scan.includeColumnStats(); + Scan scan = (Scan) tableScan.filter(filterExpression); + // Use stats to populate fileStatisticsDomain if there are predicated columns. Otherwise, skip them. + if (!predicatedColumnIds.isEmpty()) { + Schema schema = tableScan.schema(); + scan = (Scan) scan.includeColumnStats( + predicatedColumnIds.stream() + .map(schema::findColumnName) + // Newly added column may not be found in current snapshot schema until new files are added + .filter(Objects::nonNull) + .collect(toImmutableList())); + } + + synchronized (closer) { + checkState(!closed, "split source is closed"); + this.fileScanIterable = closer.register(scan.planFiles()); + this.targetSplitSize = getSplitSize(session) + .map(DataSize::toBytes) + .orElseGet(tableScan::targetSplitSize); + this.fileScanIterator = closer.register(fileScanIterable.iterator()); + this.fileTasksIterator = emptyIterator(); } - this.fileScanIterable = closer.register(scan.planFiles()); - this.targetSplitSize = tableScan.targetSplitSize(); - this.fileScanIterator = closer.register(fileScanIterable.iterator()); - this.fileTasksIterator = emptyIterator(); } TupleDomain dynamicFilterPredicate = dynamicFilter.getCurrentPredicate() .transformKeys(IcebergColumnHandle.class::cast); if (dynamicFilterPredicate.isNone()) { finish(); - return completedFuture(NO_MORE_SPLITS_BATCH); + return NO_MORE_SPLITS_BATCH; } List splits = new ArrayList<>(maxSize); while (splits.size() < maxSize && (fileTasksIterator.hasNext() || fileScanIterator.hasNext())) { if (!fileTasksIterator.hasNext()) { - FileScanTask wholeFileTask = fileScanIterator.next(); - fileTasksIterator = wholeFileTask.split(targetSplitSize).iterator(); - fileHasAnyDeletions = false; + if (limit.isPresent() && limit.getAsLong() <= outputRowsLowerBound) { + finish(); + break; + } + + List fileScanTasks = processFileScanTask(dynamicFilterPredicate); + if (fileScanTasks.isEmpty()) { + continue; + } + + fileTasksIterator = prepareFileTasksIterator(fileScanTasks); // In theory, .split() could produce empty iterator, so let's evaluate the outer loop condition again. continue; } - FileScanTask scanTask = fileTasksIterator.next(); - fileHasAnyDeletions = fileHasAnyDeletions || !scanTask.deletes().isEmpty(); - if (scanTask.deletes().isEmpty() && - maxScannedFileSizeInBytes.isPresent() && - scanTask.file().fileSizeInBytes() > maxScannedFileSizeInBytes.get()) { - continue; - } + splits.add(toIcebergSplit(fileTasksIterator.next())); + } + if (!fileScanIterator.hasNext() && !fileTasksIterator.hasNext()) { + finish(); + } + return new ConnectorSplitBatch(splits, isFinished()); + } - if (!pathDomain.includesNullableValue(utf8Slice(scanTask.file().path().toString()))) { - continue; - } - if (!fileModifiedTimeDomain.isAll()) { - long fileModifiedTime = getModificationTime(scanTask.file().path().toString()); - if (!fileModifiedTimeDomain.includesNullableValue(packDateTimeWithZone(fileModifiedTime, UTC_KEY))) { - continue; - } + private synchronized Iterator prepareFileTasksIterator(List fileScanTasks) + { + ImmutableList.Builder scanTaskBuilder = ImmutableList.builder(); + for (FileScanTaskWithDomain fileScanTaskWithDomain : fileScanTasks) { + FileScanTask wholeFileTask = fileScanTaskWithDomain.fileScanTask(); + if (recordScannedFiles) { + // Equality deletes can only be cleaned up if the whole table has been optimized. + // Equality and position deletes may apply to many files, however position deletes are always local to a partition + // https://github.com/apache/iceberg/blob/70c506ebad2dfc6d61b99c05efd59e884282bfa6/core/src/main/java/org/apache/iceberg/deletes/DeleteGranularity.java#L61 + // OPTIMIZE supports only enforced predicates which select whole partitions, so if there is no path or fileModifiedTime predicate, then we can clean up position deletes + List fullyAppliedDeletes = wholeFileTask.deletes().stream() + .filter(deleteFile -> switch (deleteFile.content()) { + case POSITION_DELETES -> partitionDomain.isAll() && pathDomain.isAll() && fileModifiedTimeDomain.isAll(); + case EQUALITY_DELETES -> tableHandle.getEnforcedPredicate().isAll(); + case DATA -> throw new IllegalStateException("Unexpected delete file: " + deleteFile); + }) + .collect(toImmutableList()); + scannedFiles.add(new DataFileWithDeleteFiles(wholeFileTask.file(), fullyAppliedDeletes)); } - IcebergSplit icebergSplit = toIcebergSplit(scanTask); - - Schema fileSchema = scanTask.spec().schema(); - Map> partitionKeys = getPartitionKeys(scanTask); - - Set identityPartitionColumns = partitionKeys.keySet().stream() - .map(fieldId -> getColumnHandle(fileSchema.findField(fieldId), typeManager)) - .collect(toImmutableSet()); - - Supplier> partitionValues = memoize(() -> { - Map bindings = new HashMap<>(); - for (IcebergColumnHandle partitionColumn : identityPartitionColumns) { - Object partitionValue = deserializePartitionValue( - partitionColumn.getType(), - partitionKeys.get(partitionColumn.getId()).orElse(null), - partitionColumn.getName()); - NullableValue bindingValue = new NullableValue(partitionColumn.getType(), partitionValue); - bindings.put(partitionColumn, bindingValue); - } - return bindings; - }); - - if (!dynamicFilterPredicate.isAll() && !dynamicFilterPredicate.equals(pushedDownDynamicFilterPredicate)) { - if (!partitionMatchesPredicate( - identityPartitionColumns, - partitionValues, - dynamicFilterPredicate)) { - continue; - } - if (!fileMatchesPredicate( - fieldIdToType, - dynamicFilterPredicate, - scanTask.file().lowerBounds(), - scanTask.file().upperBounds(), - scanTask.file().nullValueCounts())) { - continue; - } + + boolean fileHasNoDeletions = wholeFileTask.deletes().isEmpty(); + if (fileHasNoDeletions) { + // There were no deletions, so we will produce splits covering the whole file + outputRowsLowerBound = saturatedAdd(outputRowsLowerBound, wholeFileTask.file().recordCount()); } - if (!partitionMatchesConstraint(identityPartitionColumns, partitionValues, constraint)) { - continue; + + if (fileHasNoDeletions && noDataColumnsProjected(wholeFileTask)) { + scanTaskBuilder.add(fileScanTaskWithDomain); } - if (recordScannedFiles) { - // Positional and Equality deletes can only be cleaned up if the whole table has been optimized. - // Equality deletes may apply to many files, and position deletes may be grouped together. This makes it difficult to know if they are obsolete. - List fullyAppliedDeletes = tableHandle.getEnforcedPredicate().isAll() ? scanTask.deletes() : ImmutableList.of(); - scannedFiles.add(new DataFileWithDeleteFiles(scanTask.file(), fullyAppliedDeletes)); + else { + scanTaskBuilder.addAll(fileScanTaskWithDomain.split(targetSplitSize)); } - if (!fileTasksIterator.hasNext()) { - // This is the last task for this file - if (!fileHasAnyDeletions) { - // There were no deletions, so we produced splits covering the whole file - outputRowsLowerBound = saturatedAdd(outputRowsLowerBound, scanTask.file().recordCount()); - if (limit.isPresent() && limit.getAsLong() <= outputRowsLowerBound) { - finish(); - } - } + } + return scanTaskBuilder.build().iterator(); + } + + private synchronized List processFileScanTask(TupleDomain dynamicFilterPredicate) + { + FileScanTask wholeFileTask = fileScanIterator.next(); + boolean fileHasNoDeletions = wholeFileTask.deletes().isEmpty(); + FileScanTaskWithDomain fileScanTaskWithDomain = createFileScanTaskWithDomain(wholeFileTask); + if (pruneFileScanTask(fileScanTaskWithDomain, fileHasNoDeletions, dynamicFilterPredicate)) { + return ImmutableList.of(); + } + + if (!recordScannedFiles || scannedFilesByPartition == null) { + return ImmutableList.of(fileScanTaskWithDomain); + } + + // Assess if the partition that wholeFileTask belongs to should be included for OPTIMIZE + // If file was partitioned under an old spec, OPTIMIZE may be able to merge it with another file under new partitioning spec + // We don't know which partition of new spec this file belongs to, so we include all files in OPTIMIZE + if (currentSpecId != wholeFileTask.spec().specId()) { + Stream allQueuedTasks = scannedFilesByPartition.values().stream() + .filter(Optional::isPresent) + .map(Optional::get); + scannedFilesByPartition = null; + return Stream.concat(allQueuedTasks, Stream.of(fileScanTaskWithDomain)).collect(toImmutableList()); + } + StructLikeWrapperWithFieldIdToIndex structLikeWrapperWithFieldIdToIndex = createStructLikeWrapper(wholeFileTask); + Optional alreadyQueuedFileTask = scannedFilesByPartition.get(structLikeWrapperWithFieldIdToIndex); + if (alreadyQueuedFileTask != null) { + // Optional.empty() is a marker for partitions where we've seen enough files to avoid skipping them from OPTIMIZE + if (alreadyQueuedFileTask.isEmpty()) { + return ImmutableList.of(fileScanTaskWithDomain); } - splits.add(icebergSplit); + scannedFilesByPartition.put(structLikeWrapperWithFieldIdToIndex, Optional.empty()); + return ImmutableList.of(alreadyQueuedFileTask.get(), fileScanTaskWithDomain); } - return completedFuture(new ConnectorSplitBatch(splits, isFinished())); + // If file has no deletions, and it's the only file seen so far for the partition + // then we skip it from splits generation unless we encounter another file in the same partition + if (fileHasNoDeletions) { + scannedFilesByPartition.put(structLikeWrapperWithFieldIdToIndex, Optional.of(fileScanTaskWithDomain)); + return ImmutableList.of(); + } + scannedFilesByPartition.put(structLikeWrapperWithFieldIdToIndex, Optional.empty()); + return ImmutableList.of(fileScanTaskWithDomain); } - private long getModificationTime(String path) + private synchronized boolean pruneFileScanTask(FileScanTaskWithDomain fileScanTaskWithDomain, boolean fileHasNoDeletions, TupleDomain dynamicFilterPredicate) { - try { - TrinoInputFile inputFile = fileSystemFactory.create(session).newInputFile(Location.of(path)); - return inputFile.lastModified().toEpochMilli(); + BaseFileScanTask fileScanTask = (BaseFileScanTask) fileScanTaskWithDomain.fileScanTask(); + if (fileHasNoDeletions && + maxScannedFileSizeInBytes.isPresent() && + fileScanTask.file().fileSizeInBytes() > maxScannedFileSizeInBytes.get()) { + return true; } - catch (IOException e) { - throw new TrinoException(ICEBERG_FILESYSTEM_ERROR, "Failed to get file modification time: " + path, e); + + if (!partitionDomain.isAll()) { + String partition = fileScanTask.spec().partitionToPath(fileScanTask.partition()); + if (!partitionDomain.includesNullableValue(utf8Slice(partition))) { + return true; + } + } + if (!pathDomain.isAll() && !pathDomain.includesNullableValue(utf8Slice(fileScanTask.file().location()))) { + return true; } + if (!fileModifiedTimeDomain.isAll()) { + long fileModifiedTime = getModificationTime(fileScanTask.file().path().toString(), fileSystemFactory.create(session.getIdentity(), fileIoProperties)); + if (!fileModifiedTimeDomain.includesNullableValue(packDateTimeWithZone(fileModifiedTime, UTC_KEY))) { + return true; + } + } + + Schema fileSchema = fileScanTask.schema(); + Map> partitionKeys = getPartitionKeys(fileScanTask); + + Set identityPartitionColumns = partitionKeys.keySet().stream() + .map(fieldId -> getColumnHandle(fileSchema.findField(fieldId), typeManager)) + .collect(toImmutableSet()); + + Supplier> partitionValues = memoize(() -> getPartitionValues(identityPartitionColumns, partitionKeys)); + + if (!dynamicFilterPredicate.isAll() && !dynamicFilterPredicate.equals(pushedDownDynamicFilterPredicate)) { + if (!partitionMatchesPredicate( + identityPartitionColumns, + partitionValues, + dynamicFilterPredicate)) { + return true; + } + if (!fileScanTaskWithDomain.fileStatisticsDomain().overlaps(dynamicFilterPredicate)) { + return true; + } + } + + return !partitionConstraintMatcher.matches(identityPartitionColumns, partitionValues); } - private void finish() + private boolean noDataColumnsProjected(FileScanTask fileScanTask) { - close(); + return fileScanTask.spec().fields().stream() + .filter(partitionField -> partitionField.transform().isIdentity()) + .map(PartitionField::sourceId) + .collect(toImmutableSet()) + .containsAll(projectedBaseColumns); + } + + private synchronized void finish() + { + closeInternal(false); + this.finished = true; this.fileScanIterable = CloseableIterable.empty(); this.fileScanIterator = CloseableIterator.empty(); this.fileTasksIterator = emptyIterator(); @@ -326,7 +474,7 @@ private void finish() @Override public boolean isFinished() { - return fileScanIterator != null && !fileScanIterator.hasNext() && !fileTasksIterator.hasNext(); + return finished; } @Override @@ -336,37 +484,84 @@ public Optional> getTableExecuteSplitsInfo() if (!recordScannedFiles) { return Optional.empty(); } - return Optional.of(ImmutableList.copyOf(scannedFiles.build())); + long filesSkipped = 0; + List splitsInfo; + synchronized (this) { + if (scannedFilesByPartition != null) { + filesSkipped = scannedFilesByPartition.values().stream() + .filter(Optional::isPresent) + .count(); + scannedFilesByPartition = null; + } + splitsInfo = ImmutableList.copyOf(scannedFiles.build()); + } + log.info("Generated %d splits, skipped %d files for OPTIMIZE", splitsInfo.size(), filesSkipped); + return Optional.of(splitsInfo); } @Override public void close() { - try { - closer.close(); + closeInternal(true); + } + + private void closeInternal(boolean interruptIfRunning) + { + synchronized (closer) { + if (!closed) { + closed = true; + // don't cancel the current batch future during normal finishing cleanup + if (interruptIfRunning && currentBatchFuture != null) { + currentBatchFuture.cancel(true); + } + // release the reference to the current future unconditionally to avoid OOMs + currentBatchFuture = null; + try { + closer.close(); + } + catch (IOException e) { + throw new UncheckedIOException(e); + } + } } - catch (IOException e) { - throw new UncheckedIOException(e); + } + + private FileScanTaskWithDomain createFileScanTaskWithDomain(FileScanTask wholeFileTask) + { + List predicatedColumns = wholeFileTask.schema().columns().stream() + .filter(column -> predicatedColumnIds.contains(column.fieldId())) + .map(column -> getColumnHandle(column, typeManager)) + .collect(toImmutableList()); + return new FileScanTaskWithDomain( + wholeFileTask, + createFileStatisticsDomain( + fieldIdToType, + wholeFileTask.file().lowerBounds(), + wholeFileTask.file().upperBounds(), + wholeFileTask.file().nullValueCounts(), + predicatedColumns)); + } + + private record FileScanTaskWithDomain(FileScanTask fileScanTask, TupleDomain fileStatisticsDomain) + { + Iterator split(long targetSplitSize) + { + return Iterators.transform( + fileScanTask().split(targetSplitSize).iterator(), + task -> new FileScanTaskWithDomain(task, fileStatisticsDomain)); } } @VisibleForTesting - static boolean fileMatchesPredicate( - Map primitiveTypeForFieldId, - TupleDomain dynamicFilterPredicate, + static TupleDomain createFileStatisticsDomain( + Map fieldIdToType, @Nullable Map lowerBounds, @Nullable Map upperBounds, - @Nullable Map nullValueCounts) + @Nullable Map nullValueCounts, + List predicatedColumns) { - if (dynamicFilterPredicate.isNone()) { - return false; - } - Map domains = dynamicFilterPredicate.getDomains().orElseThrow(); - - for (Map.Entry domainEntry : domains.entrySet()) { - IcebergColumnHandle column = domainEntry.getKey(); - Domain domain = domainEntry.getValue(); - + ImmutableMap.Builder domainBuilder = ImmutableMap.builder(); + for (IcebergColumnHandle column : predicatedColumns) { int fieldId = column.getId(); boolean mayContainNulls; if (nullValueCounts == null) { @@ -376,17 +571,16 @@ static boolean fileMatchesPredicate( Long nullValueCount = nullValueCounts.get(fieldId); mayContainNulls = nullValueCount == null || nullValueCount > 0; } - Type type = primitiveTypeForFieldId.get(fieldId); - Domain statisticsDomain = domainForStatistics( + Type type = fieldIdToType.get(fieldId); + domainBuilder.put( column, - lowerBounds == null ? null : fromByteBuffer(type, lowerBounds.get(fieldId)), - upperBounds == null ? null : fromByteBuffer(type, upperBounds.get(fieldId)), - mayContainNulls); - if (!domain.overlaps(statisticsDomain)) { - return false; - } + domainForStatistics( + column, + lowerBounds == null ? null : fromByteBuffer(type, lowerBounds.get(fieldId)), + upperBounds == null ? null : fromByteBuffer(type, upperBounds.get(fieldId)), + mayContainNulls)); } - return true; + return TupleDomain.withColumnDomains(domainBuilder.buildOrThrow()); } private static Domain domainForStatistics( @@ -419,19 +613,38 @@ else if (upperBound != null) { return Domain.create(ValueSet.ofRanges(statisticsRange), mayContainNulls); } - static boolean partitionMatchesConstraint( - Set identityPartitionColumns, - Supplier> partitionValues, - Constraint constraint) + private static class PartitionConstraintMatcher { - // We use Constraint just to pass functional predicate here from DistributedExecutionPlanner - verify(constraint.getSummary().isAll()); + private final NonEvictableCache, Boolean> partitionConstraintResults; + private final Optional>> predicate; + private final Optional> predicateColumns; + + private PartitionConstraintMatcher(Constraint constraint) + { + // We use Constraint just to pass functional predicate here from DistributedExecutionPlanner + verify(constraint.getSummary().isAll()); + this.predicate = constraint.predicate(); + this.predicateColumns = constraint.getPredicateColumns(); + this.partitionConstraintResults = buildNonEvictableCache(CacheBuilder.newBuilder().maximumSize(1000)); + } - if (constraint.predicate().isEmpty() || - intersection(constraint.getPredicateColumns().orElseThrow(), identityPartitionColumns).isEmpty()) { - return true; + boolean matches( + Set identityPartitionColumns, + Supplier> partitionValuesSupplier) + { + if (predicate.isEmpty()) { + return true; + } + Set predicatePartitionColumns = intersection(predicateColumns.orElseThrow(), identityPartitionColumns); + if (predicatePartitionColumns.isEmpty()) { + return true; + } + Map partitionValues = partitionValuesSupplier.get(); + return uncheckedCacheGet( + partitionConstraintResults, + ImmutableMap.copyOf(Maps.filterKeys(partitionValues, predicatePartitionColumns::contains)), + () -> predicate.orElseThrow().test(partitionValues)); } - return constraint.predicate().get().test(partitionValues.get()); } @VisibleForTesting @@ -456,41 +669,62 @@ static boolean partitionMatchesPredicate( return true; } - private IcebergSplit toIcebergSplit(FileScanTask task) + private IcebergSplit toIcebergSplit(FileScanTaskWithDomain taskWithDomain) { + FileScanTask task = taskWithDomain.fileScanTask(); + Optional> partitionValues = Optional.empty(); + if (tableHandle.getTablePartitioning().isPresent()) { + PartitionSpec partitionSpec = task.spec(); + StructLike partition = task.file().partition(); + List fields = partitionSpec.fields(); + + partitionValues = Optional.of(tableHandle.getTablePartitioning().get().partitionStructFields().stream() + .map(fieldIndex -> convertIcebergValueToTrino( + partitionSpec.partitionType().field(fields.get(fieldIndex).fieldId()).type(), + partition.get(fieldIndex, Object.class))) + .toList()); + } + return new IcebergSplit( - task.file().path().toString(), + task.file().location(), task.start(), task.length(), task.file().fileSizeInBytes(), + task.file().recordCount(), IcebergFileFormat.fromIceberg(task.file().format()), + partitionValues, PartitionSpecParser.toJson(task.spec()), PartitionData.toJson(task.file().partition()), task.deletes().stream() .map(DeleteFile::fromIceberg) .collect(toImmutableList()), - SplitWeight.fromProportion(Math.min(Math.max((double) task.length() / tableScan.targetSplitSize(), minimumAssignedSplitWeight), 1.0))); + SplitWeight.fromProportion(clamp(getSplitWeight(task), minimumAssignedSplitWeight, 1.0)), + taskWithDomain.fileStatisticsDomain(), + fileIoProperties, + cachingHostAddressProvider.getHosts(task.file().location(), ImmutableList.of()), + task.file().dataSequenceNumber()); } - private static Domain getPathDomain(TupleDomain effectivePredicate) + private static double clamp(double value, double min, double max) { - IcebergColumnHandle pathColumn = pathColumnHandle(); - Domain domain = effectivePredicate.getDomains().orElseThrow(() -> new IllegalArgumentException("Unexpected NONE tuple domain")) - .get(pathColumn); - if (domain == null) { - return Domain.all(pathColumn.getType()); - } - return domain; + return Math.min(max, Math.max(value, min)); } - private static Domain getFileModifiedTimePathDomain(TupleDomain effectivePredicate) + private double getSplitWeight(FileScanTask task) { - IcebergColumnHandle fileModifiedTimeColumn = fileModifiedTimeColumnHandle(); - Domain domain = effectivePredicate.getDomains().orElseThrow(() -> new IllegalArgumentException("Unexpected NONE tuple domain")) - .get(fileModifiedTimeColumn); - if (domain == null) { - return Domain.all(fileModifiedTimeColumn.getType()); + double dataWeight = (double) task.length() / tableScan.targetSplitSize(); + double weight = dataWeight; + if (task.deletes().stream().anyMatch(deleteFile -> deleteFile.content() == POSITION_DELETES)) { + // Presence of each data position is looked up in a combined bitmap of deleted positions + weight += dataWeight; } - return domain; + + long equalityDeletes = task.deletes().stream() + .filter(deleteFile -> deleteFile.content() == EQUALITY_DELETES) + .mapToLong(ContentFile::recordCount) + .sum(); + // Every row is a separate equality predicate that must be applied to all data rows + weight += equalityDeletes * dataWeight; + return weight; } } diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergStatistics.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergStatistics.java index 8b3c80053da5..e15b0424af15 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergStatistics.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergStatistics.java @@ -13,9 +13,7 @@ */ package io.trino.plugin.iceberg; -import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; -import com.google.errorprone.annotations.Immutable; import io.trino.spi.TrinoException; import io.trino.spi.type.TypeManager; import jakarta.annotation.Nullable; @@ -45,99 +43,43 @@ import static io.trino.spi.function.InvocationConvention.simpleConvention; import static java.util.Objects.requireNonNull; -@Immutable -final class IcebergStatistics +record IcebergStatistics( + long recordCount, + long fileCount, + long size, + Map minValues, + Map maxValues, + Map nullCounts, + Map nanCounts, + Map columnSizes) { - private final long recordCount; - private final long fileCount; - private final long size; - private final Map minValues; - private final Map maxValues; - private final Map nullCounts; - private final Map nanCounts; - private final Map columnSizes; - - private IcebergStatistics( - long recordCount, - long fileCount, - long size, - Map minValues, - Map maxValues, - Map nullCounts, - Map nanCounts, - Map columnSizes) - { - this.recordCount = recordCount; - this.fileCount = fileCount; - this.size = size; - this.minValues = ImmutableMap.copyOf(requireNonNull(minValues, "minValues is null")); - this.maxValues = ImmutableMap.copyOf(requireNonNull(maxValues, "maxValues is null")); - this.nullCounts = ImmutableMap.copyOf(requireNonNull(nullCounts, "nullCounts is null")); - this.nanCounts = ImmutableMap.copyOf(requireNonNull(nanCounts, "nanCounts is null")); - this.columnSizes = ImmutableMap.copyOf(requireNonNull(columnSizes, "columnSizes is null")); - } - - public long getRecordCount() - { - return recordCount; - } - - public long getFileCount() - { - return fileCount; - } - - public long getSize() - { - return size; - } - - public Map getMinValues() - { - return minValues; - } - - public Map getMaxValues() - { - return maxValues; - } - - public Map getNullCounts() - { - return nullCounts; - } - - public Map getNanCounts() + IcebergStatistics { - return nanCounts; + minValues = ImmutableMap.copyOf(requireNonNull(minValues, "minValues is null")); + maxValues = ImmutableMap.copyOf(requireNonNull(maxValues, "maxValues is null")); + nullCounts = ImmutableMap.copyOf(requireNonNull(nullCounts, "nullCounts is null")); + nanCounts = ImmutableMap.copyOf(requireNonNull(nanCounts, "nanCounts is null")); + columnSizes = ImmutableMap.copyOf(requireNonNull(columnSizes, "columnSizes is null")); } - public Map getColumnSizes() + static class Builder { - return columnSizes; - } - - public static class Builder - { - private final List columns; private final TypeManager typeManager; - private final Map> nullCounts = new HashMap<>(); - private final Map> nanCounts = new HashMap<>(); - private final Map columnStatistics = new HashMap<>(); - private final Map columnSizes = new HashMap<>(); private final Map fieldIdToTrinoType; private long recordCount; private long fileCount; private long size; + private final Map columnStatistics = new HashMap<>(); + private final Map> nullCounts = new HashMap<>(); + private final Map> nanCounts = new HashMap<>(); + private final Map columnSizes = new HashMap<>(); public Builder( List columns, TypeManager typeManager) { - this.columns = ImmutableList.copyOf(requireNonNull(columns, "columns is null")); this.typeManager = requireNonNull(typeManager, "typeManager is null"); - this.fieldIdToTrinoType = columns.stream() .collect(toImmutableMap(Types.NestedField::fieldId, column -> toTrinoType(column.type(), typeManager))); } @@ -150,11 +92,10 @@ public void acceptDataFile(DataFile dataFile, PartitionSpec partitionSpec) Map newColumnSizes = dataFile.columnSizes(); if (newColumnSizes != null) { - for (Types.NestedField column : columns) { - int id = column.fieldId(); - Long addedSize = newColumnSizes.get(id); + for (Map.Entry entry : newColumnSizes.entrySet()) { + Long addedSize = entry.getValue(); if (addedSize != null) { - columnSizes.merge(id, addedSize, Long::sum); + columnSizes.merge(entry.getKey(), addedSize, Long::sum); } } } @@ -262,7 +203,7 @@ private void updateMinMaxStats( if (type.isOrderable() && (nullCount.isEmpty() || nullCount.get() != recordCount)) { // Capture the initial bounds during construction so there are always valid min/max values to compare to. This does make the first call to // `ColumnStatistics#updateMinMax` a no-op. - columnStatistics.computeIfAbsent(id, ignored -> { + columnStatistics.computeIfAbsent(id, ignore -> { MethodHandle comparisonHandle = typeManager.getTypeOperators() .getComparisonUnorderedLastOperator(type, simpleConvention(FAIL_ON_NULL, NEVER_NULL, NEVER_NULL)); return new ColumnStatistics(comparisonHandle, lowerBound, upperBound); diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergTableHandle.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergTableHandle.java index b3edb6165935..21e1e0535e1e 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergTableHandle.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergTableHandle.java @@ -18,6 +18,7 @@ import com.fasterxml.jackson.annotation.JsonProperty; import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; +import com.google.errorprone.annotations.DoNotCall; import io.airlift.units.DataSize; import io.trino.spi.connector.CatalogHandle; import io.trino.spi.connector.ConnectorTableHandle; @@ -55,17 +56,27 @@ public class IcebergTableHandle // Filter guaranteed to be enforced by Iceberg connector private final TupleDomain enforcedPredicate; + // Columns that are present in {@link Constraint#predicate()} applied on the table scan + private final Set constraintColumns; + // semantically limit is applied after enforcedPredicate private final OptionalLong limit; private final Set projectedColumns; private final Optional nameMappingJson; + // Coordinator-only - table partitioning applied to the table splits if available and active + private final Optional tablePartitioning; + // OPTIMIZE only. Coordinator-only private final boolean recordScannedFiles; private final Optional maxScannedFileSize; + // ANALYZE only. Coordinator-only + private final Optional forAnalyze; + @JsonCreator + @DoNotCall // For JSON deserialization only public static IcebergTableHandle fromJsonForDeserializationOnly( @JsonProperty("catalog") CatalogHandle catalog, @JsonProperty("schemaName") String schemaName, @@ -99,7 +110,10 @@ public static IcebergTableHandle fromJsonForDeserializationOnly( nameMappingJson, tableLocation, storageProperties, + Optional.empty(), false, + Optional.empty(), + ImmutableSet.of(), Optional.empty()); } @@ -119,8 +133,11 @@ public IcebergTableHandle( Optional nameMappingJson, String tableLocation, Map storageProperties, + Optional tablePartitioning, boolean recordScannedFiles, - Optional maxScannedFileSize) + Optional maxScannedFileSize, + Set constraintColumns, + Optional forAnalyze) { this.catalog = requireNonNull(catalog, "catalog is null"); this.schemaName = requireNonNull(schemaName, "schemaName is null"); @@ -137,8 +154,11 @@ public IcebergTableHandle( this.nameMappingJson = requireNonNull(nameMappingJson, "nameMappingJson is null"); this.tableLocation = requireNonNull(tableLocation, "tableLocation is null"); this.storageProperties = ImmutableMap.copyOf(requireNonNull(storageProperties, "storageProperties is null")); + this.tablePartitioning = requireNonNull(tablePartitioning, "tablePartitioning is null"); this.recordScannedFiles = recordScannedFiles; this.maxScannedFileSize = requireNonNull(maxScannedFileSize, "maxScannedFileSize is null"); + this.constraintColumns = ImmutableSet.copyOf(requireNonNull(constraintColumns, "constraintColumns is null")); + this.forAnalyze = requireNonNull(forAnalyze, "forAnalyze is null"); } @JsonProperty @@ -232,6 +252,15 @@ public Map getStorageProperties() return storageProperties; } + /** + * Get the partitioning for the table splits. + */ + @JsonIgnore + public Optional getTablePartitioning() + { + return tablePartitioning; + } + @JsonIgnore public boolean isRecordScannedFiles() { @@ -244,6 +273,18 @@ public Optional getMaxScannedFileSize() return maxScannedFileSize; } + @JsonIgnore + public Set getConstraintColumns() + { + return constraintColumns; + } + + @JsonIgnore + public Optional getForAnalyze() + { + return forAnalyze; + } + public SchemaTableName getSchemaTableName() { return new SchemaTableName(schemaName, tableName); @@ -272,8 +313,36 @@ public IcebergTableHandle withProjectedColumns(Set projecte nameMappingJson, tableLocation, storageProperties, + tablePartitioning, recordScannedFiles, - maxScannedFileSize); + maxScannedFileSize, + constraintColumns, + forAnalyze); + } + + public IcebergTableHandle forAnalyze() + { + return new IcebergTableHandle( + catalog, + schemaName, + tableName, + tableType, + snapshotId, + tableSchemaJson, + partitionSpecJson, + formatVersion, + unenforcedPredicate, + enforcedPredicate, + limit, + projectedColumns, + nameMappingJson, + tableLocation, + storageProperties, + tablePartitioning, + recordScannedFiles, + maxScannedFileSize, + constraintColumns, + Optional.of(true)); } public IcebergTableHandle forOptimize(boolean recordScannedFiles, DataSize maxScannedFileSize) @@ -294,8 +363,36 @@ public IcebergTableHandle forOptimize(boolean recordScannedFiles, DataSize maxSc nameMappingJson, tableLocation, storageProperties, + tablePartitioning, + recordScannedFiles, + Optional.of(maxScannedFileSize), + constraintColumns, + forAnalyze); + } + + public IcebergTableHandle withTablePartitioning(Optional requiredTablePartitioning) + { + return new IcebergTableHandle( + catalog, + schemaName, + tableName, + tableType, + snapshotId, + tableSchemaJson, + partitionSpecJson, + formatVersion, + unenforcedPredicate, + enforcedPredicate, + limit, + projectedColumns, + nameMappingJson, + tableLocation, + storageProperties, + requiredTablePartitioning, recordScannedFiles, - Optional.of(maxScannedFileSize)); + maxScannedFileSize, + constraintColumns, + forAnalyze); } @Override @@ -325,7 +422,9 @@ public boolean equals(Object o) Objects.equals(nameMappingJson, that.nameMappingJson) && Objects.equals(tableLocation, that.tableLocation) && Objects.equals(storageProperties, that.storageProperties) && - Objects.equals(maxScannedFileSize, that.maxScannedFileSize); + Objects.equals(maxScannedFileSize, that.maxScannedFileSize) && + Objects.equals(constraintColumns, that.constraintColumns) && + Objects.equals(forAnalyze, that.forAnalyze); } @Override @@ -348,7 +447,9 @@ public int hashCode() tableLocation, storageProperties, recordScannedFiles, - maxScannedFileSize); + maxScannedFileSize, + constraintColumns, + forAnalyze); } @Override diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergTableName.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergTableName.java index 32716889471d..63790f48bb05 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergTableName.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergTableName.java @@ -13,14 +13,15 @@ */ package io.trino.plugin.iceberg; -import io.trino.spi.TrinoException; - -import java.util.Optional; import java.util.regex.Matcher; import java.util.regex.Pattern; +import java.util.stream.Collectors; +import java.util.stream.Stream; +import static com.google.common.base.Preconditions.checkArgument; +import static com.google.common.base.Verify.verify; import static io.trino.plugin.iceberg.TableType.DATA; -import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED; +import static io.trino.plugin.iceberg.TableType.MATERIALIZED_VIEW_STORAGE; import static java.util.Locale.ENGLISH; import static java.util.Objects.requireNonNull; @@ -28,9 +29,22 @@ public final class IcebergTableName { private IcebergTableName() {} - private static final Pattern TABLE_PATTERN = Pattern.compile("" + - "(?
[^$@]+)" + - "(?:\\$(?[^@]+))?"); + private static final Pattern TABLE_PATTERN; + + static { + String referencableTableTypes = Stream.of(TableType.values()) + .filter(tableType -> tableType != DATA) + .map(tableType -> tableType.name().toLowerCase(ENGLISH)) + .collect(Collectors.joining("|")); + TABLE_PATTERN = Pattern.compile("" + + "(?
[^$@]+)" + + "(?:\\$(?(?i:" + referencableTableTypes + ")))?"); + } + + public static boolean isIcebergTableName(String tableName) + { + return TABLE_PATTERN.matcher(tableName).matches(); + } public static String tableNameWithType(String tableName, TableType tableType) { @@ -38,46 +52,38 @@ public static String tableNameWithType(String tableName, TableType tableType) return tableName + "$" + tableType.name().toLowerCase(ENGLISH); } - public static String tableNameFrom(String name) + public static String tableNameFrom(String validIcebergTableName) { - Matcher match = TABLE_PATTERN.matcher(name); - if (!match.matches()) { - throw new TrinoException(NOT_SUPPORTED, "Invalid Iceberg table name: " + name); - } - + Matcher match = TABLE_PATTERN.matcher(validIcebergTableName); + checkArgument(match.matches(), "Invalid Iceberg table name: %s", validIcebergTableName); return match.group("table"); } - public static Optional tableTypeFrom(String name) + public static TableType tableTypeFrom(String validIcebergTableName) { - Matcher match = TABLE_PATTERN.matcher(name); - if (!match.matches()) { - throw new TrinoException(NOT_SUPPORTED, "Invalid Iceberg table name: " + name); - } + Matcher match = TABLE_PATTERN.matcher(validIcebergTableName); + checkArgument(match.matches(), "Invalid Iceberg table name: %s", validIcebergTableName); + String typeString = match.group("type"); if (typeString == null) { - return Optional.of(DATA); - } - try { - TableType parsedType = TableType.valueOf(typeString.toUpperCase(ENGLISH)); - if (parsedType == DATA) { - // $data cannot be encoded in table name - return Optional.empty(); - } - return Optional.of(parsedType); - } - catch (IllegalArgumentException e) { - return Optional.empty(); + return DATA; } + TableType parsedType = TableType.valueOf(typeString.toUpperCase(ENGLISH)); + // $data cannot be encoded in table name + verify(parsedType != DATA, "parsedType is unexpectedly DATA"); + return parsedType; } - public static boolean isDataTable(String name) + public static boolean isDataTable(String validIcebergTableName) { - Matcher match = TABLE_PATTERN.matcher(name); - if (!match.matches()) { - throw new TrinoException(NOT_SUPPORTED, "Invalid Iceberg table name: " + name); - } + Matcher match = TABLE_PATTERN.matcher(validIcebergTableName); + checkArgument(match.matches(), "Invalid Iceberg table name: %s", validIcebergTableName); String typeString = match.group("type"); return typeString == null; } + + public static boolean isMaterializedViewStorage(String validIcebergTableName) + { + return tableTypeFrom(validIcebergTableName) == MATERIALIZED_VIEW_STORAGE; + } } diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergTablePartitioning.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergTablePartitioning.java new file mode 100644 index 000000000000..5d5562e50c49 --- /dev/null +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergTablePartitioning.java @@ -0,0 +1,48 @@ +/* + * 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.iceberg; + +import com.google.common.collect.ImmutableList; +import io.trino.spi.connector.ConnectorTablePartitioning; + +import java.util.List; +import java.util.Optional; + +import static com.google.common.base.Preconditions.checkArgument; +import static java.util.Objects.requireNonNull; + +public record IcebergTablePartitioning( + boolean active, + IcebergPartitioningHandle partitioningHandle, + List partitioningColumns, + List partitionStructFields) +{ + public IcebergTablePartitioning + { + requireNonNull(partitioningHandle, "partitioningHandle is null"); + partitioningColumns = ImmutableList.copyOf(requireNonNull(partitioningColumns, "partitioningColumns is null")); + partitionStructFields = ImmutableList.copyOf(requireNonNull(partitionStructFields, "partitionStructFields is null")); + checkArgument(partitioningHandle.partitionFunctions().size() == partitionStructFields.size(), "partitioningColumns and partitionStructFields must have the same size"); + } + + public IcebergTablePartitioning activate() + { + return new IcebergTablePartitioning(true, partitioningHandle, partitioningColumns, partitionStructFields); + } + + public Optional toConnectorTablePartitioning() + { + return active ? Optional.of(new ConnectorTablePartitioning(partitioningHandle, ImmutableList.copyOf(partitioningColumns))) : Optional.empty(); + } +} diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergTableProperties.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergTableProperties.java index 9de78f2280e1..91f7a1c73eba 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergTableProperties.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergTableProperties.java @@ -14,20 +14,27 @@ package io.trino.plugin.iceberg; import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableSet; import com.google.inject.Inject; import io.trino.plugin.hive.orc.OrcWriterConfig; import io.trino.spi.TrinoException; import io.trino.spi.session.PropertyMetadata; import io.trino.spi.type.ArrayType; +import io.trino.spi.type.MapType; +import io.trino.spi.type.TypeManager; import java.util.List; import java.util.Map; import java.util.Optional; +import java.util.Set; import static com.google.common.collect.ImmutableList.toImmutableList; +import static com.google.common.collect.ImmutableMap.toImmutableMap; import static io.trino.plugin.iceberg.IcebergConfig.FORMAT_VERSION_SUPPORT_MAX; import static io.trino.plugin.iceberg.IcebergConfig.FORMAT_VERSION_SUPPORT_MIN; import static io.trino.spi.StandardErrorCode.INVALID_TABLE_PROPERTY; +import static io.trino.spi.session.PropertyMetadata.booleanProperty; import static io.trino.spi.session.PropertyMetadata.doubleProperty; import static io.trino.spi.session.PropertyMetadata.enumProperty; import static io.trino.spi.session.PropertyMetadata.integerProperty; @@ -35,6 +42,12 @@ import static io.trino.spi.type.VarcharType.VARCHAR; import static java.lang.String.format; import static java.util.Locale.ENGLISH; +import static org.apache.iceberg.TableProperties.COMMIT_NUM_RETRIES_DEFAULT; +import static org.apache.iceberg.TableProperties.DEFAULT_FILE_FORMAT; +import static org.apache.iceberg.TableProperties.FORMAT_VERSION; +import static org.apache.iceberg.TableProperties.ORC_BLOOM_FILTER_COLUMNS; +import static org.apache.iceberg.TableProperties.ORC_BLOOM_FILTER_FPP; +import static org.apache.iceberg.TableProperties.RESERVED_PROPERTIES; public class IcebergTableProperties { @@ -43,15 +56,45 @@ public class IcebergTableProperties public static final String SORTED_BY_PROPERTY = "sorted_by"; public static final String LOCATION_PROPERTY = "location"; public static final String FORMAT_VERSION_PROPERTY = "format_version"; - public static final String ORC_BLOOM_FILTER_COLUMNS = "orc_bloom_filter_columns"; - public static final String ORC_BLOOM_FILTER_FPP = "orc_bloom_filter_fpp"; + public static final String MAX_COMMIT_RETRY = "max_commit_retry"; + public static final String ORC_BLOOM_FILTER_COLUMNS_PROPERTY = "orc_bloom_filter_columns"; + public static final String ORC_BLOOM_FILTER_FPP_PROPERTY = "orc_bloom_filter_fpp"; + public static final String PARQUET_BLOOM_FILTER_COLUMNS_PROPERTY = "parquet_bloom_filter_columns"; + public static final String OBJECT_STORE_LAYOUT_ENABLED_PROPERTY = "object_store_layout_enabled"; + public static final String DATA_LOCATION_PROPERTY = "data_location"; + public static final String EXTRA_PROPERTIES_PROPERTY = "extra_properties"; + + public static final Set SUPPORTED_PROPERTIES = ImmutableSet.builder() + .add(FILE_FORMAT_PROPERTY) + .add(PARTITIONING_PROPERTY) + .add(SORTED_BY_PROPERTY) + .add(LOCATION_PROPERTY) + .add(FORMAT_VERSION_PROPERTY) + .add(MAX_COMMIT_RETRY) + .add(ORC_BLOOM_FILTER_COLUMNS_PROPERTY) + .add(ORC_BLOOM_FILTER_FPP_PROPERTY) + .add(OBJECT_STORE_LAYOUT_ENABLED_PROPERTY) + .add(DATA_LOCATION_PROPERTY) + .add(EXTRA_PROPERTIES_PROPERTY) + .add(PARQUET_BLOOM_FILTER_COLUMNS_PROPERTY) + .build(); + + // These properties are used by Trino or Iceberg internally and cannot be set directly by users through extra_properties + public static final Set PROTECTED_ICEBERG_NATIVE_PROPERTIES = ImmutableSet.builder() + .addAll(RESERVED_PROPERTIES) + .add(ORC_BLOOM_FILTER_COLUMNS) + .add(ORC_BLOOM_FILTER_FPP) + .add(DEFAULT_FILE_FORMAT) + .add(FORMAT_VERSION) + .build(); private final List> tableProperties; @Inject public IcebergTableProperties( IcebergConfig icebergConfig, - OrcWriterConfig orcWriterConfig) + OrcWriterConfig orcWriterConfig, + TypeManager typeManager) { tableProperties = ImmutableList.>builder() .add(enumProperty( @@ -89,8 +132,18 @@ public IcebergTableProperties( icebergConfig.getFormatVersion(), IcebergTableProperties::validateFormatVersion, false)) + .add(integerProperty( + MAX_COMMIT_RETRY, + "Number of times to retry a commit before failing", + icebergConfig.getMaxCommitRetry(), + value -> { + if (value < 0) { + throw new TrinoException(INVALID_TABLE_PROPERTY, "max_commit_retry must be greater than or equal to 0"); + } + }, + false)) .add(new PropertyMetadata<>( - ORC_BLOOM_FILTER_COLUMNS, + ORC_BLOOM_FILTER_COLUMNS_PROPERTY, "ORC Bloom filter index columns", new ArrayType(VARCHAR), List.class, @@ -102,11 +155,53 @@ public IcebergTableProperties( .collect(toImmutableList()), value -> value)) .add(doubleProperty( - ORC_BLOOM_FILTER_FPP, + ORC_BLOOM_FILTER_FPP_PROPERTY, "ORC Bloom filter false positive probability", orcWriterConfig.getDefaultBloomFilterFpp(), IcebergTableProperties::validateOrcBloomFilterFpp, false)) + .add(new PropertyMetadata<>( + PARQUET_BLOOM_FILTER_COLUMNS_PROPERTY, + "Parquet Bloom filter index columns", + new ArrayType(VARCHAR), + List.class, + ImmutableList.of(), + false, + value -> ((List) value).stream() + .map(String.class::cast) + .map(name -> name.toLowerCase(ENGLISH)) + .collect(toImmutableList()), + value -> value)) + .add(new PropertyMetadata<>( + EXTRA_PROPERTIES_PROPERTY, + "Extra table properties", + new MapType(VARCHAR, VARCHAR, typeManager.getTypeOperators()), + Map.class, + ImmutableMap.of(), + true, // currently not shown in SHOW CREATE TABLE + value -> { + Map extraProperties = (Map) value; + if (extraProperties.containsValue(null)) { + throw new TrinoException(INVALID_TABLE_PROPERTY, format("Extra table property value cannot be null '%s'", extraProperties)); + } + if (extraProperties.containsKey(null)) { + throw new TrinoException(INVALID_TABLE_PROPERTY, format("Extra table property key cannot be null '%s'", extraProperties)); + } + + return extraProperties.entrySet().stream() + .collect(toImmutableMap(entry -> entry.getKey().toLowerCase(ENGLISH), Map.Entry::getValue)); + }, + value -> value)) + .add(booleanProperty( + OBJECT_STORE_LAYOUT_ENABLED_PROPERTY, + "Set to true to enable Iceberg object store file layout", + icebergConfig.isObjectStoreLayoutEnabled(), + false)) + .add(stringProperty( + DATA_LOCATION_PROPERTY, + "File system location URI for the table's data files", + null, + false)) .build(); } @@ -144,6 +239,11 @@ public static int getFormatVersion(Map tableProperties) return (int) tableProperties.get(FORMAT_VERSION_PROPERTY); } + public static int getMaxCommitRetry(Map tableProperties) + { + return (int) tableProperties.getOrDefault(MAX_COMMIT_RETRY, COMMIT_NUM_RETRIES_DEFAULT); + } + private static void validateFormatVersion(int version) { if (version < FORMAT_VERSION_SUPPORT_MIN || version > FORMAT_VERSION_SUPPORT_MAX) { @@ -154,13 +254,13 @@ private static void validateFormatVersion(int version) public static List getOrcBloomFilterColumns(Map tableProperties) { - List orcBloomFilterColumns = (List) tableProperties.get(ORC_BLOOM_FILTER_COLUMNS); + List orcBloomFilterColumns = (List) tableProperties.get(ORC_BLOOM_FILTER_COLUMNS_PROPERTY); return orcBloomFilterColumns == null ? ImmutableList.of() : ImmutableList.copyOf(orcBloomFilterColumns); } public static Double getOrcBloomFilterFpp(Map tableProperties) { - return (Double) tableProperties.get(ORC_BLOOM_FILTER_FPP); + return (Double) tableProperties.get(ORC_BLOOM_FILTER_FPP_PROPERTY); } private static void validateOrcBloomFilterFpp(double fpp) @@ -169,4 +269,25 @@ private static void validateOrcBloomFilterFpp(double fpp) throw new TrinoException(INVALID_TABLE_PROPERTY, "Bloom filter fpp value must be between 0.0 and 1.0"); } } + + public static List getParquetBloomFilterColumns(Map tableProperties) + { + List parquetBloomFilterColumns = (List) tableProperties.get(PARQUET_BLOOM_FILTER_COLUMNS_PROPERTY); + return parquetBloomFilterColumns == null ? ImmutableList.of() : ImmutableList.copyOf(parquetBloomFilterColumns); + } + + public static boolean getObjectStoreLayoutEnabled(Map tableProperties) + { + return (boolean) tableProperties.getOrDefault(OBJECT_STORE_LAYOUT_ENABLED_PROPERTY, false); + } + + public static Optional getDataLocation(Map tableProperties) + { + return Optional.ofNullable((String) tableProperties.get(DATA_LOCATION_PROPERTY)); + } + + public static Optional> getExtraProperties(Map tableProperties) + { + return Optional.ofNullable((Map) tableProperties.get(EXTRA_PROPERTIES_PROPERTY)); + } } diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergUtil.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergUtil.java index 78b6c1a254f4..26f824fda19e 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergUtil.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergUtil.java @@ -18,25 +18,36 @@ import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; +import com.google.common.collect.Lists; import com.google.common.collect.Sets; import io.airlift.slice.Slice; import io.airlift.slice.SliceUtf8; import io.airlift.slice.Slices; +import io.trino.filesystem.Location; +import io.trino.filesystem.TrinoFileSystem; +import io.trino.filesystem.TrinoInputFile; import io.trino.plugin.iceberg.PartitionTransforms.ColumnTransform; import io.trino.plugin.iceberg.catalog.IcebergTableOperations; import io.trino.plugin.iceberg.catalog.IcebergTableOperationsProvider; import io.trino.plugin.iceberg.catalog.TrinoCatalog; +import io.trino.plugin.iceberg.util.DefaultLocationProvider; +import io.trino.plugin.iceberg.util.ObjectStoreLocationProvider; import io.trino.spi.TrinoException; +import io.trino.spi.connector.ColumnHandle; import io.trino.spi.connector.ColumnMetadata; import io.trino.spi.connector.ConnectorSession; import io.trino.spi.connector.ConnectorTableMetadata; +import io.trino.spi.connector.ConnectorViewDefinition.ViewColumn; import io.trino.spi.connector.SchemaTableName; import io.trino.spi.function.InvocationConvention; import io.trino.spi.predicate.Domain; +import io.trino.spi.predicate.NullableValue; import io.trino.spi.predicate.Range; +import io.trino.spi.predicate.TupleDomain; import io.trino.spi.predicate.ValueSet; import io.trino.spi.type.DecimalType; import io.trino.spi.type.Int128; +import io.trino.spi.type.RowType; import io.trino.spi.type.Type; import io.trino.spi.type.TypeManager; import io.trino.spi.type.TypeOperators; @@ -62,9 +73,12 @@ import org.apache.iceberg.Transaction; import org.apache.iceberg.io.LocationProvider; import org.apache.iceberg.types.Type.PrimitiveType; +import org.apache.iceberg.types.TypeUtil; +import org.apache.iceberg.types.Types; import org.apache.iceberg.types.Types.NestedField; import org.apache.iceberg.types.Types.StructType; +import java.io.IOException; import java.lang.invoke.MethodHandle; import java.math.BigDecimal; import java.math.BigInteger; @@ -76,43 +90,52 @@ import java.util.Locale; import java.util.Map; import java.util.Map.Entry; +import java.util.Objects; import java.util.Optional; import java.util.Set; import java.util.UUID; import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.Predicate; import java.util.regex.Matcher; import java.util.regex.Pattern; import java.util.stream.Stream; import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.base.Verify.verify; +import static com.google.common.collect.ImmutableList.builderWithExpectedSize; import static com.google.common.collect.ImmutableList.toImmutableList; import static com.google.common.collect.ImmutableMap.toImmutableMap; import static com.google.common.collect.ImmutableSet.toImmutableSet; import static com.google.common.collect.Maps.immutableEntry; import static com.google.common.collect.Streams.mapWithIndex; import static io.airlift.slice.Slices.utf8Slice; +import static io.trino.parquet.writer.ParquetWriter.SUPPORTED_BLOOM_FILTER_TYPES; import static io.trino.plugin.base.io.ByteBuffers.getWrappedBytes; -import static io.trino.plugin.hive.HiveMetadata.TABLE_COMMENT; +import static io.trino.plugin.hive.metastore.Table.TABLE_COMMENT; import static io.trino.plugin.iceberg.ColumnIdentity.createColumnIdentity; +import static io.trino.plugin.iceberg.IcebergColumnHandle.fileModifiedTimeColumnHandle; import static io.trino.plugin.iceberg.IcebergColumnHandle.fileModifiedTimeColumnMetadata; +import static io.trino.plugin.iceberg.IcebergColumnHandle.partitionColumnHandle; +import static io.trino.plugin.iceberg.IcebergColumnHandle.pathColumnHandle; import static io.trino.plugin.iceberg.IcebergColumnHandle.pathColumnMetadata; import static io.trino.plugin.iceberg.IcebergErrorCode.ICEBERG_BAD_DATA; +import static io.trino.plugin.iceberg.IcebergErrorCode.ICEBERG_FILESYSTEM_ERROR; import static io.trino.plugin.iceberg.IcebergErrorCode.ICEBERG_INVALID_PARTITION_VALUE; -import static io.trino.plugin.iceberg.IcebergMetadata.ORC_BLOOM_FILTER_COLUMNS_KEY; -import static io.trino.plugin.iceberg.IcebergMetadata.ORC_BLOOM_FILTER_FPP_KEY; +import static io.trino.plugin.iceberg.IcebergTableProperties.DATA_LOCATION_PROPERTY; import static io.trino.plugin.iceberg.IcebergTableProperties.FILE_FORMAT_PROPERTY; import static io.trino.plugin.iceberg.IcebergTableProperties.FORMAT_VERSION_PROPERTY; import static io.trino.plugin.iceberg.IcebergTableProperties.LOCATION_PROPERTY; -import static io.trino.plugin.iceberg.IcebergTableProperties.ORC_BLOOM_FILTER_COLUMNS; -import static io.trino.plugin.iceberg.IcebergTableProperties.ORC_BLOOM_FILTER_FPP; +import static io.trino.plugin.iceberg.IcebergTableProperties.MAX_COMMIT_RETRY; +import static io.trino.plugin.iceberg.IcebergTableProperties.OBJECT_STORE_LAYOUT_ENABLED_PROPERTY; +import static io.trino.plugin.iceberg.IcebergTableProperties.ORC_BLOOM_FILTER_COLUMNS_PROPERTY; +import static io.trino.plugin.iceberg.IcebergTableProperties.ORC_BLOOM_FILTER_FPP_PROPERTY; +import static io.trino.plugin.iceberg.IcebergTableProperties.PARQUET_BLOOM_FILTER_COLUMNS_PROPERTY; import static io.trino.plugin.iceberg.IcebergTableProperties.PARTITIONING_PROPERTY; +import static io.trino.plugin.iceberg.IcebergTableProperties.PROTECTED_ICEBERG_NATIVE_PROPERTIES; import static io.trino.plugin.iceberg.IcebergTableProperties.SORTED_BY_PROPERTY; -import static io.trino.plugin.iceberg.IcebergTableProperties.getOrcBloomFilterColumns; -import static io.trino.plugin.iceberg.IcebergTableProperties.getOrcBloomFilterFpp; +import static io.trino.plugin.iceberg.IcebergTableProperties.SUPPORTED_PROPERTIES; import static io.trino.plugin.iceberg.IcebergTableProperties.getPartitioning; import static io.trino.plugin.iceberg.IcebergTableProperties.getSortOrder; -import static io.trino.plugin.iceberg.IcebergTableProperties.getTableLocation; import static io.trino.plugin.iceberg.PartitionFields.parsePartitionFields; import static io.trino.plugin.iceberg.PartitionFields.toPartitionFields; import static io.trino.plugin.iceberg.SortFieldUtils.parseSortFields; @@ -140,6 +163,7 @@ import static io.trino.spi.type.TimestampWithTimeZoneType.TIMESTAMP_TZ_MICROS; import static io.trino.spi.type.Timestamps.PICOSECONDS_PER_MICROSECOND; import static io.trino.spi.type.UuidType.javaUuidToTrinoUuid; +import static java.lang.Boolean.parseBoolean; import static java.lang.Double.parseDouble; import static java.lang.Float.floatToRawIntBits; import static java.lang.Float.parseFloat; @@ -149,28 +173,36 @@ import static java.math.RoundingMode.UNNECESSARY; import static java.util.Comparator.comparing; import static java.util.Objects.requireNonNull; -import static org.apache.iceberg.LocationProviders.locationsFor; import static org.apache.iceberg.MetadataTableUtils.createMetadataTableInstance; +import static org.apache.iceberg.TableProperties.COMMIT_NUM_RETRIES; import static org.apache.iceberg.TableProperties.DEFAULT_FILE_FORMAT; import static org.apache.iceberg.TableProperties.DEFAULT_FILE_FORMAT_DEFAULT; import static org.apache.iceberg.TableProperties.FORMAT_VERSION; +import static org.apache.iceberg.TableProperties.OBJECT_STORE_ENABLED; +import static org.apache.iceberg.TableProperties.OBJECT_STORE_ENABLED_DEFAULT; import static org.apache.iceberg.TableProperties.OBJECT_STORE_PATH; +import static org.apache.iceberg.TableProperties.ORC_BLOOM_FILTER_COLUMNS; +import static org.apache.iceberg.TableProperties.ORC_BLOOM_FILTER_FPP; +import static org.apache.iceberg.TableProperties.PARQUET_BLOOM_FILTER_COLUMN_ENABLED_PREFIX; import static org.apache.iceberg.TableProperties.WRITE_DATA_LOCATION; import static org.apache.iceberg.TableProperties.WRITE_LOCATION_PROVIDER_IMPL; import static org.apache.iceberg.TableProperties.WRITE_METADATA_LOCATION; import static org.apache.iceberg.types.Type.TypeID.BINARY; import static org.apache.iceberg.types.Type.TypeID.FIXED; +import static org.apache.iceberg.util.PropertyUtil.propertyAsBoolean; public final class IcebergUtil { public static final String TRINO_TABLE_METADATA_INFO_VALID_FOR = "trino_table_metadata_info_valid_for"; + public static final String TRINO_TABLE_COMMENT_CACHE_PREVENTED = "trino_table_comment_cache_prevented"; public static final String COLUMN_TRINO_NOT_NULL_PROPERTY = "trino_not_null"; public static final String COLUMN_TRINO_TYPE_ID_PROPERTY = "trino_type_id"; public static final String METADATA_FOLDER_NAME = "metadata"; public static final String METADATA_FILE_EXTENSION = ".metadata.json"; private static final Pattern SIMPLE_NAME = Pattern.compile("[a-z][a-z0-9]*"); - static final String TRINO_QUERY_ID_NAME = "trino_query_id"; + public static final String TRINO_QUERY_ID_NAME = "trino_query_id"; + public static final String TRINO_USER_NAME = "trino_user"; // Metadata file name examples // - 00001-409702ba-4735-4645-8f14-09537cc0b2c8.metadata.json // - 00001-409702ba-4735-4645-8f14-09537cc0b2c8.gz.metadata.json (https://github.com/apache/iceberg/blob/ab398a0d5ff195f763f8c7a4358ac98fa38a8de7/core/src/main/java/org/apache/iceberg/TableMetadataParser.java#L141) @@ -184,7 +216,7 @@ public final class IcebergUtil private IcebergUtil() {} - public static Table loadIcebergTable(TrinoCatalog catalog, IcebergTableOperationsProvider tableOperationsProvider, ConnectorSession session, SchemaTableName table) + public static BaseTable loadIcebergTable(TrinoCatalog catalog, IcebergTableOperationsProvider tableOperationsProvider, ConnectorSession session, SchemaTableName table) { TableOperations operations = tableOperationsProvider.createTableOperations( catalog, @@ -196,7 +228,7 @@ public static Table loadIcebergTable(TrinoCatalog catalog, IcebergTableOperation return new BaseTable(operations, quotedTableName(table), TRINO_METRICS_REPORTER); } - public static Table getIcebergTableWithMetadata( + public static BaseTable getIcebergTableWithMetadata( TrinoCatalog catalog, IcebergTableOperationsProvider tableOperationsProvider, ConnectorSession session, @@ -214,7 +246,59 @@ public static Table getIcebergTableWithMetadata( return new BaseTable(operations, quotedTableName(table), TRINO_METRICS_REPORTER); } - public static Map getIcebergTableProperties(Table icebergTable) + public static List getProjectedColumns(Schema schema, TypeManager typeManager) + { + Map indexById = TypeUtil.indexById(schema.asStruct()); + return getProjectedColumns(schema, typeManager, indexById, indexById.keySet() /* project all columns */); + } + + public static List getProjectedColumns(Schema schema, TypeManager typeManager, Set fieldIds) + { + Map indexById = TypeUtil.indexById(schema.asStruct()); + return getProjectedColumns(schema, typeManager, indexById, fieldIds /* project selected columns */); + } + + private static List getProjectedColumns(Schema schema, TypeManager typeManager, Map indexById, Set fieldIds) + { + ImmutableList.Builder columns = builderWithExpectedSize(fieldIds.size()); + Map indexParents = TypeUtil.indexParents(schema.asStruct()); + Map> indexPaths = indexById.entrySet().stream() + .collect(toImmutableMap(Entry::getKey, entry -> ImmutableList.copyOf(buildPath(indexParents, entry.getKey())))); + + for (int fieldId : fieldIds) { + columns.add(createColumnHandle(typeManager, fieldId, indexById, indexPaths)); + } + return columns.build(); + } + + public static IcebergColumnHandle createColumnHandle(TypeManager typeManager, int fieldId, Map indexById, Map> indexPaths) + { + NestedField childField = indexById.get(fieldId); + NestedField baseField = childField; + + List path = requireNonNull(indexPaths.get(fieldId)); + if (!path.isEmpty()) { + baseField = indexById.get(path.get(0)); + path = ImmutableList.builder() + .addAll(path.subList(1, path.size())) // Base column id shouldn't exist in IcebergColumnHandle.path + .add(fieldId) // Append the leaf field id + .build(); + } + return createColumnHandle(baseField, childField, typeManager, path); + } + + public static List buildPath(Map indexParents, int fieldId) + { + List path = new ArrayList<>(); + while (indexParents.containsKey(fieldId)) { + int parentId = indexParents.get(fieldId); + path.add(parentId); + fieldId = parentId; + } + return ImmutableList.copyOf(Lists.reverse(path)); + } + + public static Map getIcebergTableProperties(BaseTable icebergTable) { ImmutableMap.Builder properties = ImmutableMap.builder(); properties.put(FILE_FORMAT_PROPERTY, getFileFormat(icebergTable)); @@ -233,22 +317,68 @@ public static Map getIcebergTableProperties(Table icebergTable) properties.put(LOCATION_PROPERTY, icebergTable.location()); } - int formatVersion = ((BaseTable) icebergTable).operations().current().formatVersion(); + int formatVersion = icebergTable.operations().current().formatVersion(); properties.put(FORMAT_VERSION_PROPERTY, formatVersion); + if (icebergTable.properties().containsKey(COMMIT_NUM_RETRIES)) { + int commitNumRetries = parseInt(icebergTable.properties().get(COMMIT_NUM_RETRIES)); + properties.put(MAX_COMMIT_RETRY, commitNumRetries); + } + // iceberg ORC format bloom filter properties - String orcBloomFilterColumns = icebergTable.properties().get(ORC_BLOOM_FILTER_COLUMNS_KEY); - if (orcBloomFilterColumns != null) { - properties.put(ORC_BLOOM_FILTER_COLUMNS, Splitter.on(',').trimResults().omitEmptyStrings().splitToList(orcBloomFilterColumns)); + Optional orcBloomFilterColumns = getOrcBloomFilterColumns(icebergTable.properties()); + if (orcBloomFilterColumns.isPresent()) { + properties.put(ORC_BLOOM_FILTER_COLUMNS_PROPERTY, Splitter.on(',').trimResults().omitEmptyStrings().splitToList(orcBloomFilterColumns.get())); } - String orcBloomFilterFpp = icebergTable.properties().get(ORC_BLOOM_FILTER_FPP_KEY); - if (orcBloomFilterFpp != null) { - properties.put(ORC_BLOOM_FILTER_FPP, Double.parseDouble(orcBloomFilterFpp)); + // iceberg ORC format bloom filter properties + Optional orcBloomFilterFpp = getOrcBloomFilterFpp(icebergTable.properties()); + if (orcBloomFilterFpp.isPresent()) { + properties.put(ORC_BLOOM_FILTER_FPP_PROPERTY, Double.parseDouble(orcBloomFilterFpp.get())); + } + + // iceberg Parquet format bloom filter properties + Set parquetBloomFilterColumns = getParquetBloomFilterColumns(icebergTable.properties()); + if (!parquetBloomFilterColumns.isEmpty()) { + properties.put(PARQUET_BLOOM_FILTER_COLUMNS_PROPERTY, ImmutableList.copyOf(parquetBloomFilterColumns)); + } + + if (parseBoolean(icebergTable.properties().getOrDefault(OBJECT_STORE_ENABLED, "false"))) { + properties.put(OBJECT_STORE_LAYOUT_ENABLED_PROPERTY, true); } + Optional dataLocation = Optional.ofNullable(icebergTable.properties().get(WRITE_DATA_LOCATION)); + dataLocation.ifPresent(location -> properties.put(DATA_LOCATION_PROPERTY, location)); + return properties.buildOrThrow(); } + // Version 382-438 set incorrect table properties: https://github.com/trinodb/trino/commit/b89aac68c43e5392f23b8d6ba053bbeb6df85028#diff-2af3e19a6b656640a7d0bb73114ef224953a2efa04e569b1fe4da953b2cc6d15R418-R419 + // `orc.bloom.filter.columns` was set instead of `write.orc.bloom.filter.columns`, and `orc.bloom.filter.fpp` instead of `write.orc.bloom.filter.fpp` + // These methods maintain backward compatibility for existing table. + public static Optional getOrcBloomFilterColumns(Map properties) + { + return Stream.of( + properties.get(ORC_BLOOM_FILTER_COLUMNS)) + .filter(Objects::nonNull) + .findFirst(); + } + + public static Set getParquetBloomFilterColumns(Map properties) + { + return properties.entrySet().stream() + .filter(entry -> entry.getKey().startsWith(PARQUET_BLOOM_FILTER_COLUMN_ENABLED_PREFIX) && "true".equals(entry.getValue())) + .map(entry -> entry.getKey().substring(PARQUET_BLOOM_FILTER_COLUMN_ENABLED_PREFIX.length())) + .collect(toImmutableSet()); + } + + public static Optional getOrcBloomFilterFpp(Map properties) + { + return Stream.of( + properties.get(ORC_BLOOM_FILTER_FPP)) + .filter(Objects::nonNull) + .findFirst(); + } + public static List getColumns(Schema schema, TypeManager typeManager) { return schema.columns().stream() @@ -256,10 +386,17 @@ public static List getColumns(Schema schema, TypeManager ty .collect(toImmutableList()); } + public static List getTopLevelColumns(Schema schema, TypeManager typeManager) + { + return schema.columns().stream() + .map(column -> getColumnHandle(column, typeManager)) + .collect(toImmutableList()); + } + public static List getColumnMetadatas(Schema schema, TypeManager typeManager) { List icebergColumns = schema.columns(); - ImmutableList.Builder columns = ImmutableList.builderWithExpectedSize(icebergColumns.size() + 2); + ImmutableList.Builder columns = builderWithExpectedSize(icebergColumns.size() + 2); icebergColumns.stream() .map(column -> @@ -275,23 +412,45 @@ public static List getColumnMetadatas(Schema schema, TypeManager return columns.build(); } + public static Schema updateColumnComment(Schema schema, String columnName, String comment) + { + NestedField fieldToUpdate = schema.findField(columnName); + checkArgument(fieldToUpdate != null, "Field %s does not exist", columnName); + NestedField updatedField = NestedField.of(fieldToUpdate.fieldId(), fieldToUpdate.isOptional(), fieldToUpdate.name(), fieldToUpdate.type(), comment); + List newFields = schema.columns().stream() + .map(field -> (field.fieldId() == updatedField.fieldId()) ? updatedField : field) + .toList(); + + return new Schema(newFields, schema.getAliases(), schema.identifierFieldIds()); + } + public static IcebergColumnHandle getColumnHandle(NestedField column, TypeManager typeManager) { - Type type = toTrinoType(column.type(), typeManager); + return createColumnHandle(column, column, typeManager, ImmutableList.of()); + } + + private static IcebergColumnHandle createColumnHandle(NestedField baseColumn, NestedField childColumn, TypeManager typeManager, List path) + { return new IcebergColumnHandle( - createColumnIdentity(column), - type, - ImmutableList.of(), - type, - Optional.ofNullable(column.doc())); + createColumnIdentity(baseColumn), + toTrinoType(baseColumn.type(), typeManager), + path, + toTrinoType(childColumn.type(), typeManager), + childColumn.isOptional(), + Optional.ofNullable(childColumn.doc())); } public static Schema schemaFromHandles(List columns) + { + return structTypeFromHandles(columns).asSchema(); + } + + public static StructType structTypeFromHandles(List columns) { List icebergColumns = columns.stream() .map(column -> NestedField.optional(column.getId(), column.getName(), toIcebergType(column.getType(), column.getColumnIdentity()))) .collect(toImmutableList()); - return new Schema(StructType.of(icebergColumns).asStructType().fields()); + return StructType.of(icebergColumns); } public static Map getIdentityPartitions(PartitionSpec partitionSpec) @@ -307,6 +466,33 @@ public static Map getIdentityPartitions(PartitionSpec p return columns.buildOrThrow(); } + public static List primitiveFields(Schema schema) + { + return primitiveFields(schema.columns()) + .collect(toImmutableList()); + } + + private static Stream primitiveFields(List nestedFields) + { + return nestedFields.stream() + .flatMap(IcebergUtil::primitiveFields); + } + + private static Stream primitiveFields(NestedField nestedField) + { + org.apache.iceberg.types.Type type = nestedField.type(); + if (type.isPrimitiveType()) { + return Stream.of(nestedField); + } + + if (type.isNestedType()) { + return primitiveFields(type.asNestedType().fields()) + .map(field -> Types.NestedField.of(field.fieldId(), field.isOptional(), nestedField.name() + "." + field.name(), field.type(), field.doc())); + } + + throw new IllegalStateException("Unsupported field type: " + nestedField); + } + public static Map primitiveFieldTypes(Schema schema) { return primitiveFieldTypes(schema.columns()) @@ -370,9 +556,15 @@ public static boolean canEnforceColumnConstraintInSpecs( IcebergColumnHandle columnHandle, Domain domain) { - return table.specs().values().stream() + List partitionSpecs = table.specs().values().stream() .filter(partitionSpec -> partitionSpecIds.contains(partitionSpec.specId())) - .allMatch(spec -> canEnforceConstraintWithinPartitioningSpec(typeOperators, spec, columnHandle, domain)); + .collect(toImmutableList()); + + if (partitionSpecs.isEmpty()) { + return canEnforceConstraintWithinPartitioningSpec(typeOperators, table.spec(), columnHandle, domain); + } + + return partitionSpecs.stream().allMatch(spec -> canEnforceConstraintWithinPartitioningSpec(typeOperators, spec, columnHandle, domain)); } private static boolean canEnforceConstraintWithinPartitioningSpec(TypeOperators typeOperators, PartitionSpec spec, IcebergColumnHandle column, Domain domain) @@ -585,13 +777,34 @@ public static Map> getPartitionKeys(StructLike partiti return partitionKeys.buildOrThrow(); } + public static Map getPartitionValues( + Set identityPartitionColumns, + Map> partitionKeys) + { + ImmutableMap.Builder bindings = ImmutableMap.builder(); + for (IcebergColumnHandle partitionColumn : identityPartitionColumns) { + Object partitionValue = deserializePartitionValue( + partitionColumn.getType(), + partitionKeys.get(partitionColumn.getId()).orElse(null), + partitionColumn.getName()); + NullableValue bindingValue = new NullableValue(partitionColumn.getType(), partitionValue); + bindings.put(partitionColumn, bindingValue); + } + return bindings.buildOrThrow(); + } + public static LocationProvider getLocationProvider(SchemaTableName schemaTableName, String tableLocation, Map storageProperties) { if (storageProperties.containsKey(WRITE_LOCATION_PROVIDER_IMPL)) { throw new TrinoException(NOT_SUPPORTED, "Table " + schemaTableName + " specifies " + storageProperties.get(WRITE_LOCATION_PROVIDER_IMPL) + " as a location provider. Writing to Iceberg tables with custom location provider is not supported."); } - return locationsFor(tableLocation, storageProperties); + + if (propertyAsBoolean(storageProperties, OBJECT_STORE_ENABLED, OBJECT_STORE_ENABLED_DEFAULT)) { + return new ObjectStoreLocationProvider(tableLocation, storageProperties); + } + + return new DefaultLocationProvider(tableLocation, storageProperties); } public static Schema schemaFromMetadata(List columns) @@ -611,34 +824,112 @@ public static Schema schemaFromMetadata(List columns) return new Schema(icebergSchema.asStructType().fields()); } - public static Transaction newCreateTableTransaction(TrinoCatalog catalog, ConnectorTableMetadata tableMetadata, ConnectorSession session) + public static Schema schemaFromViewColumns(TypeManager typeManager, List columns) + { + List icebergColumns = new ArrayList<>(); + AtomicInteger nextFieldId = new AtomicInteger(1); + for (ViewColumn column : columns) { + Type trinoType = typeManager.getType(column.getType()); + org.apache.iceberg.types.Type type = toIcebergTypeForNewColumn(trinoType, nextFieldId); + NestedField field = NestedField.of(nextFieldId.getAndIncrement(), false, column.getName(), type, column.getComment().orElse(null)); + icebergColumns.add(field); + } + org.apache.iceberg.types.Type icebergSchema = StructType.of(icebergColumns); + return new Schema(icebergSchema.asStructType().fields()); + } + + public static List viewColumnsFromSchema(TypeManager typeManager, Schema schema) + { + return IcebergUtil.getTopLevelColumns(schema, typeManager).stream() + .map(column -> new ViewColumn(column.getName(), column.getType().getTypeId(), column.getComment())) + .toList(); + } + + public static Transaction newCreateTableTransaction(TrinoCatalog catalog, ConnectorTableMetadata tableMetadata, ConnectorSession session, String tableLocation, Predicate allowedExtraProperties) { SchemaTableName schemaTableName = tableMetadata.getTable(); Schema schema = schemaFromMetadata(tableMetadata.getColumns()); PartitionSpec partitionSpec = parsePartitionFields(schema, getPartitioning(tableMetadata.getProperties())); SortOrder sortOrder = parseSortFields(schema, getSortOrder(tableMetadata.getProperties())); - String targetPath = getTableLocation(tableMetadata.getProperties()) - .orElseGet(() -> catalog.defaultTableLocation(session, schemaTableName)); + return catalog.newCreateTableTransaction(session, schemaTableName, schema, partitionSpec, sortOrder, Optional.ofNullable(tableLocation), createTableProperties(tableMetadata, allowedExtraProperties)); + } + + public static Map createTableProperties(ConnectorTableMetadata tableMetadata, Predicate allowedExtraProperties) + { ImmutableMap.Builder propertiesBuilder = ImmutableMap.builder(); IcebergFileFormat fileFormat = IcebergTableProperties.getFileFormat(tableMetadata.getProperties()); propertiesBuilder.put(DEFAULT_FILE_FORMAT, fileFormat.toIceberg().toString()); propertiesBuilder.put(FORMAT_VERSION, Integer.toString(IcebergTableProperties.getFormatVersion(tableMetadata.getProperties()))); + propertiesBuilder.put(COMMIT_NUM_RETRIES, Integer.toString(IcebergTableProperties.getMaxCommitRetry(tableMetadata.getProperties()))); + + boolean objectStoreLayoutEnabled = IcebergTableProperties.getObjectStoreLayoutEnabled(tableMetadata.getProperties()); + if (objectStoreLayoutEnabled) { + propertiesBuilder.put(OBJECT_STORE_ENABLED, "true"); + } + Optional dataLocation = IcebergTableProperties.getDataLocation(tableMetadata.getProperties()); + dataLocation.ifPresent(location -> { + if (!objectStoreLayoutEnabled) { + throw new TrinoException(INVALID_TABLE_PROPERTY, "Data location can only be set when object store layout is enabled"); + } + propertiesBuilder.put(WRITE_DATA_LOCATION, location); + }); // iceberg ORC format bloom filter properties used by create table - List columns = getOrcBloomFilterColumns(tableMetadata.getProperties()); - if (!columns.isEmpty()) { - checkFormatForProperty(fileFormat.toIceberg(), FileFormat.ORC, ORC_BLOOM_FILTER_COLUMNS); - validateOrcBloomFilterColumns(tableMetadata, columns); - propertiesBuilder.put(ORC_BLOOM_FILTER_COLUMNS_KEY, Joiner.on(",").join(columns)); - propertiesBuilder.put(ORC_BLOOM_FILTER_FPP_KEY, String.valueOf(getOrcBloomFilterFpp(tableMetadata.getProperties()))); + List orcBloomFilterColumns = IcebergTableProperties.getOrcBloomFilterColumns(tableMetadata.getProperties()); + if (!orcBloomFilterColumns.isEmpty()) { + checkFormatForProperty(fileFormat.toIceberg(), FileFormat.ORC, ORC_BLOOM_FILTER_COLUMNS_PROPERTY); + validateOrcBloomFilterColumns(tableMetadata.getColumns(), orcBloomFilterColumns); + propertiesBuilder.put(ORC_BLOOM_FILTER_COLUMNS, Joiner.on(",").join(orcBloomFilterColumns)); + propertiesBuilder.put(ORC_BLOOM_FILTER_FPP, String.valueOf(IcebergTableProperties.getOrcBloomFilterFpp(tableMetadata.getProperties()))); + } + + // iceberg Parquet format bloom filter properties used by create table + List parquetBloomFilterColumns = IcebergTableProperties.getParquetBloomFilterColumns(tableMetadata.getProperties()); + if (!parquetBloomFilterColumns.isEmpty()) { + checkFormatForProperty(fileFormat.toIceberg(), FileFormat.PARQUET, PARQUET_BLOOM_FILTER_COLUMNS_PROPERTY); + validateParquetBloomFilterColumns(tableMetadata.getColumns(), parquetBloomFilterColumns); + for (String column : parquetBloomFilterColumns) { + propertiesBuilder.put(PARQUET_BLOOM_FILTER_COLUMN_ENABLED_PREFIX + column, "true"); + } } if (tableMetadata.getComment().isPresent()) { propertiesBuilder.put(TABLE_COMMENT, tableMetadata.getComment().get()); } - return catalog.newCreateTableTransaction(session, schemaTableName, schema, partitionSpec, sortOrder, targetPath, propertiesBuilder.buildOrThrow()); + Map baseProperties = propertiesBuilder.buildOrThrow(); + Map extraProperties = IcebergTableProperties.getExtraProperties(tableMetadata.getProperties()).orElseGet(ImmutableMap::of); + + verifyExtraProperties(baseProperties.keySet(), extraProperties, allowedExtraProperties); + + return ImmutableMap.builder() + .putAll(baseProperties) + .putAll(extraProperties) + .buildOrThrow(); + } + + public static void verifyExtraProperties(Set basePropertyKeys, Map extraProperties, Predicate allowedExtraProperties) + { + Set illegalExtraProperties = ImmutableSet.builder() + .addAll(Sets.intersection( + ImmutableSet.builder() + .add(TABLE_COMMENT) + .addAll(basePropertyKeys) + .addAll(SUPPORTED_PROPERTIES) + .addAll(PROTECTED_ICEBERG_NATIVE_PROPERTIES) + .build(), + extraProperties.keySet())) + .addAll(extraProperties.keySet().stream() + .filter(name -> !allowedExtraProperties.test(name)) + .collect(toImmutableSet())) + .build(); + + if (!illegalExtraProperties.isEmpty()) { + throw new TrinoException( + INVALID_TABLE_PROPERTY, + format("Illegal keys in extra_properties: %s", illegalExtraProperties)); + } } /** @@ -718,16 +1009,16 @@ public static void validateTableCanBeDropped(Table table) } } - private static void checkFormatForProperty(FileFormat actualStorageFormat, FileFormat expectedStorageFormat, String propertyName) + public static void checkFormatForProperty(FileFormat actualStorageFormat, FileFormat expectedStorageFormat, String propertyName) { if (actualStorageFormat != expectedStorageFormat) { throw new TrinoException(INVALID_TABLE_PROPERTY, format("Cannot specify %s table property for storage format: %s", propertyName, actualStorageFormat)); } } - private static void validateOrcBloomFilterColumns(ConnectorTableMetadata tableMetadata, List orcBloomFilterColumns) + public static void validateOrcBloomFilterColumns(List columns, List orcBloomFilterColumns) { - Set allColumns = tableMetadata.getColumns().stream() + Set allColumns = columns.stream() .map(ColumnMetadata::getName) .collect(toImmutableSet()); if (!allColumns.containsAll(orcBloomFilterColumns)) { @@ -735,6 +1026,21 @@ private static void validateOrcBloomFilterColumns(ConnectorTableMetadata tableMe } } + public static void validateParquetBloomFilterColumns(List columns, List parquetBloomFilterColumns) + { + Map columnTypes = columns.stream() + .collect(toImmutableMap(ColumnMetadata::getName, ColumnMetadata::getType)); + for (String column : parquetBloomFilterColumns) { + Type type = columnTypes.get(column); + if (type == null) { + throw new TrinoException(INVALID_TABLE_PROPERTY, format("Parquet Bloom filter column %s not present in schema", column)); + } + if (!SUPPORTED_BLOOM_FILTER_TYPES.contains(type)) { + throw new TrinoException(INVALID_TABLE_PROPERTY, format("Parquet Bloom filter column %s has unsupported type %s", column, type.getDisplayName())); + } + } + } + public static int parseVersion(String metadataFileName) throws TrinoException { @@ -774,6 +1080,7 @@ public static String fileName(String path) public static void commit(SnapshotUpdate update, ConnectorSession session) { update.set(TRINO_QUERY_ID_NAME, session.getQueryId()); + update.set(TRINO_USER_NAME, session.getUser()); update.commit(); } @@ -788,4 +1095,77 @@ public static Map columnNameToPositionInSchema(Schema schema) (column, position) -> immutableEntry(column.name(), Long.valueOf(position).intValue())) .collect(toImmutableMap(Entry::getKey, Entry::getValue)); } + + public static Domain getPartitionDomain(TupleDomain effectivePredicate) + { + IcebergColumnHandle partitionColumn = partitionColumnHandle(); + Domain domain = effectivePredicate.getDomains().orElseThrow(() -> new IllegalArgumentException("Unexpected NONE tuple domain")) + .get(partitionColumn); + if (domain == null) { + return Domain.all(partitionColumn.getType()); + } + return domain; + } + + public static Domain getPathDomain(TupleDomain effectivePredicate) + { + IcebergColumnHandle pathColumn = pathColumnHandle(); + Domain domain = effectivePredicate.getDomains().orElseThrow(() -> new IllegalArgumentException("Unexpected NONE tuple domain")) + .get(pathColumn); + if (domain == null) { + return Domain.all(pathColumn.getType()); + } + return domain; + } + + public static Domain getFileModifiedTimePathDomain(TupleDomain effectivePredicate) + { + IcebergColumnHandle fileModifiedTimeColumn = fileModifiedTimeColumnHandle(); + Domain domain = effectivePredicate.getDomains().orElseThrow(() -> new IllegalArgumentException("Unexpected NONE tuple domain")) + .get(fileModifiedTimeColumn); + if (domain == null) { + return Domain.all(fileModifiedTimeColumn.getType()); + } + return domain; + } + + public static long getModificationTime(String path, TrinoFileSystem fileSystem) + { + try { + TrinoInputFile inputFile = fileSystem.newInputFile(Location.of(path)); + return inputFile.lastModified().toEpochMilli(); + } + catch (IOException e) { + throw new TrinoException(ICEBERG_FILESYSTEM_ERROR, "Failed to get file modification time: " + path, e); + } + } + + public static Optional getPartitionColumnType(List fields, Schema schema, TypeManager typeManager) + { + if (fields.isEmpty()) { + return Optional.empty(); + } + List partitionFields = fields.stream() + .map(field -> RowType.field( + field.name(), + toTrinoType(field.transform().getResultType(schema.findType(field.sourceId())), typeManager))) + .collect(toImmutableList()); + List fieldIds = fields.stream() + .map(PartitionField::fieldId) + .collect(toImmutableList()); + return Optional.of(new IcebergPartitionColumn(RowType.from(partitionFields), fieldIds)); + } + + public static List partitionTypes( + List partitionFields, + Map idToPrimitiveTypeMapping) + { + ImmutableList.Builder partitionTypeBuilder = ImmutableList.builder(); + for (PartitionField partitionField : partitionFields) { + org.apache.iceberg.types.Type.PrimitiveType sourceType = idToPrimitiveTypeMapping.get(partitionField.sourceId()); + org.apache.iceberg.types.Type type = partitionField.transform().getResultType(sourceType); + partitionTypeBuilder.add(type); + } + return partitionTypeBuilder.build(); + } } diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergWritableTableHandle.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergWritableTableHandle.java index 22c885b4c465..f46cf4640309 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergWritableTableHandle.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergWritableTableHandle.java @@ -41,6 +41,7 @@ public class IcebergWritableTableHandle private final IcebergFileFormat fileFormat; private final Map storageProperties; private final RetryMode retryMode; + private final Map fileIoProperties; @JsonCreator public IcebergWritableTableHandle( @@ -53,7 +54,8 @@ public IcebergWritableTableHandle( @JsonProperty("outputPath") String outputPath, @JsonProperty("fileFormat") IcebergFileFormat fileFormat, @JsonProperty("properties") Map storageProperties, - @JsonProperty("retryMode") RetryMode retryMode) + @JsonProperty("retryMode") RetryMode retryMode, + @JsonProperty("fileIoProperties") Map fileIoProperties) { this.name = requireNonNull(name, "name is null"); this.schemaAsJson = requireNonNull(schemaAsJson, "schemaAsJson is null"); @@ -66,6 +68,7 @@ public IcebergWritableTableHandle( this.storageProperties = ImmutableMap.copyOf(requireNonNull(storageProperties, "storageProperties is null")); this.retryMode = requireNonNull(retryMode, "retryMode is null"); checkArgument(partitionsSpecsAsJson.containsKey(partitionSpecId), "partitionSpecId missing from partitionSpecs"); + this.fileIoProperties = ImmutableMap.copyOf(requireNonNull(fileIoProperties, "fileIoProperties is null")); } @JsonProperty @@ -128,6 +131,12 @@ public RetryMode getRetryMode() return retryMode; } + @JsonProperty + public Map getFileIoProperties() + { + return fileIoProperties; + } + @Override public String toString() { diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/InternalIcebergConnectorFactory.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/InternalIcebergConnectorFactory.java deleted file mode 100644 index 716abdcebf60..000000000000 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/InternalIcebergConnectorFactory.java +++ /dev/null @@ -1,153 +0,0 @@ -/* - * 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.iceberg; - -import com.google.inject.Injector; -import com.google.inject.Key; -import com.google.inject.Module; -import com.google.inject.TypeLiteral; -import io.airlift.bootstrap.Bootstrap; -import io.airlift.bootstrap.LifeCycleManager; -import io.airlift.event.client.EventModule; -import io.airlift.json.JsonModule; -import io.opentelemetry.api.OpenTelemetry; -import io.opentelemetry.api.trace.Tracer; -import io.trino.filesystem.TrinoFileSystemFactory; -import io.trino.filesystem.manager.FileSystemModule; -import io.trino.hdfs.HdfsModule; -import io.trino.hdfs.authentication.HdfsAuthenticationModule; -import io.trino.hdfs.azure.HiveAzureModule; -import io.trino.hdfs.gcs.HiveGcsModule; -import io.trino.plugin.base.CatalogName; -import io.trino.plugin.base.classloader.ClassLoaderSafeConnectorPageSinkProvider; -import io.trino.plugin.base.classloader.ClassLoaderSafeConnectorPageSourceProvider; -import io.trino.plugin.base.classloader.ClassLoaderSafeConnectorSplitManager; -import io.trino.plugin.base.classloader.ClassLoaderSafeNodePartitioningProvider; -import io.trino.plugin.base.jmx.ConnectorObjectNameGeneratorModule; -import io.trino.plugin.base.jmx.MBeanServerModule; -import io.trino.plugin.base.session.SessionPropertiesProvider; -import io.trino.plugin.hive.NodeVersion; -import io.trino.plugin.iceberg.catalog.IcebergCatalogModule; -import io.trino.spi.NodeManager; -import io.trino.spi.PageIndexerFactory; -import io.trino.spi.PageSorter; -import io.trino.spi.classloader.ThreadContextClassLoader; -import io.trino.spi.connector.CatalogHandle; -import io.trino.spi.connector.Connector; -import io.trino.spi.connector.ConnectorAccessControl; -import io.trino.spi.connector.ConnectorContext; -import io.trino.spi.connector.ConnectorNodePartitioningProvider; -import io.trino.spi.connector.ConnectorPageSinkProvider; -import io.trino.spi.connector.ConnectorPageSourceProvider; -import io.trino.spi.connector.ConnectorSplitManager; -import io.trino.spi.connector.TableProcedureMetadata; -import io.trino.spi.procedure.Procedure; -import io.trino.spi.session.PropertyMetadata; -import io.trino.spi.type.TypeManager; -import org.weakref.jmx.guice.MBeanModule; - -import java.util.Collection; -import java.util.List; -import java.util.Map; -import java.util.Optional; -import java.util.Set; -import java.util.stream.Stream; - -import static com.google.common.collect.ImmutableList.toImmutableList; - -public final class InternalIcebergConnectorFactory -{ - private InternalIcebergConnectorFactory() {} - - public static Connector createConnector( - String catalogName, - Map config, - ConnectorContext context, - Module module, - Optional icebergCatalogModule, - Optional fileSystemFactory) - { - ClassLoader classLoader = InternalIcebergConnectorFactory.class.getClassLoader(); - try (ThreadContextClassLoader ignored = new ThreadContextClassLoader(classLoader)) { - Bootstrap app = new Bootstrap( - new EventModule(), - new MBeanModule(), - new ConnectorObjectNameGeneratorModule("io.trino.plugin.iceberg", "trino.plugin.iceberg"), - new JsonModule(), - new IcebergModule(), - new IcebergSecurityModule(), - icebergCatalogModule.orElse(new IcebergCatalogModule()), - new HdfsModule(), - new HiveGcsModule(), - new HiveAzureModule(), - new HdfsAuthenticationModule(), - new MBeanServerModule(), - fileSystemFactory - .map(factory -> (Module) binder -> binder.bind(TrinoFileSystemFactory.class).toInstance(factory)) - .orElseGet(FileSystemModule::new), - binder -> { - binder.bind(OpenTelemetry.class).toInstance(context.getOpenTelemetry()); - binder.bind(Tracer.class).toInstance(context.getTracer()); - binder.bind(NodeVersion.class).toInstance(new NodeVersion(context.getNodeManager().getCurrentNode().getVersion())); - binder.bind(NodeManager.class).toInstance(context.getNodeManager()); - binder.bind(TypeManager.class).toInstance(context.getTypeManager()); - binder.bind(PageIndexerFactory.class).toInstance(context.getPageIndexerFactory()); - binder.bind(CatalogHandle.class).toInstance(context.getCatalogHandle()); - binder.bind(CatalogName.class).toInstance(new CatalogName(catalogName)); - binder.bind(PageSorter.class).toInstance(context.getPageSorter()); - }, - module); - - Injector injector = app - .doNotInitializeLogging() - .setRequiredConfigurationProperties(config) - .initialize(); - - LifeCycleManager lifeCycleManager = injector.getInstance(LifeCycleManager.class); - IcebergTransactionManager transactionManager = injector.getInstance(IcebergTransactionManager.class); - ConnectorSplitManager splitManager = injector.getInstance(ConnectorSplitManager.class); - ConnectorPageSourceProvider connectorPageSource = injector.getInstance(ConnectorPageSourceProvider.class); - ConnectorPageSinkProvider pageSinkProvider = injector.getInstance(ConnectorPageSinkProvider.class); - ConnectorNodePartitioningProvider connectorDistributionProvider = injector.getInstance(ConnectorNodePartitioningProvider.class); - Set sessionPropertiesProviders = injector.getInstance(Key.get(new TypeLiteral>() {})); - IcebergTableProperties icebergTableProperties = injector.getInstance(IcebergTableProperties.class); - IcebergMaterializedViewAdditionalProperties materializedViewAdditionalProperties = injector.getInstance(IcebergMaterializedViewAdditionalProperties.class); - IcebergAnalyzeProperties icebergAnalyzeProperties = injector.getInstance(IcebergAnalyzeProperties.class); - Set procedures = injector.getInstance(Key.get(new TypeLiteral>() {})); - Set tableProcedures = injector.getInstance(Key.get(new TypeLiteral>() {})); - Optional accessControl = injector.getInstance(Key.get(new TypeLiteral>() {})); - // Materialized view should allow configuring all the supported iceberg table properties for the storage table - List> materializedViewProperties = Stream.of(icebergTableProperties.getTableProperties(), materializedViewAdditionalProperties.getMaterializedViewProperties()) - .flatMap(Collection::stream) - .collect(toImmutableList()); - - return new IcebergConnector( - injector, - lifeCycleManager, - transactionManager, - new ClassLoaderSafeConnectorSplitManager(splitManager, classLoader), - new ClassLoaderSafeConnectorPageSourceProvider(connectorPageSource, classLoader), - new ClassLoaderSafeConnectorPageSinkProvider(pageSinkProvider, classLoader), - new ClassLoaderSafeNodePartitioningProvider(connectorDistributionProvider, classLoader), - sessionPropertiesProviders, - IcebergSchemaProperties.SCHEMA_PROPERTIES, - icebergTableProperties.getTableProperties(), - materializedViewProperties, - icebergAnalyzeProperties.getAnalyzeProperties(), - accessControl, - procedures, - tableProcedures); - } - } -} diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/MetadataLogEntriesTable.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/MetadataLogEntriesTable.java new file mode 100644 index 000000000000..fd6de0b40bb6 --- /dev/null +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/MetadataLogEntriesTable.java @@ -0,0 +1,76 @@ +/* + * 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.iceberg; + +import com.google.common.collect.ImmutableList; +import io.trino.plugin.iceberg.util.PageListBuilder; +import io.trino.spi.connector.ColumnMetadata; +import io.trino.spi.connector.ConnectorTableMetadata; +import io.trino.spi.connector.SchemaTableName; +import io.trino.spi.type.TimeZoneKey; +import org.apache.iceberg.Table; + +import java.util.concurrent.ExecutorService; + +import static io.trino.spi.type.BigintType.BIGINT; +import static io.trino.spi.type.IntegerType.INTEGER; +import static io.trino.spi.type.TimestampWithTimeZoneType.TIMESTAMP_TZ_MILLIS; +import static io.trino.spi.type.Timestamps.MICROSECONDS_PER_MILLISECOND; +import static io.trino.spi.type.VarcharType.VARCHAR; +import static java.util.Objects.requireNonNull; +import static org.apache.iceberg.MetadataTableType.METADATA_LOG_ENTRIES; + +public class MetadataLogEntriesTable + extends BaseSystemTable +{ + private static final String TIMESTAMP_COLUMN_NAME = "timestamp"; + private static final String FILE_COLUMN_NAME = "file"; + private static final String LATEST_SNAPSHOT_ID_COLUMN_NAME = "latest_snapshot_id"; + private static final String LATEST_SCHEMA_ID_COLUMN_NAME = "latest_schema_id"; + private static final String LATEST_SEQUENCE_NUMBER_COLUMN_NAME = "latest_sequence_number"; + + public MetadataLogEntriesTable(SchemaTableName tableName, Table icebergTable, ExecutorService executor) + { + super( + requireNonNull(icebergTable, "icebergTable is null"), + createConnectorTableMetadata(requireNonNull(tableName, "tableName is null")), + METADATA_LOG_ENTRIES, + executor); + } + + private static ConnectorTableMetadata createConnectorTableMetadata(SchemaTableName tableName) + { + return new ConnectorTableMetadata( + tableName, + ImmutableList.builder() + .add(new ColumnMetadata(TIMESTAMP_COLUMN_NAME, TIMESTAMP_TZ_MILLIS)) + .add(new ColumnMetadata(FILE_COLUMN_NAME, VARCHAR)) + .add(new ColumnMetadata(LATEST_SNAPSHOT_ID_COLUMN_NAME, BIGINT)) + .add(new ColumnMetadata(LATEST_SCHEMA_ID_COLUMN_NAME, INTEGER)) + .add(new ColumnMetadata(LATEST_SEQUENCE_NUMBER_COLUMN_NAME, BIGINT)) + .build()); + } + + @Override + protected void addRow(PageListBuilder pagesBuilder, Row row, TimeZoneKey timeZoneKey) + { + pagesBuilder.beginRow(); + pagesBuilder.appendTimestampTzMillis(row.get(TIMESTAMP_COLUMN_NAME, Long.class) / MICROSECONDS_PER_MILLISECOND, timeZoneKey); + pagesBuilder.appendVarchar(row.get(FILE_COLUMN_NAME, String.class)); + pagesBuilder.appendBigint(row.get(LATEST_SNAPSHOT_ID_COLUMN_NAME, Long.class)); + pagesBuilder.appendInteger(row.get(LATEST_SCHEMA_ID_COLUMN_NAME, Integer.class)); + pagesBuilder.appendBigint(row.get(LATEST_SEQUENCE_NUMBER_COLUMN_NAME, Long.class)); + pagesBuilder.endRow(); + } +} diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/PartitionFields.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/PartitionFields.java index 8c7a2bf64fac..f1e54ec7d95a 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/PartitionFields.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/PartitionFields.java @@ -51,8 +51,8 @@ public final class PartitionFields private static final Pattern TRUNCATE_PATTERN = Pattern.compile("truncate" + FUNCTION_ARGUMENT_NAME_AND_INT, CASE_INSENSITIVE); private static final Pattern VOID_PATTERN = Pattern.compile("void" + FUNCTION_ARGUMENT_NAME, CASE_INSENSITIVE); - private static final Pattern ICEBERG_BUCKET_PATTERN = Pattern.compile("bucket\\[(\\d+)]"); - private static final Pattern ICEBERG_TRUNCATE_PATTERN = Pattern.compile("truncate\\[(\\d+)]"); + public static final Pattern ICEBERG_BUCKET_PATTERN = Pattern.compile("bucket\\[(\\d+)]"); + public static final Pattern ICEBERG_TRUNCATE_PATTERN = Pattern.compile("truncate\\[(\\d+)]"); private PartitionFields() {} @@ -61,7 +61,7 @@ public static PartitionSpec parsePartitionFields(Schema schema, List fie try { PartitionSpec.Builder builder = PartitionSpec.builderFor(schema); for (String field : fields) { - parsePartitionField(builder, field); + parsePartitionFields(schema, fields, builder, field); } return builder.build(); } @@ -70,19 +70,59 @@ public static PartitionSpec parsePartitionFields(Schema schema, List fie } } - public static void parsePartitionField(PartitionSpec.Builder builder, String field) + private static void parsePartitionFields(Schema schema, List fields, PartitionSpec.Builder builder, String field) { - @SuppressWarnings("PointlessBooleanExpression") - boolean matched = false || - tryMatch(field, IDENTITY_PATTERN, match -> builder.identity(fromIdentifierToColumn(match.group()))) || - tryMatch(field, YEAR_PATTERN, match -> builder.year(fromIdentifierToColumn(match.group(1)))) || - tryMatch(field, MONTH_PATTERN, match -> builder.month(fromIdentifierToColumn(match.group(1)))) || - tryMatch(field, DAY_PATTERN, match -> builder.day(fromIdentifierToColumn(match.group(1)))) || - tryMatch(field, HOUR_PATTERN, match -> builder.hour(fromIdentifierToColumn(match.group(1)))) || - tryMatch(field, BUCKET_PATTERN, match -> builder.bucket(fromIdentifierToColumn(match.group(1)), parseInt(match.group(2)))) || - tryMatch(field, TRUNCATE_PATTERN, match -> builder.truncate(fromIdentifierToColumn(match.group(1)), parseInt(match.group(2)))) || - tryMatch(field, VOID_PATTERN, match -> builder.alwaysNull(fromIdentifierToColumn(match.group(1)))) || - false; + for (int i = 1; i < schema.columns().size() + fields.size(); i++) { + try { + parsePartitionField(builder, field, i == 1 ? "" : "_" + i); + return; + } + catch (IllegalArgumentException e) { + if (e.getMessage().contains("Cannot create partition from name that exists in schema") + || e.getMessage().contains("Cannot create identity partition sourced from different field in schema")) { + continue; + } + throw e; + } + } + throw new IllegalArgumentException("Cannot resolve partition field: " + field); + } + + public static void parsePartitionField(PartitionSpec.Builder builder, String field, String suffix) + { + boolean matched = + tryMatch(field, IDENTITY_PATTERN, match -> { + // identity doesn't allow specifying an alias + builder.identity(fromIdentifierToColumn(match.group())); + }) || + tryMatch(field, YEAR_PATTERN, match -> { + String column = fromIdentifierToColumn(match.group(1)); + builder.year(column, column + "_year" + suffix); + }) || + tryMatch(field, MONTH_PATTERN, match -> { + String column = fromIdentifierToColumn(match.group(1)); + builder.month(column, column + "_month" + suffix); + }) || + tryMatch(field, DAY_PATTERN, match -> { + String column = fromIdentifierToColumn(match.group(1)); + builder.day(column, column + "_day" + suffix); + }) || + tryMatch(field, HOUR_PATTERN, match -> { + String column = fromIdentifierToColumn(match.group(1)); + builder.hour(column, column + "_hour" + suffix); + }) || + tryMatch(field, BUCKET_PATTERN, match -> { + String column = fromIdentifierToColumn(match.group(1)); + builder.bucket(column, parseInt(match.group(2)), column + "_bucket" + suffix); + }) || + tryMatch(field, TRUNCATE_PATTERN, match -> { + String column = fromIdentifierToColumn(match.group(1)); + builder.truncate(column, parseInt(match.group(2)), column + "_trunc" + suffix); + }) || + tryMatch(field, VOID_PATTERN, match -> { + String column = fromIdentifierToColumn(match.group(1)); + builder.alwaysNull(column, column + "_null" + suffix); + }); if (!matched) { throw new IllegalArgumentException("Invalid partition field declaration: " + field); } @@ -91,13 +131,6 @@ public static void parsePartitionField(PartitionSpec.Builder builder, String fie public static String fromIdentifierToColumn(String identifier) { if (QUOTED_IDENTIFIER_PATTERN.matcher(identifier).matches()) { - // We only support lowercase quoted identifiers for now. - // See https://github.com/trinodb/trino/issues/12226#issuecomment-1128839259 - // TODO: Enhance quoted identifiers support in Iceberg partitioning to support mixed case identifiers - // See https://github.com/trinodb/trino/issues/12668 - if (!identifier.toLowerCase(ENGLISH).equals(identifier)) { - throw new IllegalArgumentException(format("Uppercase characters in identifier '%s' are not supported.", identifier)); - } return identifier.substring(1, identifier.length() - 1).replace("\"\"", "\""); } // Currently, all Iceberg columns are stored in lowercase in the Iceberg metadata files. @@ -159,7 +192,7 @@ private static String fromColumnToIdentifier(String column) public static String quotedName(String name) { - if (UNQUOTED_IDENTIFIER_PATTERN.matcher(name).matches()) { + if (UNQUOTED_IDENTIFIER_PATTERN.matcher(name).matches() && name.toLowerCase(ENGLISH).equals(name)) { return name; } return '"' + name.replace("\"", "\"\"") + '"'; diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/PartitionTable.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/PartitionTable.java index 22108bd9348f..7a3c9ad26c80 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/PartitionTable.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/PartitionTable.java @@ -13,9 +13,7 @@ */ package io.trino.plugin.iceberg; -import com.google.common.annotations.VisibleForTesting; import com.google.common.collect.ImmutableList; -import com.google.common.collect.ImmutableMap; import io.trino.spi.block.Block; import io.trino.spi.connector.ColumnMetadata; import io.trino.spi.connector.ConnectorSession; @@ -32,14 +30,11 @@ import org.apache.iceberg.FileScanTask; import org.apache.iceberg.PartitionField; import org.apache.iceberg.Schema; -import org.apache.iceberg.StructLike; import org.apache.iceberg.Table; import org.apache.iceberg.TableScan; import org.apache.iceberg.io.CloseableIterable; import org.apache.iceberg.types.Type; -import org.apache.iceberg.types.Types; import org.apache.iceberg.types.Types.NestedField; -import org.apache.iceberg.util.StructLikeWrapper; import java.io.IOException; import java.io.UncheckedIOException; @@ -51,7 +46,6 @@ import java.util.Objects; import java.util.Optional; import java.util.Set; -import java.util.stream.IntStream; import java.util.stream.Stream; import static com.google.common.collect.ImmutableList.toImmutableList; @@ -59,6 +53,7 @@ import static io.trino.plugin.iceberg.IcebergTypes.convertIcebergValueToTrino; import static io.trino.plugin.iceberg.IcebergUtil.getIdentityPartitions; import static io.trino.plugin.iceberg.IcebergUtil.primitiveFieldTypes; +import static io.trino.plugin.iceberg.StructLikeWrapperWithFieldIdToIndex.createStructLikeWrapper; import static io.trino.plugin.iceberg.TypeConverter.toTrinoType; import static io.trino.spi.block.RowValueBuilder.buildRowValue; import static io.trino.spi.type.BigintType.BIGINT; @@ -222,14 +217,11 @@ private Map getStatistic Map partitions = new HashMap<>(); for (FileScanTask fileScanTask : fileScanTasks) { DataFile dataFile = fileScanTask.file(); - Types.StructType structType = fileScanTask.spec().partitionType(); - StructLike partitionStruct = dataFile.partition(); - StructLikeWrapper partitionWrapper = StructLikeWrapper.forType(structType).set(partitionStruct); - StructLikeWrapperWithFieldIdToIndex structLikeWrapperWithFieldIdToIndex = new StructLikeWrapperWithFieldIdToIndex(partitionWrapper, structType); + StructLikeWrapperWithFieldIdToIndex structLikeWrapperWithFieldIdToIndex = createStructLikeWrapper(fileScanTask); partitions.computeIfAbsent( structLikeWrapperWithFieldIdToIndex, - ignored -> new IcebergStatistics.Builder(icebergTable.schema().columns(), typeManager)) + ignore -> new IcebergStatistics.Builder(icebergTable.schema().columns(), typeManager)) .acceptDataFile(dataFile, fileScanTask.spec()); } @@ -265,10 +257,10 @@ private RecordCursor buildRecordCursor(Map { @@ -286,10 +278,10 @@ private RecordCursor buildRecordCursor(Map { for (int i = 0; i < columnMetricTypes.size(); i++) { Integer fieldId = nonPartitionPrimitiveColumns.get(i).fieldId(); - Object min = icebergStatistics.getMinValues().get(fieldId); - Object max = icebergStatistics.getMaxValues().get(fieldId); - Long nullCount = icebergStatistics.getNullCounts().get(fieldId); - Long nanCount = icebergStatistics.getNanCounts().get(fieldId); + Object min = icebergStatistics.minValues().get(fieldId); + Object max = icebergStatistics.maxValues().get(fieldId); + Long nullCount = icebergStatistics.nullCounts().get(fieldId); + Long nanCount = icebergStatistics.nanCounts().get(fieldId); if (min == null && max == null && nullCount == null) { throw new MissingColumnMetricsException(); } @@ -336,43 +328,6 @@ private static Block getColumnMetricBlock(RowType columnMetricType, Object min, }); } - @VisibleForTesting - static class StructLikeWrapperWithFieldIdToIndex - { - private final StructLikeWrapper structLikeWrapper; - private final Map fieldIdToIndex; - - public StructLikeWrapperWithFieldIdToIndex(StructLikeWrapper structLikeWrapper, Types.StructType structType) - { - this.structLikeWrapper = structLikeWrapper; - ImmutableMap.Builder fieldIdToIndex = ImmutableMap.builder(); - List fields = structType.fields(); - IntStream.range(0, fields.size()) - .forEach(i -> fieldIdToIndex.put(fields.get(i).fieldId(), i)); - this.fieldIdToIndex = fieldIdToIndex.buildOrThrow(); - } - - @Override - public boolean equals(Object o) - { - if (this == o) { - return true; - } - if (o == null || getClass() != o.getClass()) { - return false; - } - StructLikeWrapperWithFieldIdToIndex that = (StructLikeWrapperWithFieldIdToIndex) o; - // Due to bogus implementation of equals in StructLikeWrapper https://github.com/apache/iceberg/issues/5064 order here matters. - return Objects.equals(fieldIdToIndex, that.fieldIdToIndex) && Objects.equals(structLikeWrapper, that.structLikeWrapper); - } - - @Override - public int hashCode() - { - return Objects.hash(fieldIdToIndex, structLikeWrapper); - } - } - private static class IcebergPartitionColumn { private final RowType rowType; diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/PartitionTransforms.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/PartitionTransforms.java index f22c64aebe2d..aef88d78bdc9 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/PartitionTransforms.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/PartitionTransforms.java @@ -160,6 +160,83 @@ public static ColumnTransform getColumnTransform(PartitionField field, Type sour throw new UnsupportedOperationException("Unsupported partition transform: " + field); } + public static ColumnTransform getColumnTransform(IcebergPartitionFunction field) + { + Type type = field.type(); + return switch (field.transform()) { + case IDENTITY -> identity(type); + case YEAR -> { + if (type.equals(DATE)) { + yield yearsFromDate(); + } + if (type.equals(TIMESTAMP_MICROS)) { + yield yearsFromTimestamp(); + } + if (type.equals(TIMESTAMP_TZ_MICROS)) { + yield yearsFromTimestampWithTimeZone(); + } + throw new UnsupportedOperationException("Unsupported type for 'year': " + field); + } + case MONTH -> { + if (type.equals(DATE)) { + yield monthsFromDate(); + } + if (type.equals(TIMESTAMP_MICROS)) { + yield monthsFromTimestamp(); + } + if (type.equals(TIMESTAMP_TZ_MICROS)) { + yield monthsFromTimestampWithTimeZone(); + } + throw new UnsupportedOperationException("Unsupported type for 'month': " + field); + } + case DAY -> { + if (type.equals(DATE)) { + yield daysFromDate(); + } + if (type.equals(TIMESTAMP_MICROS)) { + yield daysFromTimestamp(); + } + if (type.equals(TIMESTAMP_TZ_MICROS)) { + yield daysFromTimestampWithTimeZone(); + } + throw new UnsupportedOperationException("Unsupported type for 'day': " + field); + } + case HOUR -> { + if (type.equals(TIMESTAMP_MICROS)) { + yield hoursFromTimestamp(); + } + if (type.equals(TIMESTAMP_TZ_MICROS)) { + yield hoursFromTimestampWithTimeZone(); + } + throw new UnsupportedOperationException("Unsupported type for 'hour': " + field); + } + case VOID -> voidTransform(type); + case BUCKET -> bucket(type, field.size().orElseThrow()); + case TRUNCATE -> { + int width = field.size().orElseThrow(); + if (type.equals(INTEGER)) { + yield truncateInteger(width); + } + if (type.equals(BIGINT)) { + yield truncateBigint(width); + } + if (type instanceof DecimalType decimalType) { + if (decimalType.isShort()) { + yield truncateShortDecimal(type, width, decimalType); + } + yield truncateLongDecimal(type, width, decimalType); + } + if (type instanceof VarcharType) { + yield truncateVarchar(width); + } + if (type.equals(VARBINARY)) { + yield truncateVarbinary(width); + } + throw new UnsupportedOperationException("Unsupported type for 'truncate': " + field); + } + }; + } + private static ColumnTransform identity(Type type) { return new ColumnTransform(type, false, true, false, Function.identity(), ValueTransform.identity(type)); diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/PartitionsTable.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/PartitionsTable.java new file mode 100644 index 000000000000..02e06935df34 --- /dev/null +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/PartitionsTable.java @@ -0,0 +1,332 @@ +/* + * 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.iceberg; + +import com.google.common.collect.ImmutableList; +import io.trino.spi.block.Block; +import io.trino.spi.connector.ColumnMetadata; +import io.trino.spi.connector.ConnectorSession; +import io.trino.spi.connector.ConnectorTableMetadata; +import io.trino.spi.connector.ConnectorTransactionHandle; +import io.trino.spi.connector.InMemoryRecordSet; +import io.trino.spi.connector.RecordCursor; +import io.trino.spi.connector.SchemaTableName; +import io.trino.spi.connector.SystemTable; +import io.trino.spi.predicate.TupleDomain; +import io.trino.spi.type.RowType; +import io.trino.spi.type.TypeManager; +import org.apache.iceberg.DataFile; +import org.apache.iceberg.FileScanTask; +import org.apache.iceberg.PartitionField; +import org.apache.iceberg.Schema; +import org.apache.iceberg.Table; +import org.apache.iceberg.TableScan; +import org.apache.iceberg.io.CloseableIterable; +import org.apache.iceberg.types.Type; +import org.apache.iceberg.types.TypeUtil; +import org.apache.iceberg.types.Types.NestedField; + +import java.io.IOException; +import java.io.UncheckedIOException; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.concurrent.ExecutorService; +import java.util.stream.Stream; + +import static com.google.common.collect.ImmutableList.toImmutableList; +import static com.google.common.collect.ImmutableMap.toImmutableMap; +import static io.trino.plugin.iceberg.IcebergTypes.convertIcebergValueToTrino; +import static io.trino.plugin.iceberg.IcebergUtil.getIdentityPartitions; +import static io.trino.plugin.iceberg.IcebergUtil.primitiveFieldTypes; +import static io.trino.plugin.iceberg.StructLikeWrapperWithFieldIdToIndex.createStructLikeWrapper; +import static io.trino.plugin.iceberg.TypeConverter.toTrinoType; +import static io.trino.spi.block.RowValueBuilder.buildRowValue; +import static io.trino.spi.type.BigintType.BIGINT; +import static io.trino.spi.type.TypeUtils.writeNativeValue; +import static java.util.Objects.requireNonNull; +import static java.util.stream.Collectors.toSet; + +public class PartitionsTable + implements SystemTable +{ + private final TypeManager typeManager; + private final Table icebergTable; + private final Optional snapshotId; + private final Map idToTypeMapping; + private final List nonPartitionPrimitiveColumns; + private final Optional partitionColumnType; + private final List partitionFields; + private final Optional dataColumnType; + private final List columnMetricTypes; + private final List resultTypes; + private final ConnectorTableMetadata connectorTableMetadata; + private final ExecutorService executor; + + public PartitionsTable(SchemaTableName tableName, TypeManager typeManager, Table icebergTable, Optional snapshotId, ExecutorService executor) + { + this.typeManager = requireNonNull(typeManager, "typeManager is null"); + this.icebergTable = requireNonNull(icebergTable, "icebergTable is null"); + this.snapshotId = requireNonNull(snapshotId, "snapshotId is null"); + this.idToTypeMapping = primitiveFieldTypes(icebergTable.schema()); + + List columns = icebergTable.schema().columns(); + this.partitionFields = getAllPartitionFields(icebergTable); + + ImmutableList.Builder columnMetadataBuilder = ImmutableList.builder(); + + this.partitionColumnType = getPartitionColumnType(partitionFields, icebergTable.schema()); + partitionColumnType.ifPresent(icebergPartitionColumn -> + columnMetadataBuilder.add(new ColumnMetadata("partition", icebergPartitionColumn.rowType))); + + Stream.of("record_count", "file_count", "total_size") + .forEach(metric -> columnMetadataBuilder.add(new ColumnMetadata(metric, BIGINT))); + + Set identityPartitionIds = getIdentityPartitions(icebergTable.spec()).keySet().stream() + .map(PartitionField::sourceId) + .collect(toSet()); + + this.nonPartitionPrimitiveColumns = columns.stream() + .filter(column -> !identityPartitionIds.contains(column.fieldId()) && column.type().isPrimitiveType()) + .collect(toImmutableList()); + + this.dataColumnType = getMetricsColumnType(this.nonPartitionPrimitiveColumns); + if (dataColumnType.isPresent()) { + columnMetadataBuilder.add(new ColumnMetadata("data", dataColumnType.get())); + this.columnMetricTypes = dataColumnType.get().getFields().stream() + .map(RowType.Field::getType) + .map(RowType.class::cast) + .collect(toImmutableList()); + } + else { + this.columnMetricTypes = ImmutableList.of(); + } + + ImmutableList columnMetadata = columnMetadataBuilder.build(); + this.resultTypes = columnMetadata.stream() + .map(ColumnMetadata::getType) + .collect(toImmutableList()); + this.connectorTableMetadata = new ConnectorTableMetadata(tableName, columnMetadata); + this.executor = requireNonNull(executor, "executor is null"); + } + + @Override + public Distribution getDistribution() + { + return Distribution.SINGLE_COORDINATOR; + } + + @Override + public ConnectorTableMetadata getTableMetadata() + { + return connectorTableMetadata; + } + + static List getAllPartitionFields(Table icebergTable) + { + Set existingColumnsIds = TypeUtil.indexById(icebergTable.schema().asStruct()).keySet(); + + List visiblePartitionFields = icebergTable.specs() + .values().stream() + .flatMap(partitionSpec -> partitionSpec.fields().stream()) + // skip columns that were dropped + .filter(partitionField -> existingColumnsIds.contains(partitionField.sourceId())) + .collect(toImmutableList()); + + return filterOutDuplicates(visiblePartitionFields); + } + + private static List filterOutDuplicates(List visiblePartitionFields) + { + Set alreadyExistingFieldIds = new HashSet<>(); + List result = new ArrayList<>(); + for (PartitionField partitionField : visiblePartitionFields) { + if (!alreadyExistingFieldIds.contains(partitionField.fieldId())) { + alreadyExistingFieldIds.add(partitionField.fieldId()); + result.add(partitionField); + } + } + return result; + } + + private Optional getPartitionColumnType(List fields, Schema schema) + { + if (fields.isEmpty()) { + return Optional.empty(); + } + List partitionFields = fields.stream() + .map(field -> RowType.field( + field.name(), + toTrinoType(field.transform().getResultType(schema.findType(field.sourceId())), typeManager))) + .collect(toImmutableList()); + List fieldIds = fields.stream() + .map(PartitionField::fieldId) + .collect(toImmutableList()); + return Optional.of(new IcebergPartitionColumn(RowType.from(partitionFields), fieldIds)); + } + + private Optional getMetricsColumnType(List columns) + { + List metricColumns = columns.stream() + .map(column -> RowType.field( + column.name(), + RowType.from(ImmutableList.of( + new RowType.Field(Optional.of("min"), toTrinoType(column.type(), typeManager)), + new RowType.Field(Optional.of("max"), toTrinoType(column.type(), typeManager)), + new RowType.Field(Optional.of("null_count"), BIGINT), + new RowType.Field(Optional.of("nan_count"), BIGINT))))) + .collect(toImmutableList()); + if (metricColumns.isEmpty()) { + return Optional.empty(); + } + return Optional.of(RowType.from(metricColumns)); + } + + @Override + public RecordCursor cursor(ConnectorTransactionHandle transactionHandle, ConnectorSession session, TupleDomain constraint) + { + if (snapshotId.isEmpty()) { + return new InMemoryRecordSet(resultTypes, ImmutableList.of()).cursor(); + } + TableScan tableScan = icebergTable.newScan() + .useSnapshot(snapshotId.get()) + .includeColumnStats() + .planWith(executor); + // TODO make the cursor lazy + return buildRecordCursor(getStatisticsByPartition(tableScan)); + } + + private Map getStatisticsByPartition(TableScan tableScan) + { + try (CloseableIterable fileScanTasks = tableScan.planFiles()) { + Map partitions = new HashMap<>(); + for (FileScanTask fileScanTask : fileScanTasks) { + DataFile dataFile = fileScanTask.file(); + StructLikeWrapperWithFieldIdToIndex structLikeWrapperWithFieldIdToIndex = createStructLikeWrapper(fileScanTask); + + partitions.computeIfAbsent( + structLikeWrapperWithFieldIdToIndex, + ignore -> new IcebergStatistics.Builder(icebergTable.schema().columns(), typeManager)) + .acceptDataFile(dataFile, fileScanTask.spec()); + } + + return partitions.entrySet().stream() + .collect(toImmutableMap(Map.Entry::getKey, entry -> entry.getValue().build())); + } + catch (IOException e) { + throw new UncheckedIOException(e); + } + } + + private RecordCursor buildRecordCursor(Map partitionStatistics) + { + List partitionTypes = partitionTypes(); + List> partitionColumnClass = partitionTypes.stream() + .map(type -> type.typeId().javaClass()) + .collect(toImmutableList()); + + ImmutableList.Builder> records = ImmutableList.builder(); + + for (Map.Entry partitionEntry : partitionStatistics.entrySet()) { + StructLikeWrapperWithFieldIdToIndex partitionStruct = partitionEntry.getKey(); + IcebergStatistics icebergStatistics = partitionEntry.getValue(); + List row = new ArrayList<>(); + + // add data for partition columns + partitionColumnType.ifPresent(partitionColumnType -> { + row.add(buildRowValue(partitionColumnType.rowType, fields -> { + List partitionColumnTypes = partitionColumnType.rowType.getFields().stream() + .map(RowType.Field::getType) + .collect(toImmutableList()); + for (int i = 0; i < partitionColumnTypes.size(); i++) { + io.trino.spi.type.Type trinoType = partitionColumnType.rowType.getFields().get(i).getType(); + Object value = null; + Integer fieldId = partitionColumnType.fieldIds.get(i); + if (partitionStruct.getFieldIdToIndex().containsKey(fieldId)) { + value = convertIcebergValueToTrino( + partitionTypes.get(i), + partitionStruct.getStructLikeWrapper().get().get(partitionStruct.getFieldIdToIndex().get(fieldId), partitionColumnClass.get(i))); + } + writeNativeValue(trinoType, fields.get(i), value); + } + })); + }); + + // add the top level metrics. + row.add(icebergStatistics.recordCount()); + row.add(icebergStatistics.fileCount()); + row.add(icebergStatistics.size()); + + // add column level metrics + dataColumnType.ifPresent(dataColumnType -> { + try { + row.add(buildRowValue(dataColumnType, fields -> { + for (int i = 0; i < columnMetricTypes.size(); i++) { + Integer fieldId = nonPartitionPrimitiveColumns.get(i).fieldId(); + Object min = icebergStatistics.minValues().get(fieldId); + Object max = icebergStatistics.maxValues().get(fieldId); + Long nullCount = icebergStatistics.nullCounts().get(fieldId); + Long nanCount = icebergStatistics.nanCounts().get(fieldId); + if (min == null && max == null && nullCount == null) { + throw new MissingColumnMetricsException(); + } + + RowType columnMetricType = columnMetricTypes.get(i); + columnMetricType.writeObject(fields.get(i), getColumnMetricBlock(columnMetricType, min, max, nullCount, nanCount)); + } + })); + } + catch (MissingColumnMetricsException ignore) { + row.add(null); + } + }); + + records.add(row); + } + + return new InMemoryRecordSet(resultTypes, records.build()).cursor(); + } + + private static class MissingColumnMetricsException + extends Exception + {} + + private List partitionTypes() + { + ImmutableList.Builder partitionTypeBuilder = ImmutableList.builder(); + for (PartitionField partitionField : partitionFields) { + Type.PrimitiveType sourceType = idToTypeMapping.get(partitionField.sourceId()); + Type type = partitionField.transform().getResultType(sourceType); + partitionTypeBuilder.add(type); + } + return partitionTypeBuilder.build(); + } + + private static Block getColumnMetricBlock(RowType columnMetricType, Object min, Object max, Long nullCount, Long nanCount) + { + return buildRowValue(columnMetricType, fieldBuilders -> { + List fields = columnMetricType.getFields(); + writeNativeValue(fields.get(0).getType(), fieldBuilders.get(0), min); + writeNativeValue(fields.get(1).getType(), fieldBuilders.get(1), max); + writeNativeValue(fields.get(2).getType(), fieldBuilders.get(2), nullCount); + writeNativeValue(fields.get(3).getType(), fieldBuilders.get(3), nanCount); + }); + } + + private record IcebergPartitionColumn(RowType rowType, List fieldIds) {} +} diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/RefsTable.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/RefsTable.java index b08b6d3b7f7c..38f105281560 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/RefsTable.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/RefsTable.java @@ -15,79 +15,52 @@ import com.google.common.collect.ImmutableList; import io.trino.plugin.iceberg.util.PageListBuilder; -import io.trino.spi.Page; import io.trino.spi.connector.ColumnMetadata; -import io.trino.spi.connector.ConnectorPageSource; -import io.trino.spi.connector.ConnectorSession; import io.trino.spi.connector.ConnectorTableMetadata; -import io.trino.spi.connector.ConnectorTransactionHandle; -import io.trino.spi.connector.FixedPageSource; import io.trino.spi.connector.SchemaTableName; -import io.trino.spi.connector.SystemTable; -import io.trino.spi.predicate.TupleDomain; +import io.trino.spi.type.TimeZoneKey; import org.apache.iceberg.Table; import java.util.List; +import java.util.concurrent.ExecutorService; import static io.trino.spi.type.BigintType.BIGINT; import static io.trino.spi.type.IntegerType.INTEGER; import static io.trino.spi.type.VarcharType.VARCHAR; import static java.util.Objects.requireNonNull; +import static org.apache.iceberg.MetadataTableType.REFS; public class RefsTable - implements SystemTable + extends BaseSystemTable { - private final ConnectorTableMetadata tableMetadata; - private final Table icebergTable; + private static final List COLUMNS = ImmutableList.builder() + .add(new ColumnMetadata("name", VARCHAR)) + .add(new ColumnMetadata("type", VARCHAR)) + .add(new ColumnMetadata("snapshot_id", BIGINT)) + .add(new ColumnMetadata("max_reference_age_in_ms", BIGINT)) + .add(new ColumnMetadata("min_snapshots_to_keep", INTEGER)) + .add(new ColumnMetadata("max_snapshot_age_in_ms", BIGINT)) + .build(); - public RefsTable(SchemaTableName tableName, Table icebergTable) + public RefsTable(SchemaTableName tableName, Table icebergTable, ExecutorService executor) { - this.icebergTable = requireNonNull(icebergTable, "icebergTable is null"); - - this.tableMetadata = new ConnectorTableMetadata(requireNonNull(tableName, "tableName is null"), - ImmutableList.builder() - .add(new ColumnMetadata("name", VARCHAR)) - .add(new ColumnMetadata("type", VARCHAR)) - .add(new ColumnMetadata("snapshot_id", BIGINT)) - .add(new ColumnMetadata("max_reference_age_in_ms", BIGINT)) - .add(new ColumnMetadata("min_snapshots_to_keep", INTEGER)) - .add(new ColumnMetadata("max_snapshot_age_in_ms", BIGINT)) - .build()); - } - - @Override - public Distribution getDistribution() - { - return Distribution.SINGLE_COORDINATOR; - } - - @Override - public ConnectorTableMetadata getTableMetadata() - { - return tableMetadata; + super( + requireNonNull(icebergTable, "icebergTable is null"), + new ConnectorTableMetadata(requireNonNull(tableName, "tableName is null"), COLUMNS), + REFS, + executor); } @Override - public ConnectorPageSource pageSource(ConnectorTransactionHandle transactionHandle, ConnectorSession session, TupleDomain constraint) + protected void addRow(PageListBuilder pagesBuilder, Row row, TimeZoneKey timeZoneKey) { - return new FixedPageSource(buildPages(tableMetadata, icebergTable)); - } - - private static List buildPages(ConnectorTableMetadata tableMetadata, Table icebergTable) - { - PageListBuilder pagesBuilder = PageListBuilder.forTable(tableMetadata); - - icebergTable.refs().forEach((refName, ref) -> { - pagesBuilder.beginRow(); - pagesBuilder.appendVarchar(refName); - pagesBuilder.appendVarchar(ref.isBranch() ? "BRANCH" : "TAG"); - pagesBuilder.appendBigint(ref.snapshotId()); - pagesBuilder.appendBigint(ref.maxRefAgeMs()); - pagesBuilder.appendInteger(ref.minSnapshotsToKeep()); - pagesBuilder.appendBigint(ref.maxSnapshotAgeMs()); - pagesBuilder.endRow(); - }); - - return pagesBuilder.build(); + pagesBuilder.beginRow(); + pagesBuilder.appendVarchar(row.get("name", String.class)); + pagesBuilder.appendVarchar(row.get("type", String.class)); + pagesBuilder.appendBigint(row.get("snapshot_id", Long.class)); + pagesBuilder.appendBigint(row.get("max_reference_age_in_ms", Long.class)); + pagesBuilder.appendInteger(row.get("min_snapshots_to_keep", Integer.class)); + pagesBuilder.appendBigint(row.get("max_snapshot_age_in_ms", Long.class)); + pagesBuilder.endRow(); } } diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/SnapshotsTable.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/SnapshotsTable.java index 81f9b60699dc..74204e6b4b2b 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/SnapshotsTable.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/SnapshotsTable.java @@ -15,33 +15,17 @@ import com.google.common.collect.ImmutableList; import io.trino.plugin.iceberg.util.PageListBuilder; -import io.trino.spi.Page; import io.trino.spi.connector.ColumnMetadata; -import io.trino.spi.connector.ConnectorPageSource; -import io.trino.spi.connector.ConnectorSession; import io.trino.spi.connector.ConnectorTableMetadata; -import io.trino.spi.connector.ConnectorTransactionHandle; -import io.trino.spi.connector.FixedPageSource; import io.trino.spi.connector.SchemaTableName; -import io.trino.spi.connector.SystemTable; -import io.trino.spi.predicate.TupleDomain; import io.trino.spi.type.TimeZoneKey; import io.trino.spi.type.TypeManager; import io.trino.spi.type.TypeSignature; -import org.apache.iceberg.DataTask; -import org.apache.iceberg.FileScanTask; -import org.apache.iceberg.StructLike; import org.apache.iceberg.Table; -import org.apache.iceberg.TableScan; -import org.apache.iceberg.io.CloseableIterable; -import java.io.IOException; -import java.io.UncheckedIOException; -import java.util.List; import java.util.Map; +import java.util.concurrent.ExecutorService; -import static io.trino.plugin.iceberg.IcebergUtil.buildTableScan; -import static io.trino.plugin.iceberg.IcebergUtil.columnNameToPositionInSchema; import static io.trino.spi.type.BigintType.BIGINT; import static io.trino.spi.type.TimestampWithTimeZoneType.TIMESTAMP_TZ_MILLIS; import static io.trino.spi.type.Timestamps.MICROSECONDS_PER_MILLISECOND; @@ -50,10 +34,8 @@ import static org.apache.iceberg.MetadataTableType.SNAPSHOTS; public class SnapshotsTable - implements SystemTable + extends BaseSystemTable { - private final ConnectorTableMetadata tableMetadata; - private final Table icebergTable; private static final String COMMITTED_AT_COLUMN_NAME = "committed_at"; private static final String SNAPSHOT_ID_COLUMN_NAME = "snapshot_id"; private static final String PARENT_ID_COLUMN_NAME = "parent_id"; @@ -61,12 +43,21 @@ public class SnapshotsTable private static final String MANIFEST_LIST_COLUMN_NAME = "manifest_list"; private static final String SUMMARY_COLUMN_NAME = "summary"; - public SnapshotsTable(SchemaTableName tableName, TypeManager typeManager, Table icebergTable) + public SnapshotsTable(SchemaTableName tableName, TypeManager typeManager, Table icebergTable, ExecutorService executor) { - requireNonNull(typeManager, "typeManager is null"); + super( + requireNonNull(icebergTable, "icebergTable is null"), + createConnectorTableMetadata( + requireNonNull(tableName, "tableName is null"), + requireNonNull(typeManager, "typeManager is null")), + SNAPSHOTS, + executor); + } - this.icebergTable = requireNonNull(icebergTable, "icebergTable is null"); - tableMetadata = new ConnectorTableMetadata(requireNonNull(tableName, "tableName is null"), + private static ConnectorTableMetadata createConnectorTableMetadata(SchemaTableName tableName, TypeManager typeManager) + { + return new ConnectorTableMetadata( + tableName, ImmutableList.builder() .add(new ColumnMetadata(COMMITTED_AT_COLUMN_NAME, TIMESTAMP_TZ_MILLIS)) .add(new ColumnMetadata(SNAPSHOT_ID_COLUMN_NAME, BIGINT)) @@ -78,67 +69,16 @@ public SnapshotsTable(SchemaTableName tableName, TypeManager typeManager, Table } @Override - public Distribution getDistribution() - { - return Distribution.SINGLE_COORDINATOR; - } - - @Override - public ConnectorTableMetadata getTableMetadata() - { - return tableMetadata; - } - - @Override - public ConnectorPageSource pageSource(ConnectorTransactionHandle transactionHandle, ConnectorSession session, TupleDomain constraint) - { - return new FixedPageSource(buildPages(tableMetadata, session, icebergTable)); - } - - private static List buildPages(ConnectorTableMetadata tableMetadata, ConnectorSession session, Table icebergTable) - { - PageListBuilder pagesBuilder = PageListBuilder.forTable(tableMetadata); - - TableScan tableScan = buildTableScan(icebergTable, SNAPSHOTS); - TimeZoneKey timeZoneKey = session.getTimeZoneKey(); - - Map columnNameToPosition = columnNameToPositionInSchema(tableScan.schema()); - - try (CloseableIterable fileScanTasks = tableScan.planFiles()) { - fileScanTasks.forEach(fileScanTask -> addRows((DataTask) fileScanTask, pagesBuilder, timeZoneKey, columnNameToPosition)); - } - catch (IOException e) { - throw new UncheckedIOException(e); - } - - return pagesBuilder.build(); - } - - private static void addRows(DataTask dataTask, PageListBuilder pagesBuilder, TimeZoneKey timeZoneKey, Map columnNameToPositionInSchema) - { - try (CloseableIterable dataRows = dataTask.rows()) { - dataRows.forEach(dataTaskRow -> addRow(pagesBuilder, dataTaskRow, timeZoneKey, columnNameToPositionInSchema)); - } - catch (IOException e) { - throw new UncheckedIOException(e); - } - } - - private static void addRow(PageListBuilder pagesBuilder, StructLike structLike, TimeZoneKey timeZoneKey, Map columnNameToPositionInSchema) + protected void addRow(PageListBuilder pagesBuilder, Row row, TimeZoneKey timeZoneKey) { pagesBuilder.beginRow(); - - pagesBuilder.appendTimestampTzMillis( - structLike.get(columnNameToPositionInSchema.get(COMMITTED_AT_COLUMN_NAME), Long.class) / MICROSECONDS_PER_MILLISECOND, - timeZoneKey); - pagesBuilder.appendBigint(structLike.get(columnNameToPositionInSchema.get(SNAPSHOT_ID_COLUMN_NAME), Long.class)); - - Long parentId = structLike.get(columnNameToPositionInSchema.get(PARENT_ID_COLUMN_NAME), Long.class); - pagesBuilder.appendBigint(parentId != null ? parentId.longValue() : null); - - pagesBuilder.appendVarchar(structLike.get(columnNameToPositionInSchema.get(OPERATION_COLUMN_NAME), String.class)); - pagesBuilder.appendVarchar(structLike.get(columnNameToPositionInSchema.get(MANIFEST_LIST_COLUMN_NAME), String.class)); - pagesBuilder.appendVarcharVarcharMap(structLike.get(columnNameToPositionInSchema.get(SUMMARY_COLUMN_NAME), Map.class)); + pagesBuilder.appendTimestampTzMillis(row.get(COMMITTED_AT_COLUMN_NAME, Long.class) / MICROSECONDS_PER_MILLISECOND, timeZoneKey); + pagesBuilder.appendBigint(row.get(SNAPSHOT_ID_COLUMN_NAME, Long.class)); + pagesBuilder.appendBigint(row.get(PARENT_ID_COLUMN_NAME, Long.class)); + pagesBuilder.appendVarchar(row.get(OPERATION_COLUMN_NAME, String.class)); + pagesBuilder.appendVarchar(row.get(MANIFEST_LIST_COLUMN_NAME, String.class)); + //noinspection unchecked + pagesBuilder.appendVarcharVarcharMap(row.get(SUMMARY_COLUMN_NAME, Map.class)); pagesBuilder.endRow(); } } diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/StructLikeWrapperWithFieldIdToIndex.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/StructLikeWrapperWithFieldIdToIndex.java new file mode 100644 index 000000000000..4fa5c08f1a9c --- /dev/null +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/StructLikeWrapperWithFieldIdToIndex.java @@ -0,0 +1,79 @@ +/* + * 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.iceberg; + +import com.google.common.annotations.VisibleForTesting; +import com.google.common.collect.ImmutableMap; +import org.apache.iceberg.FileScanTask; +import org.apache.iceberg.types.Types; +import org.apache.iceberg.util.StructLikeWrapper; + +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.stream.IntStream; + +public class StructLikeWrapperWithFieldIdToIndex +{ + private final StructLikeWrapper structLikeWrapper; + private final Map fieldIdToIndex; + + public static StructLikeWrapperWithFieldIdToIndex createStructLikeWrapper(FileScanTask fileScanTask) + { + Types.StructType structType = fileScanTask.spec().partitionType(); + StructLikeWrapper partitionWrapper = StructLikeWrapper.forType(structType).set(fileScanTask.file().partition()); + return new StructLikeWrapperWithFieldIdToIndex(partitionWrapper, structType); + } + + @VisibleForTesting + StructLikeWrapperWithFieldIdToIndex(StructLikeWrapper structLikeWrapper, Types.StructType structType) + { + this.structLikeWrapper = structLikeWrapper; + ImmutableMap.Builder fieldIdToIndex = ImmutableMap.builder(); + List fields = structType.fields(); + IntStream.range(0, fields.size()) + .forEach(i -> fieldIdToIndex.put(fields.get(i).fieldId(), i)); + this.fieldIdToIndex = fieldIdToIndex.buildOrThrow(); + } + + public StructLikeWrapper getStructLikeWrapper() + { + return structLikeWrapper; + } + + public Map getFieldIdToIndex() + { + return fieldIdToIndex; + } + + @Override + public boolean equals(Object o) + { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + StructLikeWrapperWithFieldIdToIndex that = (StructLikeWrapperWithFieldIdToIndex) o; + // Due to bogus implementation of equals in StructLikeWrapper https://github.com/apache/iceberg/issues/5064 order here matters. + return Objects.equals(fieldIdToIndex, that.fieldIdToIndex) && Objects.equals(structLikeWrapper, that.structLikeWrapper); + } + + @Override + public int hashCode() + { + return Objects.hash(fieldIdToIndex, structLikeWrapper); + } +} diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/TableStatisticsReader.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/TableStatisticsReader.java index 8e00d7e0c3da..feb8dbb50b8b 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/TableStatisticsReader.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/TableStatisticsReader.java @@ -18,9 +18,11 @@ import com.google.common.collect.ImmutableMap; import com.google.common.collect.Maps; import io.airlift.log.Logger; +import io.trino.filesystem.TrinoFileSystem; import io.trino.spi.TrinoException; import io.trino.spi.connector.ColumnHandle; import io.trino.spi.connector.ConnectorSession; +import io.trino.spi.predicate.Domain; import io.trino.spi.predicate.TupleDomain; import io.trino.spi.statistics.ColumnStatistics; import io.trino.spi.statistics.DoubleRange; @@ -38,12 +40,10 @@ import org.apache.iceberg.TableScan; import org.apache.iceberg.io.CloseableIterable; import org.apache.iceberg.puffin.StandardBlobTypes; -import org.apache.iceberg.types.Type; import org.apache.iceberg.types.Types; import java.io.IOException; import java.io.UncheckedIOException; -import java.util.HashSet; import java.util.Iterator; import java.util.List; import java.util.Map; @@ -51,18 +51,23 @@ import java.util.Objects; import java.util.Optional; import java.util.Set; -import java.util.regex.Matcher; -import java.util.regex.Pattern; import static com.google.common.base.Verify.verifyNotNull; +import static com.google.common.collect.ImmutableList.toImmutableList; import static com.google.common.collect.ImmutableMap.toImmutableMap; +import static com.google.common.collect.ImmutableSet.toImmutableSet; import static com.google.common.collect.Iterables.getOnlyElement; import static com.google.common.collect.Streams.stream; +import static io.airlift.slice.Slices.utf8Slice; import static io.trino.plugin.iceberg.ExpressionConverter.toIcebergExpression; import static io.trino.plugin.iceberg.IcebergErrorCode.ICEBERG_INVALID_METADATA; import static io.trino.plugin.iceberg.IcebergMetadataColumn.isMetadataColumnId; import static io.trino.plugin.iceberg.IcebergSessionProperties.isExtendedStatisticsEnabled; -import static io.trino.plugin.iceberg.IcebergUtil.getColumns; +import static io.trino.plugin.iceberg.IcebergUtil.getFileModifiedTimePathDomain; +import static io.trino.plugin.iceberg.IcebergUtil.getModificationTime; +import static io.trino.plugin.iceberg.IcebergUtil.getPathDomain; +import static io.trino.spi.type.DateTimeEncoding.packDateTimeWithZone; +import static io.trino.spi.type.TimeZoneKey.UTC_KEY; import static io.trino.spi.type.VarbinaryType.VARBINARY; import static io.trino.spi.type.VarcharType.VARCHAR; import static java.lang.Long.parseLong; @@ -70,6 +75,7 @@ import static java.util.function.Function.identity; import static java.util.stream.Collectors.toMap; import static java.util.stream.Collectors.toUnmodifiableMap; +import static org.apache.iceberg.util.SnapshotUtil.schemaFor; public final class TableStatisticsReader { @@ -77,22 +83,15 @@ private TableStatisticsReader() {} private static final Logger log = Logger.get(TableStatisticsReader.class); - // TODO (https://github.com/trinodb/trino/issues/15397): remove support for Trino-specific statistics properties - @Deprecated - public static final String TRINO_STATS_PREFIX = "trino.stats.ndv."; - // TODO (https://github.com/trinodb/trino/issues/15397): remove support for Trino-specific statistics properties - @Deprecated - public static final String TRINO_STATS_NDV_FORMAT = TRINO_STATS_PREFIX + "%d.ndv"; - // TODO (https://github.com/trinodb/trino/issues/15397): remove support for Trino-specific statistics properties - @Deprecated - public static final Pattern TRINO_STATS_COLUMN_ID_PATTERN = Pattern.compile(Pattern.quote(TRINO_STATS_PREFIX) + "(?\\d+)\\..*"); - // TODO (https://github.com/trinodb/trino/issues/15397): remove support for Trino-specific statistics properties - @Deprecated - public static final Pattern TRINO_STATS_NDV_PATTERN = Pattern.compile(Pattern.quote(TRINO_STATS_PREFIX) + "(?\\d+)\\.ndv"); - public static final String APACHE_DATASKETCHES_THETA_V1_NDV_PROPERTY = "ndv"; - public static TableStatistics getTableStatistics(TypeManager typeManager, ConnectorSession session, IcebergTableHandle tableHandle, Table icebergTable) + public static TableStatistics getTableStatistics( + TypeManager typeManager, + ConnectorSession session, + IcebergTableHandle tableHandle, + Set projectedColumns, + Table icebergTable, + TrinoFileSystem fileSystem) { return makeTableStatistics( typeManager, @@ -100,7 +99,9 @@ public static TableStatistics getTableStatistics(TypeManager typeManager, Connec tableHandle.getSnapshotId(), tableHandle.getEnforcedPredicate(), tableHandle.getUnenforcedPredicate(), - isExtendedStatisticsEnabled(session)); + projectedColumns, + isExtendedStatisticsEnabled(session), + fileSystem); } @VisibleForTesting @@ -110,7 +111,9 @@ public static TableStatistics makeTableStatistics( Optional snapshot, TupleDomain enforcedConstraint, TupleDomain unenforcedConstraint, - boolean extendedStatisticsEnabled) + Set projectedColumns, + boolean extendedStatisticsEnabled, + TrinoFileSystem fileSystem) { if (snapshot.isEmpty()) { // No snapshot, so no data. @@ -130,25 +133,42 @@ public static TableStatistics makeTableStatistics( .build(); } - Schema icebergTableSchema = icebergTable.schema(); - List columns = icebergTableSchema.columns(); - - List columnHandles = getColumns(icebergTableSchema, typeManager); - Map idToColumnHandle = columnHandles.stream() - .collect(toUnmodifiableMap(IcebergColumnHandle::getId, identity())); + List columns = icebergTable.schema().columns(); Map idToType = columns.stream() .map(column -> Maps.immutableEntry(column.fieldId(), column.type())) .collect(toUnmodifiableMap(Map.Entry::getKey, Map.Entry::getValue)); + Set columnIds = projectedColumns.stream() + .map(IcebergColumnHandle::getId) + .collect(toImmutableSet()); + + Domain pathDomain = getPathDomain(effectivePredicate); + Domain fileModifiedTimeDomain = getFileModifiedTimePathDomain(effectivePredicate); + Schema snapshotSchema = schemaFor(icebergTable, snapshotId); TableScan tableScan = icebergTable.newScan() - // Table enforced constraint may include eg $path column predicate which is not handled by Iceberg library TODO apply $path and $file_modified_time filters here .filter(toIcebergExpression(effectivePredicate.filter((column, domain) -> !isMetadataColumnId(column.getId())))) .useSnapshot(snapshotId) - .includeColumnStats(); + .includeColumnStats( + columnIds.stream() + .map(snapshotSchema::findColumnName) + .filter(Objects::nonNull) + .collect(toImmutableList())); IcebergStatistics.Builder icebergStatisticsBuilder = new IcebergStatistics.Builder(columns, typeManager); try (CloseableIterable fileScanTasks = tableScan.planFiles()) { - fileScanTasks.forEach(fileScanTask -> icebergStatisticsBuilder.acceptDataFile(fileScanTask.file(), fileScanTask.spec())); + fileScanTasks.forEach(fileScanTask -> { + if (!pathDomain.isAll() && !pathDomain.includesNullableValue(utf8Slice(fileScanTask.file().location()))) { + return; + } + if (!fileModifiedTimeDomain.isAll()) { + long fileModifiedTime = getModificationTime(fileScanTask.file().location(), fileSystem); + if (!fileModifiedTimeDomain.includesNullableValue(packDateTimeWithZone(fileModifiedTime, UTC_KEY))) { + return; + } + } + + icebergStatisticsBuilder.acceptDataFile(fileScanTask.file(), fileScanTask.spec()); + }); } catch (IOException e) { throw new UncheckedIOException(e); @@ -156,7 +176,7 @@ public static TableStatistics makeTableStatistics( IcebergStatistics summary = icebergStatisticsBuilder.build(); - if (summary.getFileCount() == 0) { + if (summary.fileCount() == 0) { return TableStatistics.builder() .setRowCount(Estimate.of(0)) .build(); @@ -165,28 +185,24 @@ public static TableStatistics makeTableStatistics( Map ndvs = readNdvs( icebergTable, snapshotId, - // TODO We don't need NDV information for columns not involved in filters/joins. Engine should provide set of columns - // it makes sense to find NDV information for. - idToColumnHandle.keySet(), + columnIds, extendedStatisticsEnabled); ImmutableMap.Builder columnHandleBuilder = ImmutableMap.builder(); - double recordCount = summary.getRecordCount(); - for (Entry columnHandleTuple : idToColumnHandle.entrySet()) { - IcebergColumnHandle columnHandle = columnHandleTuple.getValue(); + double recordCount = summary.recordCount(); + for (IcebergColumnHandle columnHandle : projectedColumns) { int fieldId = columnHandle.getId(); ColumnStatistics.Builder columnBuilder = new ColumnStatistics.Builder(); - Long nullCount = summary.getNullCounts().get(fieldId); + Long nullCount = summary.nullCounts().get(fieldId); if (nullCount != null) { columnBuilder.setNullsFraction(Estimate.of(nullCount / recordCount)); } - if (idToType.get(columnHandleTuple.getKey()).typeId() == Type.TypeID.FIXED) { - Types.FixedType fixedType = (Types.FixedType) idToType.get(columnHandleTuple.getKey()); + if (idToType.get(columnHandle.getId()) instanceof Types.FixedType fixedType) { long columnSize = fixedType.length(); columnBuilder.setDataSize(Estimate.of(columnSize)); } - else if (summary.getColumnSizes() != null) { - Long columnSize = summary.getColumnSizes().get(fieldId); + else if (summary.columnSizes() != null) { + Long columnSize = summary.columnSizes().get(fieldId); if (columnSize != null) { // columnSize is the size on disk and Trino column stats is size in memory. // The relation between the two is type and data dependent. @@ -210,8 +226,8 @@ else if (columnHandle.getBaseType() == VARBINARY) { } } } - Object min = summary.getMinValues().get(fieldId); - Object max = summary.getMaxValues().get(fieldId); + Object min = summary.minValues().get(fieldId); + Object max = summary.maxValues().get(fieldId); if (min != null && max != null) { columnBuilder.setRange(DoubleRange.from(columnHandle.getType(), min, max)); } @@ -224,20 +240,19 @@ else if (columnHandle.getBaseType() == VARBINARY) { return new TableStatistics(Estimate.of(recordCount), columnHandleBuilder.buildOrThrow()); } - private static Map readNdvs(Table icebergTable, long snapshotId, Set columnIds, boolean extendedStatisticsEnabled) + public static Map readNdvs(Table icebergTable, long snapshotId, Set columnIds, boolean extendedStatisticsEnabled) { if (!extendedStatisticsEnabled) { return ImmutableMap.of(); } ImmutableMap.Builder ndvByColumnId = ImmutableMap.builder(); - Set remainingColumnIds = new HashSet<>(columnIds); getLatestStatisticsFile(icebergTable, snapshotId).ifPresent(statisticsFile -> { Map thetaBlobsByFieldId = statisticsFile.blobMetadata().stream() .filter(blobMetadata -> blobMetadata.type().equals(StandardBlobTypes.APACHE_DATASKETCHES_THETA_V1)) .filter(blobMetadata -> blobMetadata.fields().size() == 1) - .filter(blobMetadata -> remainingColumnIds.contains(getOnlyElement(blobMetadata.fields()))) + .filter(blobMetadata -> columnIds.contains(getOnlyElement(blobMetadata.fields()))) // Fail loud upon duplicates (there must be none) .collect(toImmutableMap(blobMetadata -> getOnlyElement(blobMetadata.fields()), identity())); @@ -247,33 +262,13 @@ private static Map readNdvs(Table icebergTable, long snapshotId, String ndv = blobMetadata.properties().get(APACHE_DATASKETCHES_THETA_V1_NDV_PROPERTY); if (ndv == null) { log.debug("Blob %s is missing %s property", blobMetadata.type(), APACHE_DATASKETCHES_THETA_V1_NDV_PROPERTY); - remainingColumnIds.remove(fieldId); } else { - remainingColumnIds.remove(fieldId); ndvByColumnId.put(fieldId, parseLong(ndv)); } } }); - // TODO (https://github.com/trinodb/trino/issues/15397): remove support for Trino-specific statistics properties - Iterator> properties = icebergTable.properties().entrySet().iterator(); - while (!remainingColumnIds.isEmpty() && properties.hasNext()) { - Entry entry = properties.next(); - String key = entry.getKey(); - String value = entry.getValue(); - if (key.startsWith(TRINO_STATS_PREFIX)) { - Matcher matcher = TRINO_STATS_NDV_PATTERN.matcher(key); - if (matcher.matches()) { - int columnId = Integer.parseInt(matcher.group("columnId")); - if (remainingColumnIds.remove(columnId)) { - long ndv = parseLong(value); - ndvByColumnId.put(columnId, ndv); - } - } - } - } - return ndvByColumnId.buildOrThrow(); } diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/TableStatisticsWriter.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/TableStatisticsWriter.java index 473df30b3316..69674474326c 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/TableStatisticsWriter.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/TableStatisticsWriter.java @@ -98,12 +98,8 @@ public StatisticsFile writeStatisticsFile( StatsUpdateMode updateMode, CollectedStatistics collectedStatistics) { - Snapshot snapshot = table.snapshot(snapshotId); TableOperations operations = ((HasTableOperations) table).operations(); FileIO fileIO = operations.io(); - long snapshotSequenceNumber = snapshot.sequenceNumber(); - Schema schema = table.schemas().get(snapshot.schemaId()); - collectedStatistics = mergeStatisticsIfNecessary( table, snapshotId, @@ -112,6 +108,23 @@ public StatisticsFile writeStatisticsFile( collectedStatistics); Map ndvSketches = collectedStatistics.ndvSketches(); + return writeStatisticsFile(session, table, fileIO, snapshotId, ndvSketches); + } + + public StatisticsFile rewriteStatisticsFile(ConnectorSession session, Table table, long snapshotId) + { + TableOperations operations = ((HasTableOperations) table).operations(); + FileIO fileIO = operations.io(); + // This will rewrite old statistics file as ndvSketches map is empty + return writeStatisticsFile(session, table, fileIO, snapshotId, Map.of()); + } + + private GenericStatisticsFile writeStatisticsFile(ConnectorSession session, Table table, FileIO fileIO, long snapshotId, Map ndvSketches) + { + Snapshot snapshot = table.snapshot(snapshotId); + long snapshotSequenceNumber = snapshot.sequenceNumber(); + TableOperations operations = ((HasTableOperations) table).operations(); + Schema schema = table.schemas().get(snapshot.schemaId()); Set validFieldIds = stream( Traverser.forTree((Types.NestedField nestedField) -> { Type type = nestedField.type(); diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/TableType.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/TableType.java index 7141f488d7da..14dcbe856e82 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/TableType.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/TableType.java @@ -17,10 +17,15 @@ public enum TableType { DATA, HISTORY, + METADATA_LOG_ENTRIES, SNAPSHOTS, + ALL_MANIFESTS, MANIFESTS, PARTITIONS, FILES, + ALL_ENTRIES, + ENTRIES, PROPERTIES, - REFS + REFS, + MATERIALIZED_VIEW_STORAGE, } diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/TypeConverter.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/TypeConverter.java index 5840560106c3..f94a693cc4e9 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/TypeConverter.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/TypeConverter.java @@ -45,19 +45,23 @@ import org.apache.iceberg.types.Types; import java.util.ArrayList; +import java.util.HashSet; import java.util.List; import java.util.Map; import java.util.Optional; +import java.util.Set; import java.util.concurrent.atomic.AtomicInteger; import static com.google.common.collect.ImmutableList.toImmutableList; import static com.google.common.collect.Iterables.getOnlyElement; +import static io.trino.spi.StandardErrorCode.DUPLICATE_COLUMN_NAME; import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED; import static io.trino.spi.type.TimeType.TIME_MICROS; import static io.trino.spi.type.TimestampType.TIMESTAMP_MICROS; import static io.trino.spi.type.TimestampWithTimeZoneType.TIMESTAMP_TZ_MICROS; import static io.trino.spi.type.UuidType.UUID; import static java.lang.String.format; +import static java.util.Locale.ENGLISH; public final class TypeConverter { @@ -195,6 +199,7 @@ private static org.apache.iceberg.types.Type fromRow(RowType type, Optional fieldNames = new HashSet<>(); List fields = new ArrayList<>(); for (int i = 0; i < type.getFields().size(); i++) { int fieldIndex = i; @@ -206,6 +211,9 @@ private static org.apache.iceberg.types.Type fromRow(RowType type, Optional new TrinoException(NOT_SUPPORTED, "Row type field does not have a name: " + type.getDisplayName())); + if (!fieldNames.add(name.toLowerCase(ENGLISH))) { + throw new TrinoException(DUPLICATE_COLUMN_NAME, "Field name '%s' specified more than once".formatted(name.toLowerCase(ENGLISH))); + } fields.add(Types.NestedField.optional(id, name, toIcebergTypeInternal(field.getType(), childColumnIdentity, nextFieldId))); } return Types.StructType.of(fields); diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/AbstractIcebergTableOperations.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/AbstractIcebergTableOperations.java index 36d09f2d0d19..e56197e7e165 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/AbstractIcebergTableOperations.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/AbstractIcebergTableOperations.java @@ -17,8 +17,11 @@ import dev.failsafe.RetryPolicy; import io.trino.annotation.NotThreadSafe; import io.trino.filesystem.Location; +import io.trino.filesystem.TrinoFileSystem; +import io.trino.plugin.hive.HiveType; import io.trino.plugin.hive.metastore.Column; import io.trino.plugin.hive.metastore.StorageFormat; +import io.trino.plugin.iceberg.IcebergExceptions; import io.trino.plugin.iceberg.util.HiveSchemaUtil; import io.trino.spi.TrinoException; import io.trino.spi.connector.ConnectorSession; @@ -27,27 +30,28 @@ import org.apache.iceberg.TableMetadata; import org.apache.iceberg.TableMetadataParser; import org.apache.iceberg.exceptions.CommitFailedException; -import org.apache.iceberg.exceptions.ValidationException; import org.apache.iceberg.io.FileIO; import org.apache.iceberg.io.LocationProvider; import org.apache.iceberg.io.OutputFile; import org.apache.iceberg.types.Types.NestedField; -import java.io.FileNotFoundException; +import java.io.UncheckedIOException; import java.time.Duration; import java.util.List; +import java.util.Map; import java.util.Objects; import java.util.Optional; import java.util.OptionalInt; +import java.util.function.Function; import static com.google.common.base.Preconditions.checkState; import static com.google.common.collect.ImmutableList.toImmutableList; -import static io.trino.plugin.hive.HiveType.toHiveType; import static io.trino.plugin.hive.util.HiveClassNames.FILE_INPUT_FORMAT_CLASS; import static io.trino.plugin.hive.util.HiveClassNames.FILE_OUTPUT_FORMAT_CLASS; import static io.trino.plugin.hive.util.HiveClassNames.LAZY_SIMPLE_SERDE_CLASS; import static io.trino.plugin.iceberg.IcebergErrorCode.ICEBERG_INVALID_METADATA; -import static io.trino.plugin.iceberg.IcebergErrorCode.ICEBERG_MISSING_METADATA; +import static io.trino.plugin.iceberg.IcebergExceptions.translateMetadataException; +import static io.trino.plugin.iceberg.IcebergTableName.isMaterializedViewStorage; import static io.trino.plugin.iceberg.IcebergUtil.METADATA_FOLDER_NAME; import static io.trino.plugin.iceberg.IcebergUtil.fixBrokenMetadataLocation; import static io.trino.plugin.iceberg.IcebergUtil.getLocationProvider; @@ -59,6 +63,7 @@ import static java.util.Objects.requireNonNull; import static java.util.UUID.randomUUID; import static org.apache.iceberg.BaseMetastoreTableOperations.METADATA_LOCATION_PROP; +import static org.apache.iceberg.CatalogUtil.deleteRemovedMetadataFiles; import static org.apache.iceberg.TableMetadataParser.getFileExtension; import static org.apache.iceberg.TableProperties.METADATA_COMPRESSION; import static org.apache.iceberg.TableProperties.METADATA_COMPRESSION_DEFAULT; @@ -152,6 +157,11 @@ public void commit(@Nullable TableMetadata base, TableMetadata metadata) return; } + if (isMaterializedViewStorage(tableName)) { + commitMaterializedViewRefresh(base, metadata); + return; + } + if (base == null) { if (PROVIDER_PROPERTY_VALUE.equals(metadata.properties().get(PROVIDER_PROPERTY_KEY))) { // Assume this is a table executing migrate procedure @@ -165,6 +175,7 @@ public void commit(@Nullable TableMetadata base, TableMetadata metadata) } else { commitToExistingTable(base, metadata); + deleteRemovedMetadataFiles(fileIo, base, metadata); } shouldRefresh = true; @@ -176,6 +187,8 @@ public void commit(@Nullable TableMetadata base, TableMetadata metadata) protected abstract void commitToExistingTable(TableMetadata base, TableMetadata metadata); + protected abstract void commitMaterializedViewRefresh(TableMetadata base, TableMetadata metadata); + @Override public FileIO io() { @@ -224,6 +237,13 @@ protected String writeNewMetadata(TableMetadata metadata, int newVersion) } protected void refreshFromMetadataLocation(String newLocation) + { + refreshFromMetadataLocation( + newLocation, + metadataLocation -> TableMetadataParser.read(fileIo, fileIo.newInputFile(metadataLocation))); + } + + protected void refreshFromMetadataLocation(String newLocation, Function metadataLoader) { // use null-safe equality check because new tables have a null metadata location if (Objects.equals(currentMetadataLocation, newLocation)) { @@ -231,24 +251,27 @@ protected void refreshFromMetadataLocation(String newLocation) return; } + // a table that is replaced doesn't need its metadata reloaded + if (newLocation == null) { + shouldRefresh = false; + return; + } + TableMetadata newMetadata; try { newMetadata = Failsafe.with(RetryPolicy.builder() - .withMaxRetries(20) + .withMaxRetries(3) .withBackoff(100, 5000, MILLIS, 4.0) - .withMaxDuration(Duration.ofMinutes(10)) - .abortOn(failure -> failure instanceof ValidationException || isNotFoundException(failure)) + .withMaxDuration(Duration.ofMinutes(3)) + .abortOn(throwable -> TrinoFileSystem.isUnrecoverableException(throwable) || IcebergExceptions.isFatalException(throwable)) .build()) - .get(() -> TableMetadataParser.read(fileIo, io().newInputFile(newLocation))); + .get(() -> metadataLoader.apply(newLocation)); + } + catch (UncheckedIOException e) { + throw new TrinoException(ICEBERG_INVALID_METADATA, "Error accessing metadata file for table %s".formatted(getSchemaTableName().toString()), e); } catch (Throwable failure) { - if (isNotFoundException(failure)) { - throw new TrinoException(ICEBERG_MISSING_METADATA, "Metadata not found in metadata location for table " + getSchemaTableName(), failure); - } - if (failure instanceof ValidationException) { - throw new TrinoException(ICEBERG_INVALID_METADATA, "Invalid metadata file for table " + getSchemaTableName(), failure); - } - throw failure; + throw translateMetadataException(failure, getSchemaTableName().toString()); } String newUUID = newMetadata.uuid(); @@ -263,14 +286,6 @@ protected void refreshFromMetadataLocation(String newLocation) shouldRefresh = false; } - private static boolean isNotFoundException(Throwable failure) - { - // qualified name, as this is NOT the io.trino.spi.connector.NotFoundException - return failure instanceof org.apache.iceberg.exceptions.NotFoundException || - // This is used in context where the code cannot throw a checked exception, so FileNotFoundException would need to be wrapped - failure.getCause() instanceof FileNotFoundException; - } - protected static String newTableMetadataFilePath(TableMetadata meta, int newVersion) { String codec = meta.property(METADATA_COMPRESSION, METADATA_COMPRESSION_DEFAULT); @@ -291,8 +306,9 @@ public static List toHiveColumns(List columns) return columns.stream() .map(column -> new Column( column.name(), - toHiveType(HiveSchemaUtil.convert(column.type())), - Optional.empty())) + HiveType.fromTypeInfo(HiveSchemaUtil.convert(column.type())), + Optional.empty(), + Map.of())) .collect(toImmutableList()); } } diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/AbstractTrinoCatalog.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/AbstractTrinoCatalog.java index 6747414e60c8..bb445045dd9e 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/AbstractTrinoCatalog.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/AbstractTrinoCatalog.java @@ -16,14 +16,20 @@ import com.google.common.collect.ImmutableMap; import dev.failsafe.Failsafe; import dev.failsafe.RetryPolicy; +import io.airlift.log.Logger; import io.trino.filesystem.Location; import io.trino.filesystem.TrinoFileSystem; +import io.trino.filesystem.TrinoFileSystemFactory; import io.trino.plugin.base.CatalogName; import io.trino.plugin.hive.HiveMetadata; +import io.trino.plugin.hive.metastore.TableInfo; import io.trino.plugin.iceberg.ColumnIdentity; import io.trino.plugin.iceberg.IcebergMaterializedViewDefinition; import io.trino.plugin.iceberg.IcebergUtil; import io.trino.plugin.iceberg.PartitionTransforms.ColumnTransform; +import io.trino.plugin.iceberg.catalog.glue.TrinoGlueCatalog; +import io.trino.plugin.iceberg.fileio.ForwardingFileIo; +import io.trino.plugin.iceberg.fileio.ForwardingOutputFile; import io.trino.spi.TrinoException; import io.trino.spi.connector.CatalogSchemaTableName; import io.trino.spi.connector.ColumnMetadata; @@ -32,6 +38,8 @@ import io.trino.spi.connector.ConnectorTableMetadata; import io.trino.spi.connector.ConnectorViewDefinition; import io.trino.spi.connector.SchemaTableName; +import io.trino.spi.connector.TableNotFoundException; +import io.trino.spi.connector.ViewNotFoundException; import io.trino.spi.type.ArrayType; import io.trino.spi.type.CharType; import io.trino.spi.type.MapType; @@ -43,11 +51,13 @@ import io.trino.spi.type.Type; import io.trino.spi.type.TypeManager; import org.apache.iceberg.AppendFiles; +import org.apache.iceberg.BaseTable; import org.apache.iceberg.PartitionSpec; import org.apache.iceberg.Schema; import org.apache.iceberg.SortOrder; import org.apache.iceberg.Table; import org.apache.iceberg.TableMetadata; +import org.apache.iceberg.TableMetadataParser; import org.apache.iceberg.TableOperations; import org.apache.iceberg.Transaction; import org.apache.iceberg.types.Types; @@ -55,7 +65,6 @@ import java.io.IOException; import java.time.Duration; import java.time.temporal.ChronoUnit; -import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.Optional; @@ -66,23 +75,29 @@ import static com.google.common.collect.ImmutableList.toImmutableList; import static com.google.common.collect.ImmutableSet.toImmutableSet; import static io.trino.plugin.hive.HiveMetadata.STORAGE_TABLE; -import static io.trino.plugin.hive.HiveMetadata.TABLE_COMMENT; -import static io.trino.plugin.hive.ViewReaderUtil.ICEBERG_MATERIALIZED_VIEW_COMMENT; import static io.trino.plugin.hive.ViewReaderUtil.PRESTO_VIEW_FLAG; -import static io.trino.plugin.hive.metastore.glue.converter.GlueToTrinoConverter.mappedCopy; +import static io.trino.plugin.hive.metastore.Table.TABLE_COMMENT; +import static io.trino.plugin.hive.metastore.TableInfo.ICEBERG_MATERIALIZED_VIEW_COMMENT; import static io.trino.plugin.hive.util.HiveUtil.escapeTableName; import static io.trino.plugin.iceberg.IcebergErrorCode.ICEBERG_FILESYSTEM_ERROR; -import static io.trino.plugin.iceberg.IcebergMaterializedViewAdditionalProperties.STORAGE_SCHEMA; -import static io.trino.plugin.iceberg.IcebergMaterializedViewAdditionalProperties.getStorageSchema; import static io.trino.plugin.iceberg.IcebergMaterializedViewDefinition.decodeMaterializedViewData; -import static io.trino.plugin.iceberg.IcebergTableProperties.FILE_FORMAT_PROPERTY; +import static io.trino.plugin.iceberg.IcebergMaterializedViewProperties.STORAGE_SCHEMA; +import static io.trino.plugin.iceberg.IcebergMaterializedViewProperties.getStorageSchema; +import static io.trino.plugin.iceberg.IcebergTableName.tableNameWithType; import static io.trino.plugin.iceberg.IcebergTableProperties.getPartitioning; +import static io.trino.plugin.iceberg.IcebergTableProperties.getSortOrder; +import static io.trino.plugin.iceberg.IcebergTableProperties.getTableLocation; +import static io.trino.plugin.iceberg.IcebergUtil.METADATA_FOLDER_NAME; +import static io.trino.plugin.iceberg.IcebergUtil.TRINO_QUERY_ID_NAME; import static io.trino.plugin.iceberg.IcebergUtil.commit; -import static io.trino.plugin.iceberg.IcebergUtil.getIcebergTableProperties; +import static io.trino.plugin.iceberg.IcebergUtil.createTableProperties; import static io.trino.plugin.iceberg.IcebergUtil.schemaFromMetadata; import static io.trino.plugin.iceberg.PartitionFields.parsePartitionFields; import static io.trino.plugin.iceberg.PartitionTransforms.getColumnTransform; +import static io.trino.plugin.iceberg.SortFieldUtils.parseSortFields; +import static io.trino.plugin.iceberg.TableType.MATERIALIZED_VIEW_STORAGE; import static io.trino.plugin.iceberg.TypeConverter.toTrinoType; +import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED; import static io.trino.spi.StandardErrorCode.TABLE_NOT_FOUND; import static io.trino.spi.type.IntegerType.INTEGER; import static io.trino.spi.type.SmallintType.SMALLINT; @@ -94,31 +109,40 @@ import static java.lang.String.format; import static java.util.Objects.requireNonNull; import static java.util.UUID.randomUUID; +import static org.apache.iceberg.BaseMetastoreTableOperations.METADATA_LOCATION_PROP; import static org.apache.iceberg.TableMetadata.newTableMetadata; -import static org.apache.iceberg.TableProperties.DEFAULT_FILE_FORMAT_DEFAULT; +import static org.apache.iceberg.TableMetadataParser.getFileExtension; +import static org.apache.iceberg.TableProperties.METADATA_COMPRESSION_DEFAULT; +import static org.apache.iceberg.Transactions.createOrReplaceTableTransaction; import static org.apache.iceberg.Transactions.createTableTransaction; public abstract class AbstractTrinoCatalog implements TrinoCatalog { + private static final Logger LOG = Logger.get(TrinoGlueCatalog.class); + public static final String TRINO_CREATED_BY_VALUE = "Trino Iceberg connector"; + public static final String ICEBERG_VIEW_RUN_AS_OWNER = "trino.run-as-owner"; + protected static final String TRINO_CREATED_BY = HiveMetadata.TRINO_CREATED_BY; - protected static final String PRESTO_QUERY_ID_NAME = HiveMetadata.PRESTO_QUERY_ID_NAME; private final CatalogName catalogName; - private final TypeManager typeManager; + protected final TypeManager typeManager; protected final IcebergTableOperationsProvider tableOperationsProvider; + private final TrinoFileSystemFactory fileSystemFactory; private final boolean useUniqueTableLocation; protected AbstractTrinoCatalog( CatalogName catalogName, TypeManager typeManager, IcebergTableOperationsProvider tableOperationsProvider, + TrinoFileSystemFactory fileSystemFactory, boolean useUniqueTableLocation) { this.catalogName = requireNonNull(catalogName, "catalogName is null"); this.typeManager = requireNonNull(typeManager, "typeManager is null"); this.tableOperationsProvider = requireNonNull(tableOperationsProvider, "tableOperationsProvider is null"); + this.fileSystemFactory = requireNonNull(fileSystemFactory, "fileSystemFactory is null"); this.useUniqueTableLocation = useUniqueTableLocation; } @@ -132,6 +156,7 @@ public void updateTableComment(ConnectorSession session, SchemaTableName schemaT else { icebergTable.updateProperties().set(TABLE_COMMENT, comment.get()).commit(); } + invalidateTableCache(schemaTableName); } @Override @@ -139,18 +164,23 @@ public void updateColumnComment(ConnectorSession session, SchemaTableName schema { Table icebergTable = loadTable(session, schemaTableName); icebergTable.updateSchema().updateColumnDoc(columnIdentity.getName(), comment.orElse(null)).commit(); + invalidateTableCache(schemaTableName); } @Override public Map getViews(ConnectorSession session, Optional namespace) { ImmutableMap.Builder views = ImmutableMap.builder(); - for (SchemaTableName name : listViews(session, namespace)) { + for (TableInfo tableInfo : listTables(session, namespace)) { + if (tableInfo.extendedRelationType() != TableInfo.ExtendedRelationType.TRINO_VIEW) { + continue; + } + SchemaTableName name = tableInfo.tableName(); try { getView(session, name).ifPresent(view -> views.put(name, view)); } catch (TrinoException e) { - if (e.getErrorCode().equals(TABLE_NOT_FOUND.toErrorCode())) { + if (e.getErrorCode().equals(TABLE_NOT_FOUND.toErrorCode()) || e instanceof TableNotFoundException || e instanceof ViewNotFoundException) { // Ignore view that was dropped during query execution (race condition) } else { @@ -187,21 +217,61 @@ protected Transaction newCreateTableTransaction( Schema schema, PartitionSpec partitionSpec, SortOrder sortOrder, - String location, + Optional location, Map properties, Optional owner) { - TableMetadata metadata = newTableMetadata(schema, partitionSpec, sortOrder, location, properties); + TableMetadata metadata = newTableMetadata(schema, partitionSpec, sortOrder, location.orElse(null), properties); TableOperations ops = tableOperationsProvider.createTableOperations( this, session, schemaTableName.getSchemaName(), schemaTableName.getTableName(), owner, - Optional.of(location)); + location); return createTableTransaction(schemaTableName.toString(), ops, metadata); } + protected Transaction newCreateOrReplaceTableTransaction( + ConnectorSession session, + SchemaTableName schemaTableName, + Schema schema, + PartitionSpec partitionSpec, + SortOrder sortOrder, + String location, + Map properties, + Optional owner) + { + BaseTable table; + Optional metadata = Optional.empty(); + try { + table = loadTable(session, new SchemaTableName(schemaTableName.getSchemaName(), schemaTableName.getTableName())); + metadata = Optional.of(table.operations().current()); + } + catch (TableNotFoundException ignore) { + // ignored + } + IcebergTableOperations operations = tableOperationsProvider.createTableOperations( + this, + session, + schemaTableName.getSchemaName(), + schemaTableName.getTableName(), + owner, + Optional.of(location)); + TableMetadata newMetaData; + if (metadata.isPresent()) { + operations.initializeFromMetadata(metadata.get()); + newMetaData = operations.current() + // don't inherit table properties from earlier snapshots + .replaceProperties(properties) + .buildReplacement(schema, partitionSpec, sortOrder, location, properties); + } + else { + newMetaData = newTableMetadata(schema, partitionSpec, sortOrder, location, properties); + } + return createOrReplaceTableTransaction(schemaTableName.toString(), operations, newMetaData); + } + protected String createNewTableName(String baseTableName) { String tableNameLocationComponent = escapeTableName(baseTableName); @@ -221,20 +291,72 @@ protected void deleteTableDirectory(TrinoFileSystem fileSystem, SchemaTableName } } - protected SchemaTableName createMaterializedViewStorageTable(ConnectorSession session, SchemaTableName viewName, ConnectorMaterializedViewDefinition definition) + protected Location createMaterializedViewStorage( + ConnectorSession session, + SchemaTableName viewName, + ConnectorMaterializedViewDefinition definition, + Map materializedViewProperties) + { + if (getStorageSchema(materializedViewProperties).isPresent()) { + throw new TrinoException(NOT_SUPPORTED, "Materialized view property '%s' is not supported when hiding materialized view storage tables is enabled".formatted(STORAGE_SCHEMA)); + } + SchemaTableName storageTableName = new SchemaTableName(viewName.getSchemaName(), tableNameWithType(viewName.getTableName(), MATERIALIZED_VIEW_STORAGE)); + String tableLocation = getTableLocation(materializedViewProperties) + .orElseGet(() -> defaultTableLocation(session, viewName)); + List columns = columnsForMaterializedView(definition, materializedViewProperties); + + Schema schema = schemaFromMetadata(columns); + PartitionSpec partitionSpec = parsePartitionFields(schema, getPartitioning(materializedViewProperties)); + SortOrder sortOrder = parseSortFields(schema, getSortOrder(materializedViewProperties)); + Map properties = createTableProperties(new ConnectorTableMetadata(storageTableName, columns, materializedViewProperties, Optional.empty()), ignore -> false); + + TableMetadata metadata = newTableMetadata(schema, partitionSpec, sortOrder, tableLocation, properties); + + String fileName = format("%05d-%s%s", 0, randomUUID(), getFileExtension(METADATA_COMPRESSION_DEFAULT)); + Location metadataFileLocation = Location.of(tableLocation).appendPath(METADATA_FOLDER_NAME).appendPath(fileName); + + TrinoFileSystem fileSystem = fileSystemFactory.create(session); + TableMetadataParser.write(metadata, new ForwardingOutputFile(fileSystem, metadataFileLocation)); + + return metadataFileLocation; + } + + protected void dropMaterializedViewStorage(TrinoFileSystem fileSystem, String storageMetadataLocation) + throws IOException + { + TableMetadata metadata = TableMetadataParser.read(new ForwardingFileIo(fileSystem), storageMetadataLocation); + String storageLocation = metadata.location(); + fileSystem.deleteDirectory(Location.of(storageLocation)); + } + + protected SchemaTableName createMaterializedViewStorageTable( + ConnectorSession session, + SchemaTableName viewName, + ConnectorMaterializedViewDefinition definition, + Map materializedViewProperties) { // Generate a storage table name and create a storage table. The properties in the definition are table properties for the // storage table as indicated in the materialized view definition. String storageTableName = "st_" + randomUUID().toString().replace("-", ""); - Map storageTableProperties = new HashMap<>(definition.getProperties()); - storageTableProperties.putIfAbsent(FILE_FORMAT_PROPERTY, DEFAULT_FILE_FORMAT_DEFAULT); - String storageSchema = getStorageSchema(definition.getProperties()).orElse(viewName.getSchemaName()); + String storageSchema = getStorageSchema(materializedViewProperties).orElse(viewName.getSchemaName()); SchemaTableName storageTable = new SchemaTableName(storageSchema, storageTableName); + List columns = columnsForMaterializedView(definition, materializedViewProperties); - Schema schemaWithTimestampTzPreserved = schemaFromMetadata(mappedCopy( - definition.getColumns(), - column -> { + ConnectorTableMetadata tableMetadata = new ConnectorTableMetadata(storageTable, columns, materializedViewProperties, Optional.empty()); + String tableLocation = getTableLocation(tableMetadata.getProperties()) + .orElseGet(() -> defaultTableLocation(session, tableMetadata.getTable())); + Transaction transaction = IcebergUtil.newCreateTableTransaction(this, tableMetadata, session, tableLocation, ignore -> false); + AppendFiles appendFiles = transaction.newAppend(); + commit(appendFiles, session); + transaction.commitTransaction(); + return storageTable; + } + + private List columnsForMaterializedView(ConnectorMaterializedViewDefinition definition, Map materializedViewProperties) + { + Schema schemaWithTimestampTzPreserved = schemaFromMetadata(definition.getColumns().stream() + .map(column -> { Type type = typeManager.getType(column.getType()); if (type instanceof TimestampWithTimeZoneType timestampTzType && timestampTzType.getPrecision() <= 6) { // For now preserve timestamptz columns so that we can parse partitioning @@ -244,8 +366,9 @@ protected SchemaTableName createMaterializedViewStorageTable(ConnectorSession se type = typeForMaterializedViewStorageTable(type); } return new ColumnMetadata(column.getName(), type); - })); - PartitionSpec partitionSpec = parsePartitionFields(schemaWithTimestampTzPreserved, getPartitioning(definition.getProperties())); + }) + .collect(toImmutableList())); + PartitionSpec partitionSpec = parsePartitionFields(schemaWithTimestampTzPreserved, getPartitioning(materializedViewProperties)); Set temporalPartitioningSources = partitionSpec.fields().stream() .flatMap(partitionField -> { Types.NestedField sourceField = schemaWithTimestampTzPreserved.findField(partitionField.sourceId()); @@ -258,9 +381,8 @@ protected SchemaTableName createMaterializedViewStorageTable(ConnectorSession se }) .collect(toImmutableSet()); - List columns = mappedCopy( - definition.getColumns(), - column -> { + return definition.getColumns().stream() + .map(column -> { Type type = typeManager.getType(column.getType()); if (type instanceof TimestampWithTimeZoneType timestampTzType && timestampTzType.getPrecision() <= 6 && temporalPartitioningSources.contains(column.getName())) { // Apply point-in-time semantics to maintain partitioning capabilities @@ -270,14 +392,8 @@ protected SchemaTableName createMaterializedViewStorageTable(ConnectorSession se type = typeForMaterializedViewStorageTable(type); } return new ColumnMetadata(column.getName(), type); - }); - - ConnectorTableMetadata tableMetadata = new ConnectorTableMetadata(storageTable, columns, storageTableProperties, Optional.empty()); - Transaction transaction = IcebergUtil.newCreateTableTransaction(this, tableMetadata, session); - AppendFiles appendFiles = transaction.newAppend(); - commit(appendFiles, session); - transaction.commitTransaction(); - return storageTable; + }) + .collect(toImmutableList()); } /** @@ -334,7 +450,6 @@ private Type typeForMaterializedViewStorageTable(Type type) } protected ConnectorMaterializedViewDefinition getMaterializedViewDefinition( - Table icebergTable, Optional owner, String viewOriginalText, SchemaTableName storageTableName) @@ -349,10 +464,7 @@ protected ConnectorMaterializedViewDefinition getMaterializedViewDefinition( definition.getGracePeriod(), definition.getComment(), owner, - ImmutableMap.builder() - .putAll(getIcebergTableProperties(icebergTable)) - .put(STORAGE_SCHEMA, storageTableName.getSchemaName()) - .buildOrThrow()); + Map.of()); } protected List toSpiMaterializedViewColumns(List columns) @@ -365,7 +477,7 @@ protected List toSpiMaterializedView protected Map createMaterializedViewProperties(ConnectorSession session, SchemaTableName storageTableName) { return ImmutableMap.builder() - .put(PRESTO_QUERY_ID_NAME, session.getQueryId()) + .put(TRINO_QUERY_ID_NAME, session.getQueryId()) .put(STORAGE_SCHEMA, storageTableName.getSchemaName()) .put(STORAGE_TABLE, storageTableName.getTableName()) .put(PRESTO_VIEW_FLAG, "true") @@ -374,6 +486,19 @@ protected Map createMaterializedViewProperties(ConnectorSession .buildOrThrow(); } + protected Map createMaterializedViewProperties(ConnectorSession session, Location storageMetadataLocation) + { + return ImmutableMap.builder() + .put(TRINO_QUERY_ID_NAME, session.getQueryId()) + .put(METADATA_LOCATION_PROP, storageMetadataLocation.toString()) + .put(PRESTO_VIEW_FLAG, "true") + .put(TRINO_CREATED_BY, TRINO_CREATED_BY_VALUE) + .put(TABLE_COMMENT, ICEBERG_MATERIALIZED_VIEW_COMMENT) + .buildOrThrow(); + } + + protected abstract void invalidateTableCache(SchemaTableName schemaTableName); + protected static class MaterializedViewMayBeBeingRemovedException extends RuntimeException { diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/IcebergHiveMetastoreModule.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/IcebergHiveMetastoreModule.java new file mode 100644 index 000000000000..73f4f8595baf --- /dev/null +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/IcebergHiveMetastoreModule.java @@ -0,0 +1,50 @@ +/* + * 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.iceberg.catalog; + +import com.google.inject.Binder; +import com.google.inject.Key; +import com.google.inject.Scopes; +import com.google.inject.multibindings.Multibinder; +import io.airlift.configuration.AbstractConfigurationAwareModule; +import io.airlift.units.Duration; +import io.trino.plugin.hive.HideDeltaLakeTables; +import io.trino.plugin.hive.metastore.CachingHiveMetastoreModule; +import io.trino.plugin.hive.metastore.cache.CachingHiveMetastoreConfig; +import io.trino.plugin.iceberg.procedure.MigrateProcedure; +import io.trino.spi.procedure.Procedure; + +import java.util.concurrent.TimeUnit; + +import static com.google.inject.multibindings.Multibinder.newSetBinder; +import static io.airlift.configuration.ConfigBinder.configBinder; + +public class IcebergHiveMetastoreModule + extends AbstractConfigurationAwareModule +{ + @Override + protected void setup(Binder binder) + { + binder.bind(Key.get(boolean.class, HideDeltaLakeTables.class)).toInstance(false); + + install(new CachingHiveMetastoreModule()); + + // ensure caching metastore wrapper isn't created, as it's not leveraged by Iceberg + configBinder(binder).bindConfigDefaults(CachingHiveMetastoreConfig.class, config -> + config.setStatsCacheTtl(new Duration(0, TimeUnit.SECONDS))); + + Multibinder procedures = newSetBinder(binder, Procedure.class); + procedures.addBinding().toProvider(MigrateProcedure.class).in(Scopes.SINGLETON); + } +} diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/TrinoCatalog.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/TrinoCatalog.java index 46282a4524aa..a26fee8d015b 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/TrinoCatalog.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/TrinoCatalog.java @@ -13,6 +13,7 @@ */ package io.trino.plugin.iceberg.catalog; +import io.trino.plugin.hive.metastore.TableInfo; import io.trino.plugin.iceberg.ColumnIdentity; import io.trino.plugin.iceberg.UnknownTableTypeException; import io.trino.spi.connector.CatalogSchemaTableName; @@ -24,6 +25,8 @@ import io.trino.spi.connector.RelationCommentMetadata; import io.trino.spi.connector.SchemaTableName; import io.trino.spi.security.TrinoPrincipal; +import jakarta.annotation.Nullable; +import org.apache.iceberg.BaseTable; import org.apache.iceberg.PartitionSpec; import org.apache.iceberg.Schema; import org.apache.iceberg.SortOrder; @@ -39,6 +42,8 @@ import java.util.function.Predicate; import java.util.function.UnaryOperator; +import static com.google.common.collect.ImmutableList.toImmutableList; + /** * An interface to allow different Iceberg catalog implementations in IcebergMetadata. *

@@ -62,6 +67,11 @@ public interface TrinoCatalog void dropNamespace(ConnectorSession session, String namespace); + default Optional getNamespaceSeparator() + { + return Optional.empty(); + } + Map loadNamespaceMetadata(ConnectorSession session, String namespace); Optional getNamespacePrincipal(ConnectorSession session, String namespace); @@ -72,7 +82,17 @@ public interface TrinoCatalog void renameNamespace(ConnectorSession session, String source, String target); - List listTables(ConnectorSession session, Optional namespace); + List listTables(ConnectorSession session, Optional namespace); + + List listIcebergTables(ConnectorSession session, Optional namespace); + + default List listViews(ConnectorSession session, Optional namespace) + { + return listTables(session, namespace).stream() + .filter(info -> info.extendedRelationType() == TableInfo.ExtendedRelationType.TRINO_VIEW) + .map(TableInfo::tableName) + .collect(toImmutableList()); + } Optional> streamRelationColumns( ConnectorSession session, @@ -86,7 +106,21 @@ Optional> streamRelationComments( UnaryOperator> relationFilter, Predicate isRedirected); + default Transaction newTransaction(Table icebergTable) + { + return icebergTable.newTransaction(); + } + Transaction newCreateTableTransaction( + ConnectorSession session, + SchemaTableName schemaTableName, + Schema schema, + PartitionSpec partitionSpec, + SortOrder sortOrder, + Optional location, + Map properties); + + Transaction newCreateOrReplaceTableTransaction( ConnectorSession session, SchemaTableName schemaTableName, Schema schema, @@ -113,7 +147,7 @@ Transaction newCreateTableTransaction( * @return Iceberg table loaded * @throws UnknownTableTypeException if table is not of Iceberg type in the metastore */ - Table loadTable(ConnectorSession session, SchemaTableName schemaTableName); + BaseTable loadTable(ConnectorSession session, SchemaTableName schemaTableName); /** * Bulk load column metadata. The returned map may contain fewer entries then asked for. @@ -128,6 +162,7 @@ Transaction newCreateTableTransaction( void updateMaterializedViewColumnComment(ConnectorSession session, SchemaTableName schemaViewName, String columnName, Optional comment); + @Nullable String defaultTableLocation(ConnectorSession session, SchemaTableName schemaTableName); void setTablePrincipal(ConnectorSession session, SchemaTableName schemaTableName, TrinoPrincipal principal); @@ -140,18 +175,15 @@ Transaction newCreateTableTransaction( void dropView(ConnectorSession session, SchemaTableName schemaViewName); - List listViews(ConnectorSession session, Optional namespace); - Map getViews(ConnectorSession session, Optional namespace); Optional getView(ConnectorSession session, SchemaTableName viewName); - List listMaterializedViews(ConnectorSession session, Optional namespace); - void createMaterializedView( ConnectorSession session, SchemaTableName viewName, ConnectorMaterializedViewDefinition definition, + Map materializedViewProperties, boolean replace, boolean ignoreExisting); @@ -159,6 +191,8 @@ void createMaterializedView( Optional getMaterializedView(ConnectorSession session, SchemaTableName viewName); + Optional getMaterializedViewStorageTable(ConnectorSession session, SchemaTableName viewName); + void renameMaterializedView(ConnectorSession session, SchemaTableName source, SchemaTableName target); void updateColumnComment(ConnectorSession session, SchemaTableName schemaTableName, ColumnIdentity columnIdentity, Optional comment); diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/file/FileMetastoreTableOperations.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/file/FileMetastoreTableOperations.java index e25fac198f77..5fc076d0b90a 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/file/FileMetastoreTableOperations.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/file/FileMetastoreTableOperations.java @@ -13,6 +13,7 @@ */ package io.trino.plugin.iceberg.catalog.file; +import com.google.common.collect.ImmutableMap; import io.trino.annotation.NotThreadSafe; import io.trino.plugin.hive.metastore.MetastoreUtil; import io.trino.plugin.hive.metastore.PrincipalPrivileges; @@ -27,11 +28,14 @@ import org.apache.iceberg.io.FileIO; import java.util.Optional; +import java.util.function.BiFunction; import static com.google.common.base.Preconditions.checkState; import static io.trino.plugin.hive.HiveErrorCode.HIVE_CONCURRENT_MODIFICATION_DETECTED; import static io.trino.plugin.hive.metastore.PrincipalPrivileges.NO_PRIVILEGES; +import static io.trino.plugin.iceberg.IcebergTableName.tableNameFrom; import static org.apache.iceberg.BaseMetastoreTableOperations.METADATA_LOCATION_PROP; +import static org.apache.iceberg.BaseMetastoreTableOperations.PREVIOUS_METADATA_LOCATION_PROP; @NotThreadSafe public class FileMetastoreTableOperations @@ -53,9 +57,24 @@ public FileMetastoreTableOperations( protected void commitToExistingTable(TableMetadata base, TableMetadata metadata) { Table currentTable = getTable(); + commitTableUpdate(currentTable, metadata, (table, newMetadataLocation) -> Table.builder(table) + .apply(builder -> updateMetastoreTable(builder, metadata, newMetadataLocation, Optional.of(currentMetadataLocation))) + .build()); + } + + @Override + protected void commitMaterializedViewRefresh(TableMetadata base, TableMetadata metadata) + { + Table materializedView = getTable(database, tableNameFrom(tableName)); + commitTableUpdate(materializedView, metadata, (table, newMetadataLocation) -> Table.builder(table) + .apply(builder -> builder.setParameter(METADATA_LOCATION_PROP, newMetadataLocation).setParameter(PREVIOUS_METADATA_LOCATION_PROP, currentMetadataLocation)) + .build()); + } + private void commitTableUpdate(Table table, TableMetadata metadata, BiFunction tableUpdateFunction) + { checkState(currentMetadataLocation != null, "No current metadata location for existing table"); - String metadataLocation = currentTable.getParameters().get(METADATA_LOCATION_PROP); + String metadataLocation = table.getParameters().get(METADATA_LOCATION_PROP); if (!currentMetadataLocation.equals(metadataLocation)) { throw new CommitFailedException("Metadata location [%s] is not same as table metadata location [%s] for %s", currentMetadataLocation, metadataLocation, getSchemaTableName()); @@ -63,15 +82,13 @@ protected void commitToExistingTable(TableMetadata base, TableMetadata metadata) String newMetadataLocation = writeNewMetadata(metadata, version.orElseThrow() + 1); - Table table = Table.builder(currentTable) - .apply(builder -> updateMetastoreTable(builder, metadata, newMetadataLocation, Optional.of(currentMetadataLocation))) - .build(); + Table updatedTable = tableUpdateFunction.apply(table, newMetadataLocation); // todo privileges should not be replaced for an alter PrincipalPrivileges privileges = table.getOwner().map(MetastoreUtil::buildInitialPrivilegeSet).orElse(NO_PRIVILEGES); try { - metastore.replaceTable(database, tableName, table, privileges); + metastore.replaceTable(database, table.getTableName(), updatedTable, privileges, ImmutableMap.of()); } catch (RuntimeException e) { if (e instanceof TrinoException trinoException && diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/file/IcebergFileMetastoreCatalogModule.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/file/IcebergFileMetastoreCatalogModule.java index e937260c7c86..e61620fc15d1 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/file/IcebergFileMetastoreCatalogModule.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/file/IcebergFileMetastoreCatalogModule.java @@ -20,11 +20,10 @@ import io.airlift.configuration.AbstractConfigurationAwareModule; import io.airlift.units.Duration; import io.trino.plugin.hive.HideDeltaLakeTables; -import io.trino.plugin.hive.metastore.DecoratedHiveMetastoreModule; +import io.trino.plugin.hive.metastore.CachingHiveMetastoreModule; import io.trino.plugin.hive.metastore.cache.CachingHiveMetastoreConfig; import io.trino.plugin.hive.metastore.file.FileMetastoreModule; import io.trino.plugin.iceberg.catalog.IcebergTableOperationsProvider; -import io.trino.plugin.iceberg.catalog.MetastoreValidator; import io.trino.plugin.iceberg.catalog.TrinoCatalogFactory; import io.trino.plugin.iceberg.catalog.hms.TrinoHiveCatalogFactory; import io.trino.plugin.iceberg.procedure.MigrateProcedure; @@ -45,9 +44,8 @@ protected void setup(Binder binder) install(new FileMetastoreModule()); binder.bind(IcebergTableOperationsProvider.class).to(FileMetastoreTableOperationsProvider.class).in(Scopes.SINGLETON); binder.bind(TrinoCatalogFactory.class).to(TrinoHiveCatalogFactory.class).in(Scopes.SINGLETON); - binder.bind(MetastoreValidator.class).asEagerSingleton(); binder.bind(Key.get(boolean.class, HideDeltaLakeTables.class)).toInstance(HIDE_DELTA_LAKE_TABLES_IN_ICEBERG); - install(new DecoratedHiveMetastoreModule(false)); + install(new CachingHiveMetastoreModule()); configBinder(binder).bindConfigDefaults(CachingHiveMetastoreConfig.class, config -> { // ensure caching metastore wrapper isn't created, as it's not leveraged by Iceberg diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/glue/GlueIcebergTableOperations.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/glue/GlueIcebergTableOperations.java index 04cc5645f934..a195251b68c6 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/glue/GlueIcebergTableOperations.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/glue/GlueIcebergTableOperations.java @@ -13,16 +13,6 @@ */ package io.trino.plugin.iceberg.catalog.glue; -import com.amazonaws.services.glue.AWSGlueAsync; -import com.amazonaws.services.glue.model.AlreadyExistsException; -import com.amazonaws.services.glue.model.ConcurrentModificationException; -import com.amazonaws.services.glue.model.CreateTableRequest; -import com.amazonaws.services.glue.model.EntityNotFoundException; -import com.amazonaws.services.glue.model.InvalidInputException; -import com.amazonaws.services.glue.model.ResourceNumberLimitExceededException; -import com.amazonaws.services.glue.model.Table; -import com.amazonaws.services.glue.model.TableInput; -import com.amazonaws.services.glue.model.UpdateTableRequest; import com.google.common.collect.ImmutableMap; import io.trino.plugin.hive.metastore.glue.GlueMetastoreStats; import io.trino.plugin.iceberg.UnknownTableTypeException; @@ -37,17 +27,31 @@ import org.apache.iceberg.exceptions.CommitFailedException; import org.apache.iceberg.exceptions.CommitStateUnknownException; import org.apache.iceberg.io.FileIO; +import software.amazon.awssdk.services.glue.GlueClient; +import software.amazon.awssdk.services.glue.model.AlreadyExistsException; +import software.amazon.awssdk.services.glue.model.ConcurrentModificationException; +import software.amazon.awssdk.services.glue.model.EntityNotFoundException; +import software.amazon.awssdk.services.glue.model.InvalidInputException; +import software.amazon.awssdk.services.glue.model.ResourceNumberLimitExceededException; +import software.amazon.awssdk.services.glue.model.StorageDescriptor; +import software.amazon.awssdk.services.glue.model.Table; +import software.amazon.awssdk.services.glue.model.TableInput; +import java.util.HashMap; import java.util.Map; import java.util.Optional; +import java.util.function.BiFunction; import static com.google.common.base.Verify.verify; -import static io.trino.plugin.hive.ViewReaderUtil.isHiveOrPrestoView; -import static io.trino.plugin.hive.ViewReaderUtil.isPrestoView; -import static io.trino.plugin.hive.metastore.glue.converter.GlueToTrinoConverter.getTableParameters; -import static io.trino.plugin.hive.metastore.glue.converter.GlueToTrinoConverter.getTableType; +import static io.trino.plugin.hive.ViewReaderUtil.isTrinoMaterializedView; +import static io.trino.plugin.hive.ViewReaderUtil.isTrinoView; +import static io.trino.plugin.hive.metastore.glue.GlueConverter.getTableType; import static io.trino.plugin.hive.util.HiveUtil.isIcebergTable; +import static io.trino.plugin.iceberg.IcebergErrorCode.ICEBERG_COMMIT_ERROR; import static io.trino.plugin.iceberg.IcebergErrorCode.ICEBERG_INVALID_METADATA; +import static io.trino.plugin.iceberg.IcebergTableName.isMaterializedViewStorage; +import static io.trino.plugin.iceberg.IcebergTableName.tableNameFrom; +import static io.trino.plugin.iceberg.catalog.glue.GlueIcebergUtil.getMaterializedViewTableInput; import static io.trino.plugin.iceberg.catalog.glue.GlueIcebergUtil.getTableInput; import static java.lang.String.format; import static java.util.Objects.requireNonNull; @@ -59,7 +63,7 @@ public class GlueIcebergTableOperations { private final TypeManager typeManager; private final boolean cacheTableMetadata; - private final AWSGlueAsync glueClient; + private final GlueClient glueClient; private final GlueMetastoreStats stats; private final GetGlueTable getGlueTable; @@ -69,7 +73,7 @@ public class GlueIcebergTableOperations protected GlueIcebergTableOperations( TypeManager typeManager, boolean cacheTableMetadata, - AWSGlueAsync glueClient, + GlueClient glueClient, GlueMetastoreStats stats, GetGlueTable getGlueTable, FileIO fileIo, @@ -90,15 +94,25 @@ protected GlueIcebergTableOperations( @Override protected String getRefreshedLocation(boolean invalidateCaches) { - Table table = getTable(invalidateCaches); - glueVersionId = table.getVersionId(); + boolean isMaterializedViewStorageTable = isMaterializedViewStorage(tableName); - Map parameters = getTableParameters(table); - if (isPrestoView(parameters) && isHiveOrPrestoView(getTableType(table))) { - // this is a Presto Hive view, hence not a table + Table table; + if (isMaterializedViewStorageTable) { + table = getTable(database, tableNameFrom(tableName), invalidateCaches); + } + else { + table = getTable(database, tableName, invalidateCaches); + } + glueVersionId = table.versionId(); + + String tableType = getTableType(table); + Map parameters = table.parameters(); + if (!isMaterializedViewStorageTable && (isTrinoView(tableType, parameters) || isTrinoMaterializedView(tableType, parameters))) { + // this is a Hive view or Trino/Presto view, or Trino materialized view, hence not a table + // TODO table operations should not be constructed for views (remove exception-driven code path) throw new TableNotFoundException(getSchemaTableName()); } - if (!isIcebergTable(parameters)) { + if (!isMaterializedViewStorageTable && !isIcebergTable(parameters)) { throw new UnknownTableTypeException(getSchemaTableName()); } @@ -114,13 +128,12 @@ protected void commitNewTable(TableMetadata metadata) { verify(version.isEmpty(), "commitNewTable called on a table which already exists"); String newMetadataLocation = writeNewMetadata(metadata, 0); - TableInput tableInput = getTableInput(typeManager, tableName, owner, metadata, newMetadataLocation, ImmutableMap.of(), cacheTableMetadata); + TableInput tableInput = getTableInput(typeManager, tableName, owner, metadata, metadata.location(), newMetadataLocation, ImmutableMap.of(), cacheTableMetadata); - CreateTableRequest createTableRequest = new CreateTableRequest() - .withDatabaseName(database) - .withTableInput(tableInput); try { - stats.getCreateTable().call(() -> glueClient.createTable(createTableRequest)); + stats.getCreateTable().call(() -> glueClient.createTable(x -> x + .databaseName(database) + .tableInput(tableInput))); } catch (AlreadyExistsException | EntityNotFoundException @@ -136,22 +149,50 @@ protected void commitNewTable(TableMetadata metadata) @Override protected void commitToExistingTable(TableMetadata base, TableMetadata metadata) { - String newMetadataLocation = writeNewMetadata(metadata, version.orElseThrow() + 1); - TableInput tableInput = getTableInput( - typeManager, - tableName, - owner, + commitTableUpdate( + getTable(database, tableName, false), + metadata, + (table, newMetadataLocation) -> + getTableInput( + typeManager, + tableName, + owner, + metadata, + Optional.ofNullable(table.storageDescriptor()).map(StorageDescriptor::location).orElse(null), + newMetadataLocation, + ImmutableMap.of(PREVIOUS_METADATA_LOCATION_PROP, currentMetadataLocation), + cacheTableMetadata)); + } + + @Override + protected void commitMaterializedViewRefresh(TableMetadata base, TableMetadata metadata) + { + commitTableUpdate( + getTable(database, tableNameFrom(tableName), false), metadata, - newMetadataLocation, - ImmutableMap.of(PREVIOUS_METADATA_LOCATION_PROP, currentMetadataLocation), - cacheTableMetadata); - - UpdateTableRequest updateTableRequest = new UpdateTableRequest() - .withDatabaseName(database) - .withTableInput(tableInput) - .withVersionId(glueVersionId); + (table, newMetadataLocation) -> { + Map parameters = new HashMap<>(table.parameters()); + parameters.put(METADATA_LOCATION_PROP, newMetadataLocation); + parameters.put(PREVIOUS_METADATA_LOCATION_PROP, currentMetadataLocation); + + return getMaterializedViewTableInput( + table.name(), + table.viewOriginalText(), + table.owner(), + parameters); + }); + } + + private void commitTableUpdate(Table table, TableMetadata metadata, BiFunction tableUpdateFunction) + { + String newMetadataLocation = writeNewMetadata(metadata, version.orElseThrow() + 1); + TableInput tableInput = tableUpdateFunction.apply(table, newMetadataLocation); + try { - stats.getUpdateTable().call(() -> glueClient.updateTable(updateTableRequest)); + stats.getUpdateTable().call(() -> glueClient.updateTable(x -> x + .databaseName(database) + .tableInput(tableInput) + .versionId(glueVersionId))); } catch (ConcurrentModificationException e) { // CommitFailedException is handled as a special case in the Iceberg library. This commit will automatically retry @@ -159,7 +200,7 @@ protected void commitToExistingTable(TableMetadata base, TableMetadata metadata) } catch (EntityNotFoundException | InvalidInputException | ResourceNumberLimitExceededException e) { // Signal a non-retriable commit failure and eventually clean up metadata files corresponding to the current transaction - throw e; + throw new TrinoException(ICEBERG_COMMIT_ERROR, "Cannot commit table update", e); } catch (RuntimeException e) { // Cannot determine whether the `updateTable` operation was successful, @@ -169,7 +210,7 @@ protected void commitToExistingTable(TableMetadata base, TableMetadata metadata) shouldRefresh = true; } - private Table getTable(boolean invalidateCaches) + private Table getTable(String database, String tableName, boolean invalidateCaches) { return getGlueTable.get(new SchemaTableName(database, tableName), invalidateCaches); } diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/glue/GlueIcebergTableOperationsProvider.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/glue/GlueIcebergTableOperationsProvider.java index 4b54259edf7a..62310372a5a0 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/glue/GlueIcebergTableOperationsProvider.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/glue/GlueIcebergTableOperationsProvider.java @@ -13,7 +13,6 @@ */ package io.trino.plugin.iceberg.catalog.glue; -import com.amazonaws.services.glue.AWSGlueAsync; import com.google.inject.Inject; import io.trino.filesystem.TrinoFileSystemFactory; import io.trino.plugin.hive.metastore.glue.GlueMetastoreStats; @@ -23,6 +22,7 @@ import io.trino.plugin.iceberg.fileio.ForwardingFileIo; import io.trino.spi.connector.ConnectorSession; import io.trino.spi.type.TypeManager; +import software.amazon.awssdk.services.glue.GlueClient; import java.util.Optional; @@ -34,7 +34,7 @@ public class GlueIcebergTableOperationsProvider private final TypeManager typeManager; private final boolean cacheTableMetadata; private final TrinoFileSystemFactory fileSystemFactory; - private final AWSGlueAsync glueClient; + private final GlueClient glueClient; private final GlueMetastoreStats stats; @Inject @@ -43,7 +43,7 @@ public GlueIcebergTableOperationsProvider( IcebergGlueCatalogConfig catalogConfig, TrinoFileSystemFactory fileSystemFactory, GlueMetastoreStats stats, - AWSGlueAsync glueClient) + GlueClient glueClient) { this.typeManager = requireNonNull(typeManager, "typeManager is null"); this.cacheTableMetadata = catalogConfig.isCacheTableMetadata(); diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/glue/GlueIcebergUtil.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/glue/GlueIcebergUtil.java index b8988d3112d2..0dfbb999acd3 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/glue/GlueIcebergUtil.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/glue/GlueIcebergUtil.java @@ -13,9 +13,6 @@ */ package io.trino.plugin.iceberg.catalog.glue; -import com.amazonaws.services.glue.model.Column; -import com.amazonaws.services.glue.model.StorageDescriptor; -import com.amazonaws.services.glue.model.TableInput; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import io.trino.plugin.iceberg.TypeConverter; @@ -24,6 +21,9 @@ import org.apache.iceberg.TableMetadata; import org.apache.iceberg.types.Type; import org.apache.iceberg.types.Types; +import software.amazon.awssdk.services.glue.model.Column; +import software.amazon.awssdk.services.glue.model.StorageDescriptor; +import software.amazon.awssdk.services.glue.model.TableInput; import java.util.HashMap; import java.util.List; @@ -35,12 +35,13 @@ import static com.google.common.base.MoreObjects.firstNonNull; import static com.google.common.collect.ImmutableList.builderWithExpectedSize; import static io.trino.plugin.hive.HiveMetadata.PRESTO_VIEW_EXPANDED_TEXT_MARKER; -import static io.trino.plugin.hive.HiveMetadata.TABLE_COMMENT; import static io.trino.plugin.hive.TableType.EXTERNAL_TABLE; import static io.trino.plugin.hive.TableType.VIRTUAL_VIEW; -import static io.trino.plugin.hive.ViewReaderUtil.ICEBERG_MATERIALIZED_VIEW_COMMENT; +import static io.trino.plugin.hive.metastore.Table.TABLE_COMMENT; +import static io.trino.plugin.hive.metastore.TableInfo.ICEBERG_MATERIALIZED_VIEW_COMMENT; import static io.trino.plugin.iceberg.IcebergUtil.COLUMN_TRINO_NOT_NULL_PROPERTY; import static io.trino.plugin.iceberg.IcebergUtil.COLUMN_TRINO_TYPE_ID_PROPERTY; +import static io.trino.plugin.iceberg.IcebergUtil.TRINO_TABLE_COMMENT_CACHE_PREVENTED; import static io.trino.plugin.iceberg.IcebergUtil.TRINO_TABLE_METADATA_INFO_VALID_FOR; import static java.util.Locale.ENGLISH; import static org.apache.iceberg.BaseMetastoreTableOperations.ICEBERG_TABLE_TYPE_VALUE; @@ -67,6 +68,7 @@ public static TableInput getTableInput( String tableName, Optional owner, TableMetadata metadata, + @Nullable String tableLocation, String newMetadataLocation, Map parameters, boolean cacheTableMetadata) @@ -76,38 +78,44 @@ public static TableInput getTableInput( parameters.put(METADATA_LOCATION_PROP, newMetadataLocation); parameters.remove(TRINO_TABLE_METADATA_INFO_VALID_FOR); // no longer valid - TableInput tableInput = new TableInput() - .withName(tableName) - .withOwner(owner.orElse(null)) + StorageDescriptor.Builder storageDescriptor = StorageDescriptor.builder() + .location(tableLocation); + + TableInput.Builder tableInput = TableInput.builder() + .name(tableName) + .owner(owner.orElse(null)) // Iceberg does not distinguish managed and external tables, all tables are treated the same and marked as EXTERNAL - .withTableType(EXTERNAL_TABLE.name()); + .tableType(EXTERNAL_TABLE.name()) + .storageDescriptor(storageDescriptor.build()); if (cacheTableMetadata) { // Store table metadata sufficient to answer information_schema.columns and system.metadata.table_comments queries, which are often queried in bulk by e.g. BI tools - String comment = metadata.properties().get(TABLE_COMMENT); Optional> glueColumns = glueColumns(typeManager, metadata); - boolean canPersistComment = (comment == null || comment.length() <= GLUE_TABLE_PARAMETER_LENGTH_LIMIT); - boolean canPersistColumnInfo = glueColumns.isPresent(); - boolean canPersistMetadata = canPersistComment && canPersistColumnInfo; - - if (canPersistMetadata) { - tableInput.withStorageDescriptor(new StorageDescriptor() - .withColumns(glueColumns.get())); + glueColumns.ifPresent(columns -> tableInput.storageDescriptor( + storageDescriptor.columns(columns).build())); - if (comment != null) { + String comment = metadata.properties().get(TABLE_COMMENT); + if (comment != null) { + if (comment.length() <= GLUE_TABLE_PARAMETER_LENGTH_LIMIT) { parameters.put(TABLE_COMMENT, comment); + parameters.remove(TRINO_TABLE_COMMENT_CACHE_PREVENTED); } else { parameters.remove(TABLE_COMMENT); + parameters.put(TRINO_TABLE_COMMENT_CACHE_PREVENTED, "true"); } - parameters.put(TRINO_TABLE_METADATA_INFO_VALID_FOR, newMetadataLocation); } + else { + parameters.remove(TABLE_COMMENT); + parameters.remove(TRINO_TABLE_COMMENT_CACHE_PREVENTED); + } + parameters.put(TRINO_TABLE_METADATA_INFO_VALID_FOR, newMetadataLocation); } - tableInput.withParameters(parameters); + tableInput.parameters(parameters); - return tableInput; + return tableInput.build(); } private static Optional> glueColumns(TypeManager typeManager, TableMetadata metadata) @@ -124,10 +132,10 @@ private static Optional> glueColumns(TypeManager typeManager, Table return Optional.empty(); } String trinoTypeId = TypeConverter.toTrinoType(icebergColumn.type(), typeManager).getTypeId().getId(); - Column column = new Column() - .withName(icebergColumn.name()) - .withType(glueTypeString) - .withComment(icebergColumn.doc()); + Column.Builder column = Column.builder() + .name(icebergColumn.name()) + .type(glueTypeString) + .comment(icebergColumn.doc()); ImmutableMap.Builder parameters = ImmutableMap.builder(); if (icebergColumn.isRequired()) { @@ -140,8 +148,8 @@ private static Optional> glueColumns(TypeManager typeManager, Table // Store type parameter for some (first) column so that we can later detect whether column parameters weren't erased by something. parameters.put(COLUMN_TRINO_TYPE_ID_PROPERTY, trinoTypeId); } - column.setParameters(parameters.buildOrThrow()); - glueColumns.add(column); + column.parameters(parameters.buildOrThrow()); + glueColumns.add(column.build()); firstColumn = false; } @@ -198,23 +206,25 @@ private static String toGlueTypeStringLossy(Type type) public static TableInput getViewTableInput(String viewName, String viewOriginalText, @Nullable String owner, Map parameters) { - return new TableInput() - .withName(viewName) - .withTableType(VIRTUAL_VIEW.name()) - .withViewOriginalText(viewOriginalText) - .withViewExpandedText(PRESTO_VIEW_EXPANDED_TEXT_MARKER) - .withOwner(owner) - .withParameters(parameters); + return TableInput.builder() + .name(viewName) + .tableType(VIRTUAL_VIEW.name()) + .viewOriginalText(viewOriginalText) + .viewExpandedText(PRESTO_VIEW_EXPANDED_TEXT_MARKER) + .owner(owner) + .parameters(parameters) + .build(); } public static TableInput getMaterializedViewTableInput(String viewName, String viewOriginalText, String owner, Map parameters) { - return new TableInput() - .withName(viewName) - .withTableType(VIRTUAL_VIEW.name()) - .withViewOriginalText(viewOriginalText) - .withViewExpandedText(ICEBERG_MATERIALIZED_VIEW_COMMENT) - .withOwner(owner) - .withParameters(parameters); + return TableInput.builder() + .name(viewName) + .tableType(VIRTUAL_VIEW.name()) + .viewOriginalText(viewOriginalText) + .viewExpandedText(ICEBERG_MATERIALIZED_VIEW_COMMENT) + .owner(owner) + .parameters(parameters) + .build(); } } diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/glue/IcebergGlueCatalogModule.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/glue/IcebergGlueCatalogModule.java index 32f3811d0b48..72900ffb9593 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/glue/IcebergGlueCatalogModule.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/glue/IcebergGlueCatalogModule.java @@ -13,31 +13,15 @@ */ package io.trino.plugin.iceberg.catalog.glue; -import com.amazonaws.auth.AWSCredentialsProvider; -import com.amazonaws.handlers.RequestHandler2; -import com.amazonaws.services.glue.model.Table; import com.google.inject.Binder; -import com.google.inject.Key; import com.google.inject.Scopes; -import com.google.inject.TypeLiteral; -import com.google.inject.multibindings.Multibinder; import io.airlift.configuration.AbstractConfigurationAwareModule; -import io.trino.plugin.hive.HideDeltaLakeTables; -import io.trino.plugin.hive.metastore.glue.ForGlueHiveMetastore; -import io.trino.plugin.hive.metastore.glue.GlueCredentialsProvider; import io.trino.plugin.hive.metastore.glue.GlueHiveMetastoreConfig; import io.trino.plugin.hive.metastore.glue.GlueMetastoreModule; -import io.trino.plugin.hive.metastore.glue.GlueMetastoreStats; +import io.trino.plugin.iceberg.catalog.IcebergHiveMetastoreModule; import io.trino.plugin.iceberg.catalog.IcebergTableOperationsProvider; import io.trino.plugin.iceberg.catalog.TrinoCatalogFactory; -import io.trino.plugin.iceberg.procedure.MigrateProcedure; -import io.trino.spi.procedure.Procedure; -import java.util.function.Predicate; - -import static com.google.inject.multibindings.Multibinder.newSetBinder; -import static com.google.inject.multibindings.OptionalBinder.newOptionalBinder; -import static io.airlift.configuration.ConditionalModule.conditionalModule; import static io.airlift.configuration.ConfigBinder.configBinder; import static org.weakref.jmx.guice.ExportBinder.newExporter; @@ -49,24 +33,11 @@ protected void setup(Binder binder) { configBinder(binder).bindConfig(GlueHiveMetastoreConfig.class); configBinder(binder).bindConfig(IcebergGlueCatalogConfig.class); - binder.bind(GlueMetastoreStats.class).in(Scopes.SINGLETON); - newExporter(binder).export(GlueMetastoreStats.class).withGeneratedName(); - binder.bind(AWSCredentialsProvider.class).toProvider(GlueCredentialsProvider.class).in(Scopes.SINGLETON); binder.bind(IcebergTableOperationsProvider.class).to(GlueIcebergTableOperationsProvider.class).in(Scopes.SINGLETON); binder.bind(TrinoCatalogFactory.class).to(TrinoGlueCatalogFactory.class).in(Scopes.SINGLETON); newExporter(binder).export(TrinoCatalogFactory.class).withGeneratedName(); - install(conditionalModule( - IcebergGlueCatalogConfig.class, - IcebergGlueCatalogConfig::isSkipArchive, - internalBinder -> newSetBinder(internalBinder, RequestHandler2.class, ForGlueHiveMetastore.class).addBinding().toInstance(new SkipArchiveRequestHandler()))); - - // Required to inject HiveMetastoreFactory for migrate procedure - binder.bind(Key.get(boolean.class, HideDeltaLakeTables.class)).toInstance(false); - newOptionalBinder(binder, Key.get(new TypeLiteral>() {}, ForGlueHiveMetastore.class)) - .setBinding().toInstance(table -> true); + install(new IcebergHiveMetastoreModule()); install(new GlueMetastoreModule()); - Multibinder procedures = newSetBinder(binder, Procedure.class); - procedures.addBinding().toProvider(MigrateProcedure.class).in(Scopes.SINGLETON); } } diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/glue/SkipArchiveRequestHandler.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/glue/SkipArchiveRequestHandler.java deleted file mode 100644 index 5d04d7059b01..000000000000 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/glue/SkipArchiveRequestHandler.java +++ /dev/null @@ -1,49 +0,0 @@ -/* - * 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.iceberg.catalog.glue; - -import com.amazonaws.AmazonWebServiceRequest; -import com.amazonaws.handlers.RequestHandler2; -import com.amazonaws.services.glue.model.CreateDatabaseRequest; -import com.amazonaws.services.glue.model.CreateTableRequest; -import com.amazonaws.services.glue.model.DeleteDatabaseRequest; -import com.amazonaws.services.glue.model.DeleteTableRequest; -import com.amazonaws.services.glue.model.GetDatabaseRequest; -import com.amazonaws.services.glue.model.GetDatabasesRequest; -import com.amazonaws.services.glue.model.GetTableRequest; -import com.amazonaws.services.glue.model.GetTablesRequest; -import com.amazonaws.services.glue.model.UpdateTableRequest; - -public class SkipArchiveRequestHandler - extends RequestHandler2 -{ - @Override - public AmazonWebServiceRequest beforeExecution(AmazonWebServiceRequest request) - { - if (request instanceof UpdateTableRequest updateTableRequest) { - return updateTableRequest.withSkipArchive(true); - } - if (request instanceof CreateDatabaseRequest || - request instanceof DeleteDatabaseRequest || - request instanceof GetDatabasesRequest || - request instanceof GetDatabaseRequest || - request instanceof CreateTableRequest || - request instanceof DeleteTableRequest || - request instanceof GetTablesRequest || - request instanceof GetTableRequest) { - return request; - } - throw new IllegalArgumentException("Unsupported request: " + request); - } -} diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/glue/TrinoGlueCatalog.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/glue/TrinoGlueCatalog.java index a3a0a2cb96f7..c7ae8e4ecd0b 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/glue/TrinoGlueCatalog.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/glue/TrinoGlueCatalog.java @@ -13,27 +13,8 @@ */ package io.trino.plugin.iceberg.catalog.glue; -import com.amazonaws.AmazonServiceException; -import com.amazonaws.services.glue.AWSGlueAsync; -import com.amazonaws.services.glue.model.AccessDeniedException; -import com.amazonaws.services.glue.model.AlreadyExistsException; -import com.amazonaws.services.glue.model.Column; -import com.amazonaws.services.glue.model.CreateDatabaseRequest; -import com.amazonaws.services.glue.model.CreateTableRequest; -import com.amazonaws.services.glue.model.Database; -import com.amazonaws.services.glue.model.DatabaseInput; -import com.amazonaws.services.glue.model.DeleteDatabaseRequest; -import com.amazonaws.services.glue.model.DeleteTableRequest; -import com.amazonaws.services.glue.model.EntityNotFoundException; -import com.amazonaws.services.glue.model.GetDatabaseRequest; -import com.amazonaws.services.glue.model.GetDatabasesRequest; -import com.amazonaws.services.glue.model.GetDatabasesResult; -import com.amazonaws.services.glue.model.GetTableRequest; -import com.amazonaws.services.glue.model.GetTablesRequest; -import com.amazonaws.services.glue.model.GetTablesResult; -import com.amazonaws.services.glue.model.TableInput; -import com.amazonaws.services.glue.model.UpdateTableRequest; import com.google.common.cache.Cache; +import com.google.common.collect.AbstractIterator; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.common.util.concurrent.UncheckedExecutionException; @@ -41,18 +22,23 @@ import dev.failsafe.RetryPolicy; import io.airlift.log.Logger; import io.trino.cache.EvictableCacheBuilder; +import io.trino.filesystem.Location; +import io.trino.filesystem.TrinoFileSystem; import io.trino.filesystem.TrinoFileSystemFactory; import io.trino.plugin.base.CatalogName; import io.trino.plugin.hive.SchemaAlreadyExistsException; import io.trino.plugin.hive.TrinoViewUtil; import io.trino.plugin.hive.ViewAlreadyExistsException; import io.trino.plugin.hive.ViewReaderUtil; +import io.trino.plugin.hive.metastore.TableInfo; import io.trino.plugin.hive.metastore.glue.GlueMetastoreStats; import io.trino.plugin.iceberg.IcebergMaterializedViewDefinition; import io.trino.plugin.iceberg.IcebergMetadata; import io.trino.plugin.iceberg.UnknownTableTypeException; import io.trino.plugin.iceberg.catalog.AbstractTrinoCatalog; +import io.trino.plugin.iceberg.catalog.IcebergTableOperations; import io.trino.plugin.iceberg.catalog.IcebergTableOperationsProvider; +import io.trino.plugin.iceberg.fileio.ForwardingFileIo; import io.trino.spi.TrinoException; import io.trino.spi.connector.CatalogSchemaTableName; import io.trino.spi.connector.ColumnMetadata; @@ -62,6 +48,7 @@ import io.trino.spi.connector.MaterializedViewNotFoundException; import io.trino.spi.connector.RelationColumnsMetadata; import io.trino.spi.connector.RelationCommentMetadata; +import io.trino.spi.connector.RelationType; import io.trino.spi.connector.SchemaNotFoundException; import io.trino.spi.connector.SchemaTableName; import io.trino.spi.connector.TableNotFoundException; @@ -75,22 +62,40 @@ import org.apache.iceberg.PartitionSpec; import org.apache.iceberg.Schema; import org.apache.iceberg.SortOrder; -import org.apache.iceberg.Table; import org.apache.iceberg.TableMetadata; import org.apache.iceberg.TableMetadataParser; import org.apache.iceberg.TableOperations; import org.apache.iceberg.Transaction; +import org.apache.iceberg.exceptions.NotFoundException; import org.apache.iceberg.io.FileIO; - +import software.amazon.awssdk.core.exception.SdkException; +import software.amazon.awssdk.services.glue.GlueClient; +import software.amazon.awssdk.services.glue.model.AccessDeniedException; +import software.amazon.awssdk.services.glue.model.AlreadyExistsException; +import software.amazon.awssdk.services.glue.model.Column; +import software.amazon.awssdk.services.glue.model.Database; +import software.amazon.awssdk.services.glue.model.DatabaseInput; +import software.amazon.awssdk.services.glue.model.EntityNotFoundException; +import software.amazon.awssdk.services.glue.model.GetDatabasesResponse; +import software.amazon.awssdk.services.glue.model.GetTablesResponse; +import software.amazon.awssdk.services.glue.model.StorageDescriptor; +import software.amazon.awssdk.services.glue.model.Table; +import software.amazon.awssdk.services.glue.model.TableInput; + +import java.io.IOException; import java.time.Duration; +import java.util.Collection; import java.util.HashMap; import java.util.Iterator; import java.util.List; import java.util.Map; +import java.util.Map.Entry; import java.util.Objects; import java.util.Optional; import java.util.Set; -import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.Callable; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.Executor; import java.util.function.Consumer; import java.util.function.Predicate; import java.util.function.UnaryOperator; @@ -100,34 +105,37 @@ import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.base.Preconditions.checkState; import static com.google.common.base.Throwables.throwIfInstanceOf; +import static com.google.common.base.Throwables.throwIfUnchecked; +import static com.google.common.base.Verify.verify; import static com.google.common.collect.ImmutableList.toImmutableList; import static com.google.common.collect.ImmutableSet.toImmutableSet; +import static com.google.common.collect.Streams.stream; import static io.trino.cache.CacheUtils.uncheckedCacheGet; import static io.trino.filesystem.Locations.appendPath; +import static io.trino.plugin.base.util.ExecutorUtil.processWithAdditionalThreads; import static io.trino.plugin.hive.HiveErrorCode.HIVE_DATABASE_LOCATION_ERROR; import static io.trino.plugin.hive.HiveErrorCode.HIVE_METASTORE_ERROR; import static io.trino.plugin.hive.HiveMetadata.STORAGE_TABLE; -import static io.trino.plugin.hive.HiveMetadata.TABLE_COMMENT; import static io.trino.plugin.hive.TableType.VIRTUAL_VIEW; import static io.trino.plugin.hive.TrinoViewUtil.createViewProperties; import static io.trino.plugin.hive.ViewReaderUtil.encodeViewData; import static io.trino.plugin.hive.ViewReaderUtil.isPrestoView; import static io.trino.plugin.hive.ViewReaderUtil.isTrinoMaterializedView; -import static io.trino.plugin.hive.metastore.glue.AwsSdkUtil.getPaginatedResults; -import static io.trino.plugin.hive.metastore.glue.converter.GlueToTrinoConverter.getColumnParameters; -import static io.trino.plugin.hive.metastore.glue.converter.GlueToTrinoConverter.getTableParameters; -import static io.trino.plugin.hive.metastore.glue.converter.GlueToTrinoConverter.getTableType; -import static io.trino.plugin.hive.metastore.glue.converter.GlueToTrinoConverter.getTableTypeNullable; +import static io.trino.plugin.hive.ViewReaderUtil.isTrinoView; +import static io.trino.plugin.hive.metastore.Table.TABLE_COMMENT; +import static io.trino.plugin.hive.metastore.glue.GlueConverter.getTableType; +import static io.trino.plugin.hive.metastore.glue.GlueConverter.getTableTypeNullable; import static io.trino.plugin.hive.util.HiveUtil.isHiveSystemSchema; import static io.trino.plugin.hive.util.HiveUtil.isIcebergTable; import static io.trino.plugin.iceberg.IcebergErrorCode.ICEBERG_BAD_DATA; import static io.trino.plugin.iceberg.IcebergErrorCode.ICEBERG_CATALOG_ERROR; import static io.trino.plugin.iceberg.IcebergErrorCode.ICEBERG_INVALID_METADATA; -import static io.trino.plugin.iceberg.IcebergMaterializedViewAdditionalProperties.STORAGE_SCHEMA; import static io.trino.plugin.iceberg.IcebergMaterializedViewDefinition.decodeMaterializedViewData; import static io.trino.plugin.iceberg.IcebergMaterializedViewDefinition.encodeMaterializedViewData; import static io.trino.plugin.iceberg.IcebergMaterializedViewDefinition.fromConnectorMaterializedViewDefinition; +import static io.trino.plugin.iceberg.IcebergMaterializedViewProperties.STORAGE_SCHEMA; import static io.trino.plugin.iceberg.IcebergSchemaProperties.LOCATION_PROPERTY; +import static io.trino.plugin.iceberg.IcebergTableName.tableNameWithType; import static io.trino.plugin.iceberg.IcebergUtil.COLUMN_TRINO_NOT_NULL_PROPERTY; import static io.trino.plugin.iceberg.IcebergUtil.COLUMN_TRINO_TYPE_ID_PROPERTY; import static io.trino.plugin.iceberg.IcebergUtil.TRINO_TABLE_METADATA_INFO_VALID_FOR; @@ -136,6 +144,7 @@ import static io.trino.plugin.iceberg.IcebergUtil.getTableComment; import static io.trino.plugin.iceberg.IcebergUtil.quotedTableName; import static io.trino.plugin.iceberg.IcebergUtil.validateTableCanBeDropped; +import static io.trino.plugin.iceberg.TableType.MATERIALIZED_VIEW_STORAGE; import static io.trino.plugin.iceberg.TrinoMetricsReporter.TRINO_METRICS_REPORTER; import static io.trino.plugin.iceberg.catalog.glue.GlueIcebergUtil.getMaterializedViewTableInput; import static io.trino.plugin.iceberg.catalog.glue.GlueIcebergUtil.getTableInput; @@ -148,6 +157,7 @@ import static java.lang.Boolean.parseBoolean; import static java.lang.String.format; import static java.util.Locale.ENGLISH; +import static java.util.Map.entry; import static java.util.Objects.requireNonNull; import static org.apache.iceberg.BaseMetastoreTableOperations.METADATA_LOCATION_PROP; import static org.apache.iceberg.CatalogUtil.dropTableData; @@ -160,20 +170,28 @@ public class TrinoGlueCatalog private static final int PER_QUERY_CACHE_SIZE = 1000; private final String trinoVersion; - private final TypeManager typeManager; private final boolean cacheTableMetadata; private final TrinoFileSystemFactory fileSystemFactory; private final Optional defaultSchemaLocation; - private final AWSGlueAsync glueClient; + private final GlueClient glueClient; private final GlueMetastoreStats stats; + private final boolean hideMaterializedViewStorageTable; + private final boolean isUsingSystemSecurity; + private final Executor metadataFetchingExecutor; - private final Cache glueTableCache = EvictableCacheBuilder.newBuilder() + private final Cache glueTableCache = EvictableCacheBuilder.newBuilder() // Even though this is query-scoped, this still needs to be bounded. information_schema queries can access large number of tables. .maximumSize(Math.max(PER_QUERY_CACHE_SIZE, IcebergMetadata.GET_METADATA_BATCH_SIZE)) .build(); - private final Map tableMetadataCache = new ConcurrentHashMap<>(); - private final Map viewCache = new ConcurrentHashMap<>(); - private final Map materializedViewCache = new ConcurrentHashMap<>(); + private final Cache tableMetadataCache = EvictableCacheBuilder.newBuilder() + .maximumSize(PER_QUERY_CACHE_SIZE) + .build(); + private final Cache viewCache = EvictableCacheBuilder.newBuilder() + .maximumSize(PER_QUERY_CACHE_SIZE) + .build(); + private final Cache materializedViewCache = EvictableCacheBuilder.newBuilder() + .maximumSize(PER_QUERY_CACHE_SIZE) + .build(); public TrinoGlueCatalog( CatalogName catalogName, @@ -182,19 +200,24 @@ public TrinoGlueCatalog( boolean cacheTableMetadata, IcebergTableOperationsProvider tableOperationsProvider, String trinoVersion, - AWSGlueAsync glueClient, + GlueClient glueClient, GlueMetastoreStats stats, + boolean isUsingSystemSecurity, Optional defaultSchemaLocation, - boolean useUniqueTableLocation) + boolean useUniqueTableLocation, + boolean hideMaterializedViewStorageTable, + Executor metadataFetchingExecutor) { - super(catalogName, typeManager, tableOperationsProvider, useUniqueTableLocation); + super(catalogName, typeManager, tableOperationsProvider, fileSystemFactory, useUniqueTableLocation); this.trinoVersion = requireNonNull(trinoVersion, "trinoVersion is null"); - this.typeManager = requireNonNull(typeManager, "typeManager is null"); this.cacheTableMetadata = cacheTableMetadata; this.fileSystemFactory = requireNonNull(fileSystemFactory, "fileSystemFactory is null"); this.glueClient = requireNonNull(glueClient, "glueClient is null"); this.stats = requireNonNull(stats, "stats is null"); + this.isUsingSystemSecurity = isUsingSystemSecurity; this.defaultSchemaLocation = requireNonNull(defaultSchemaLocation, "defaultSchemaLocation is null"); + this.hideMaterializedViewStorageTable = hideMaterializedViewStorageTable; + this.metadataFetchingExecutor = requireNonNull(metadataFetchingExecutor, "metadataFetchingExecutor is null"); } @Override @@ -207,13 +230,13 @@ public boolean namespaceExists(ConnectorSession session, String namespace) } return stats.getGetDatabase().call(() -> { try { - glueClient.getDatabase(new GetDatabaseRequest().withName(namespace)); + glueClient.getDatabase(x -> x.name(namespace)); return true; } catch (EntityNotFoundException e) { return false; } - catch (AmazonServiceException e) { + catch (SdkException e) { throw new TrinoException(ICEBERG_CATALOG_ERROR, e); } }); @@ -223,18 +246,14 @@ public boolean namespaceExists(ConnectorSession session, String namespace) public List listNamespaces(ConnectorSession session) { try { - return getPaginatedResults( - glueClient::getDatabases, - new GetDatabasesRequest(), - GetDatabasesRequest::setNextToken, - GetDatabasesResult::getNextToken, - stats.getGetDatabases()) - .map(GetDatabasesResult::getDatabaseList) - .flatMap(List::stream) - .map(com.amazonaws.services.glue.model.Database::getName) - .collect(toImmutableList()); + return stats.getGetDatabases().call(() -> + glueClient.getDatabasesPaginator(x -> {}).stream() + .map(GetDatabasesResponse::databaseList) + .flatMap(List::stream) + .map(Database::name) + .collect(toImmutableList())); } - catch (AmazonServiceException e) { + catch (SdkException e) { throw new TrinoException(ICEBERG_CATALOG_ERROR, e); } } @@ -253,12 +272,12 @@ public void dropNamespace(ConnectorSession session, String namespace) try { glueTableCache.invalidateAll(); stats.getDeleteDatabase().call(() -> - glueClient.deleteDatabase(new DeleteDatabaseRequest().withName(namespace))); + glueClient.deleteDatabase(x -> x.name(namespace))); } catch (EntityNotFoundException e) { - throw new SchemaNotFoundException(namespace); + throw new SchemaNotFoundException(namespace, e); } - catch (AmazonServiceException e) { + catch (SdkException e) { throw new TrinoException(ICEBERG_CATALOG_ERROR, e); } } @@ -267,22 +286,21 @@ public void dropNamespace(ConnectorSession session, String namespace) public Map loadNamespaceMetadata(ConnectorSession session, String namespace) { try { - GetDatabaseRequest getDatabaseRequest = new GetDatabaseRequest().withName(namespace); Database database = stats.getGetDatabase().call(() -> - glueClient.getDatabase(getDatabaseRequest).getDatabase()); + glueClient.getDatabase(x -> x.name(namespace)).database()); ImmutableMap.Builder metadata = ImmutableMap.builder(); - if (database.getLocationUri() != null) { - metadata.put(LOCATION_PROPERTY, database.getLocationUri()); + if (database.locationUri() != null) { + metadata.put(LOCATION_PROPERTY, database.locationUri()); } - if (database.getParameters() != null) { - metadata.putAll(database.getParameters()); + if (database.parameters() != null) { + metadata.putAll(database.parameters()); } return metadata.buildOrThrow(); } catch (EntityNotFoundException e) { - throw new SchemaNotFoundException(namespace); + throw new SchemaNotFoundException(namespace, e); } - catch (AmazonServiceException e) { + catch (SdkException e) { throw new TrinoException(ICEBERG_CATALOG_ERROR, e); } } @@ -301,28 +319,28 @@ public void createNamespace(ConnectorSession session, String namespace, Map - glueClient.createDatabase(new CreateDatabaseRequest() - .withDatabaseInput(createDatabaseInput(namespace, properties)))); + glueClient.createDatabase(x -> x + .databaseInput(createDatabaseInput(namespace, properties)))); } catch (AlreadyExistsException e) { - throw new SchemaAlreadyExistsException(namespace); + throw new SchemaAlreadyExistsException(namespace, e); } - catch (AmazonServiceException e) { + catch (SdkException e) { throw new TrinoException(ICEBERG_CATALOG_ERROR, e); } } private DatabaseInput createDatabaseInput(String namespace, Map properties) { - DatabaseInput databaseInput = new DatabaseInput().withName(namespace); + DatabaseInput.Builder databaseInput = DatabaseInput.builder().name(namespace); properties.forEach((property, value) -> { switch (property) { - case LOCATION_PROPERTY -> databaseInput.setLocationUri((String) value); + case LOCATION_PROPERTY -> databaseInput.locationUri((String) value); default -> throw new IllegalArgumentException("Unrecognized property: " + property); } }); - return databaseInput; + return databaseInput.build(); } @Override @@ -338,27 +356,45 @@ public void renameNamespace(ConnectorSession session, String source, String targ } @Override - public List listTables(ConnectorSession session, Optional namespace) + public List listTables(ConnectorSession session, Optional namespace) + { + return listTables(session, namespace, ignore -> true); + } + + @Override + public List listIcebergTables(ConnectorSession session, Optional namespace) { - ImmutableList.Builder tables = ImmutableList.builder(); + return listTables(session, namespace, table -> isIcebergTable(table.parameters())).stream() + .map(TableInfo::tableName) + .collect(toImmutableList()); + } + + private List listTables( + ConnectorSession session, + Optional namespace, + Predicate

tablePredicate) + { + List>> tasks = listNamespaces(session, namespace).stream() + .map(glueNamespace -> (Callable>) () -> getGlueTablesWithExceptionHandling(glueNamespace) + .filter(tablePredicate) + .map(table -> mapToTableInfo(glueNamespace, table)) + .collect(toImmutableList())) + .collect(toImmutableList()); try { - List namespaces = listNamespaces(session, namespace); - for (String glueNamespace : namespaces) { - try { - // Add all tables from a namespace together, in case it is removed while fetching paginated results - tables.addAll(getGlueTables(glueNamespace) - .map(table -> new SchemaTableName(glueNamespace, table.getName())) - .collect(toImmutableList())); - } - catch (EntityNotFoundException | AccessDeniedException e) { - // Namespace may have been deleted or permission denied - } - } + return processWithAdditionalThreads(tasks, metadataFetchingExecutor).stream() + .flatMap(Collection::stream) + .collect(toImmutableList()); } - catch (AmazonServiceException e) { - throw new TrinoException(ICEBERG_CATALOG_ERROR, e); + catch (ExecutionException e) { + throw new RuntimeException(e.getCause()); } - return tables.build(); + } + + private TableInfo mapToTableInfo(String glueNamespace, Table table) + { + return new TableInfo( + new SchemaTableName(glueNamespace, table.name()), + TableInfo.ExtendedRelationType.fromTableTypeAndComment(getTableType(table), table.parameters().get(TABLE_COMMENT))); } @Override @@ -370,22 +406,22 @@ public Optional> streamRelationColumns( { ImmutableList.Builder unfilteredResult = ImmutableList.builder(); ImmutableList.Builder filteredResult = ImmutableList.builder(); - Map unprocessed = new HashMap<>(); + Map unprocessed = new HashMap<>(); listNamespaces(session, namespace).stream() .flatMap(glueNamespace -> getGlueTables(glueNamespace) - .map(table -> Map.entry(new SchemaTableName(glueNamespace, table.getName()), table))) + .map(table -> entry(new SchemaTableName(glueNamespace, table.name()), table))) .forEach(entry -> { SchemaTableName name = entry.getKey(); - com.amazonaws.services.glue.model.Table table = entry.getValue(); + Table table = entry.getValue(); String tableType = getTableType(table); - Map tableParameters = getTableParameters(table); + Map tableParameters = table.parameters(); if (isTrinoMaterializedView(tableType, tableParameters)) { - IcebergMaterializedViewDefinition definition = decodeMaterializedViewData(table.getViewOriginalText()); + IcebergMaterializedViewDefinition definition = decodeMaterializedViewData(table.viewOriginalText()); unfilteredResult.add(RelationColumnsMetadata.forMaterializedView(name, toSpiMaterializedViewColumns(definition.getColumns()))); } else if (isPrestoView(tableParameters)) { - ConnectorViewDefinition definition = ViewReaderUtil.PrestoViewReader.decodeViewData(table.getViewOriginalText()); + ConnectorViewDefinition definition = ViewReaderUtil.PrestoViewReader.decodeViewData(table.viewOriginalText()); unfilteredResult.add(RelationColumnsMetadata.forView(name, definition.getColumns())); } else if (isRedirected.test(name)) { @@ -427,12 +463,12 @@ else if (!isIcebergTable(tableParameters)) { private void getColumnsFromIcebergMetadata( ConnectorSession session, - Map glueTables, // only Iceberg tables + Map glueTables, // only Iceberg tables UnaryOperator> relationFilter, Consumer resultsCollector) { for (SchemaTableName tableName : relationFilter.apply(glueTables.keySet())) { - com.amazonaws.services.glue.model.Table table = glueTables.get(tableName); + Table table = glueTables.get(tableName); // potentially racy with invalidation, but TrinoGlueCatalog is session-scoped uncheckedCacheGet(glueTableCache, tableName, () -> table); List columns; @@ -461,22 +497,22 @@ public Optional> streamRelationComments( ImmutableList.Builder unfilteredResult = ImmutableList.builder(); ImmutableList.Builder filteredResult = ImmutableList.builder(); - Map unprocessed = new HashMap<>(); + Map unprocessed = new HashMap<>(); listNamespaces(session, namespace).stream() .flatMap(glueNamespace -> getGlueTables(glueNamespace) - .map(table -> Map.entry(new SchemaTableName(glueNamespace, table.getName()), table))) + .map(table -> entry(new SchemaTableName(glueNamespace, table.name()), table))) .forEach(entry -> { SchemaTableName name = entry.getKey(); - com.amazonaws.services.glue.model.Table table = entry.getValue(); + Table table = entry.getValue(); String tableType = getTableType(table); - Map tableParameters = getTableParameters(table); + Map tableParameters = table.parameters(); if (isTrinoMaterializedView(tableType, tableParameters)) { - Optional comment = decodeMaterializedViewData(table.getViewOriginalText()).getComment(); + Optional comment = decodeMaterializedViewData(table.viewOriginalText()).getComment(); unfilteredResult.add(RelationCommentMetadata.forTable(name, comment)); } else if (isPrestoView(tableParameters)) { - Optional comment = ViewReaderUtil.PrestoViewReader.decodeViewData(table.getViewOriginalText()).getComment(); + Optional comment = ViewReaderUtil.PrestoViewReader.decodeViewData(table.viewOriginalText()).getComment(); unfilteredResult.add(RelationCommentMetadata.forTable(name, comment)); } else if (isRedirected.test(name)) { @@ -521,12 +557,12 @@ else if (!isIcebergTable(tableParameters)) { private void getCommentsFromIcebergMetadata( ConnectorSession session, - Map glueTables, // only Iceberg tables + Map glueTables, // only Iceberg tables UnaryOperator> relationFilter, Consumer resultsCollector) { for (SchemaTableName tableName : relationFilter.apply(glueTables.keySet())) { - com.amazonaws.services.glue.model.Table table = glueTables.get(tableName); + Table table = glueTables.get(tableName); // potentially racy with invalidation, but TrinoGlueCatalog is session-scoped uncheckedCacheGet(glueTableCache, tableName, () -> table); Optional comment; @@ -543,24 +579,32 @@ private void getCommentsFromIcebergMetadata( } @Override - public Table loadTable(ConnectorSession session, SchemaTableName table) + public BaseTable loadTable(ConnectorSession session, SchemaTableName table) { - if (viewCache.containsKey(table) || materializedViewCache.containsKey(table)) { + if (viewCache.asMap().containsKey(table) || materializedViewCache.asMap().containsKey(table)) { throw new TableNotFoundException(table); } - TableMetadata metadata = tableMetadataCache.computeIfAbsent( - table, - ignore -> { - TableOperations operations = tableOperationsProvider.createTableOperations( - this, - session, - table.getSchemaName(), - table.getTableName(), - Optional.empty(), - Optional.empty()); - return new BaseTable(operations, quotedTableName(table), TRINO_METRICS_REPORTER).operations().current(); - }); + TableMetadata metadata; + try { + metadata = uncheckedCacheGet( + tableMetadataCache, + table, + () -> { + TableOperations operations = tableOperationsProvider.createTableOperations( + this, + session, + table.getSchemaName(), + table.getTableName(), + Optional.empty(), + Optional.empty()); + return new BaseTable(operations, quotedTableName(table), TRINO_METRICS_REPORTER).operations().current(); + }); + } + catch (UncheckedExecutionException e) { + throwIfUnchecked(e.getCause()); + throw e; + } return getIcebergTableWithMetadata( this, @@ -599,45 +643,45 @@ public Map> tryGetColumnMetadata(Connector private Optional> getCachedColumnMetadata(SchemaTableName tableName) { - if (!cacheTableMetadata || viewCache.containsKey(tableName) || materializedViewCache.containsKey(tableName)) { + if (!cacheTableMetadata || viewCache.asMap().containsKey(tableName) || materializedViewCache.asMap().containsKey(tableName)) { return Optional.empty(); } - com.amazonaws.services.glue.model.Table glueTable = getTable(tableName, false); + Table glueTable = getTable(tableName, false); return getCachedColumnMetadata(glueTable); } - private Optional> getCachedColumnMetadata(com.amazonaws.services.glue.model.Table glueTable) + private Optional> getCachedColumnMetadata(Table glueTable) { if (!cacheTableMetadata) { return Optional.empty(); } - Map tableParameters = getTableParameters(glueTable); + Map tableParameters = glueTable.parameters(); String metadataLocation = tableParameters.get(METADATA_LOCATION_PROP); String metadataValidForMetadata = tableParameters.get(TRINO_TABLE_METADATA_INFO_VALID_FOR); if (metadataLocation == null || !metadataLocation.equals(metadataValidForMetadata) || - glueTable.getStorageDescriptor() == null || - glueTable.getStorageDescriptor().getColumns() == null) { + glueTable.storageDescriptor() == null || + glueTable.storageDescriptor().columns() == null) { return Optional.empty(); } - List glueColumns = glueTable.getStorageDescriptor().getColumns(); - if (glueColumns.stream().noneMatch(column -> getColumnParameters(column).containsKey(COLUMN_TRINO_TYPE_ID_PROPERTY))) { + List glueColumns = glueTable.storageDescriptor().columns(); + if (glueColumns.stream().noneMatch(column -> column.parameters().containsKey(COLUMN_TRINO_TYPE_ID_PROPERTY))) { // No column has type parameter, maybe the parameters were erased return Optional.empty(); } ImmutableList.Builder columns = ImmutableList.builderWithExpectedSize(glueColumns.size()); for (Column glueColumn : glueColumns) { - Map columnParameters = getColumnParameters(glueColumn); - String trinoTypeId = columnParameters.getOrDefault(COLUMN_TRINO_TYPE_ID_PROPERTY, glueColumn.getType()); + Map columnParameters = glueColumn.parameters(); + String trinoTypeId = columnParameters.getOrDefault(COLUMN_TRINO_TYPE_ID_PROPERTY, glueColumn.type()); boolean notNull = parseBoolean(columnParameters.getOrDefault(COLUMN_TRINO_NOT_NULL_PROPERTY, "false")); Type type = typeManager.getType(TypeId.of(trinoTypeId)); columns.add(ColumnMetadata.builder() - .setName(glueColumn.getName()) + .setName(glueColumn.name()) .setType(type) - .setComment(Optional.ofNullable(glueColumn.getComment())) + .setComment(Optional.ofNullable(glueColumn.comment())) .setNullable(!notNull) .build()); } @@ -652,7 +696,7 @@ public void dropTable(ConnectorSession session, SchemaTableName schemaTableName) try { deleteTable(schemaTableName.getSchemaName(), schemaTableName.getTableName()); } - catch (AmazonServiceException e) { + catch (SdkException e) { throw new TrinoException(HIVE_METASTORE_ERROR, e); } try { @@ -664,18 +708,20 @@ public void dropTable(ConnectorSession session, SchemaTableName schemaTableName) LOG.warn(e, "Failed to delete table data referenced by metadata"); } deleteTableDirectory(fileSystemFactory.create(session), schemaTableName, table.location()); + invalidateTableCache(schemaTableName); } @Override public void dropCorruptedTable(ConnectorSession session, SchemaTableName schemaTableName) { - com.amazonaws.services.glue.model.Table table = dropTableFromMetastore(session, schemaTableName); - String metadataLocation = getTableParameters(table).get(METADATA_LOCATION_PROP); + Table table = dropTableFromMetastore(session, schemaTableName); + String metadataLocation = table.parameters().get(METADATA_LOCATION_PROP); if (metadataLocation == null) { throw new TrinoException(ICEBERG_INVALID_METADATA, format("Table %s is missing [%s] property", schemaTableName, METADATA_LOCATION_PROP)); } String tableLocation = metadataLocation.replaceFirst("/metadata/[^/]*$", ""); deleteTableDirectory(fileSystemFactory.create(session), schemaTableName, tableLocation); + invalidateTableCache(schemaTableName); } @Override @@ -685,7 +731,7 @@ public Transaction newCreateTableTransaction( Schema schema, PartitionSpec partitionSpec, SortOrder sortOrder, - String location, + Optional location, Map properties) { return newCreateTableTransaction( @@ -699,6 +745,27 @@ public Transaction newCreateTableTransaction( Optional.of(session.getUser())); } + @Override + public Transaction newCreateOrReplaceTableTransaction( + ConnectorSession session, + SchemaTableName schemaTableName, + Schema schema, + PartitionSpec partitionSpec, + SortOrder sortOrder, + String location, + Map properties) + { + return newCreateOrReplaceTableTransaction( + session, + schemaTableName, + schema, + partitionSpec, + sortOrder, + location, + properties, + Optional.of(session.getUser())); + } + @Override public void registerTable(ConnectorSession session, SchemaTableName schemaTableName, TableMetadata tableMetadata) throws TrinoException @@ -708,6 +775,7 @@ public void registerTable(ConnectorSession session, SchemaTableName schemaTableN schemaTableName.getTableName(), Optional.of(session.getUser()), tableMetadata, + tableMetadata.location(), tableMetadata.metadataFileLocation(), ImmutableMap.of(), cacheTableMetadata); @@ -718,20 +786,21 @@ public void registerTable(ConnectorSession session, SchemaTableName schemaTableN public void unregisterTable(ConnectorSession session, SchemaTableName schemaTableName) { dropTableFromMetastore(session, schemaTableName); + invalidateTableCache(schemaTableName); } - private com.amazonaws.services.glue.model.Table dropTableFromMetastore(ConnectorSession session, SchemaTableName schemaTableName) + private Table dropTableFromMetastore(ConnectorSession session, SchemaTableName schemaTableName) { - com.amazonaws.services.glue.model.Table table = getTableAndCacheMetadata(session, schemaTableName) + Table table = getTableAndCacheMetadata(session, schemaTableName) .orElseThrow(() -> new TableNotFoundException(schemaTableName)); - if (!isIcebergTable(getTableParameters(table))) { + if (!isIcebergTable(table.parameters())) { throw new UnknownTableTypeException(schemaTableName); } try { deleteTable(schemaTableName.getSchemaName(), schemaTableName.getTableName()); } - catch (AmazonServiceException e) { + catch (SdkException e) { throw new TrinoException(HIVE_METASTORE_ERROR, e); } return table; @@ -742,9 +811,9 @@ public void renameTable(ConnectorSession session, SchemaTableName from, SchemaTa { boolean newTableCreated = false; try { - com.amazonaws.services.glue.model.Table table = getTableAndCacheMetadata(session, from) + Table table = getTableAndCacheMetadata(session, from) .orElseThrow(() -> new TableNotFoundException(from)); - Map tableParameters = new HashMap<>(getTableParameters(table)); + Map tableParameters = new HashMap<>(table.parameters()); FileIO io = loadTable(session, from).io(); String metadataLocation = tableParameters.remove(METADATA_LOCATION_PROP); if (metadataLocation == null) { @@ -754,14 +823,16 @@ public void renameTable(ConnectorSession session, SchemaTableName from, SchemaTa TableInput tableInput = getTableInput( typeManager, to.getTableName(), - Optional.ofNullable(table.getOwner()), + Optional.ofNullable(table.owner()), metadata, + Optional.ofNullable(table.storageDescriptor()).map(StorageDescriptor::location).orElse(null), metadataLocation, tableParameters, cacheTableMetadata); createTable(to.getSchemaName(), tableInput); newTableCreated = true; deleteTable(from.getSchemaName(), from.getTableName()); + invalidateTableCache(from); } catch (RuntimeException e) { if (newTableCreated) { @@ -778,9 +849,9 @@ public void renameTable(ConnectorSession session, SchemaTableName from, SchemaTa } } - private Optional getTableAndCacheMetadata(ConnectorSession session, SchemaTableName schemaTableName) + private Optional
getTableAndCacheMetadata(ConnectorSession session, SchemaTableName schemaTableName) { - com.amazonaws.services.glue.model.Table table; + Table table; try { table = getTable(schemaTableName, false); } @@ -788,54 +859,56 @@ private Optional getTableAndCacheMetada return Optional.empty(); } - Map parameters = getTableParameters(table); - if (isIcebergTable(parameters) && !tableMetadataCache.containsKey(schemaTableName)) { - if (viewCache.containsKey(schemaTableName) || materializedViewCache.containsKey(schemaTableName)) { + String tableType = getTableType(table); + Map parameters = table.parameters(); + if (isIcebergTable(parameters) && !tableMetadataCache.asMap().containsKey(schemaTableName)) { + if (viewCache.asMap().containsKey(schemaTableName) || materializedViewCache.asMap().containsKey(schemaTableName)) { throw new TrinoException(GENERIC_INTERNAL_ERROR, "Glue table cache inconsistency. Table cannot also be a view/materialized view"); } String metadataLocation = parameters.get(METADATA_LOCATION_PROP); try { // Cache the TableMetadata while we have the Table retrieved anyway - TableOperations operations = tableOperationsProvider.createTableOperations( - this, - session, - schemaTableName.getSchemaName(), - schemaTableName.getTableName(), - Optional.empty(), - Optional.empty()); - FileIO io = operations.io(); - tableMetadataCache.put(schemaTableName, TableMetadataParser.read(io, io.newInputFile(metadataLocation))); + // Note: this is racy from cache invalidation perspective, but it should not matter here + uncheckedCacheGet(tableMetadataCache, schemaTableName, () -> TableMetadataParser.read(new ForwardingFileIo(fileSystemFactory.create(session)), metadataLocation)); } catch (RuntimeException e) { LOG.warn(e, "Failed to cache table metadata from table at %s", metadataLocation); } } else if (isTrinoMaterializedView(getTableType(table), parameters)) { - if (viewCache.containsKey(schemaTableName) || tableMetadataCache.containsKey(schemaTableName)) { + if (viewCache.asMap().containsKey(schemaTableName) || tableMetadataCache.asMap().containsKey(schemaTableName)) { throw new TrinoException(GENERIC_INTERNAL_ERROR, "Glue table cache inconsistency. Materialized View cannot also be a table or view"); } try { - createMaterializedViewDefinition(session, schemaTableName, table) - .ifPresent(materializedView -> materializedViewCache.put(schemaTableName, materializedView)); + // Note: this is racy from cache invalidation perspective, but it should not matter here + uncheckedCacheGet(materializedViewCache, schemaTableName, () -> { + ConnectorMaterializedViewDefinition materializedView = createMaterializedViewDefinition(session, schemaTableName, table); + return new MaterializedViewData( + materializedView, + Optional.ofNullable(parameters.get(METADATA_LOCATION_PROP))); + }); } catch (RuntimeException e) { LOG.warn(e, "Failed to cache materialized view from %s", schemaTableName); } } - else if (isPrestoView(parameters) && !viewCache.containsKey(schemaTableName)) { - if (materializedViewCache.containsKey(schemaTableName) || tableMetadataCache.containsKey(schemaTableName)) { + else if (isTrinoView(tableType, parameters) && !viewCache.asMap().containsKey(schemaTableName)) { + if (materializedViewCache.asMap().containsKey(schemaTableName) || tableMetadataCache.asMap().containsKey(schemaTableName)) { throw new TrinoException(GENERIC_INTERNAL_ERROR, "Glue table cache inconsistency. View cannot also be a materialized view or table"); } try { - TrinoViewUtil.getView(schemaTableName, - Optional.ofNullable(table.getViewOriginalText()), - getTableType(table), + TrinoViewUtil.getView( + Optional.ofNullable(table.viewOriginalText()), + tableType, parameters, - Optional.ofNullable(table.getOwner())) - .ifPresent(viewDefinition -> viewCache.put(schemaTableName, viewDefinition)); + Optional.ofNullable(table.owner())) + .ifPresent(viewDefinition -> { + // Note: this is racy from cache invalidation perspective, but it should not matter here + uncheckedCacheGet(viewCache, schemaTableName, () -> viewDefinition); + }); } catch (RuntimeException e) { LOG.warn(e, "Failed to cache view from %s", schemaTableName); @@ -848,12 +921,9 @@ else if (isPrestoView(parameters) && !viewCache.containsKey(schemaTableName)) { @Override public String defaultTableLocation(ConnectorSession session, SchemaTableName schemaTableName) { - GetDatabaseRequest getDatabaseRequest = new GetDatabaseRequest() - .withName(schemaTableName.getSchemaName()); String databaseLocation = stats.getGetDatabase().call(() -> - glueClient.getDatabase(getDatabaseRequest) - .getDatabase() - .getLocationUri()); + glueClient.getDatabase(x -> x.name(schemaTableName.getSchemaName())) + .database().locationUri()); String tableName = createNewTableName(schemaTableName.getTableName()); @@ -898,9 +968,9 @@ public void createView(ConnectorSession session, SchemaTableName schemaViewName, private void doCreateView(ConnectorSession session, SchemaTableName schemaViewName, TableInput viewTableInput, boolean replace) { - Optional existing = getTableAndCacheMetadata(session, schemaViewName); + Optional
existing = getTableAndCacheMetadata(session, schemaViewName); if (existing.isPresent()) { - if (!replace || !isPrestoView(getTableParameters(existing.get()))) { + if (!replace || !isTrinoView(getTableType(existing.get()), existing.get().parameters())) { // TODO: ViewAlreadyExists is misleading if the name is used by a table https://github.com/trinodb/trino/issues/10037 throw new ViewAlreadyExistsException(schemaViewName); } @@ -922,13 +992,13 @@ public void renameView(ConnectorSession session, SchemaTableName source, SchemaT { boolean newTableCreated = false; try { - com.amazonaws.services.glue.model.Table existingView = getTableAndCacheMetadata(session, source) + Table existingView = getTableAndCacheMetadata(session, source) .orElseThrow(() -> new TableNotFoundException(source)); - viewCache.remove(source); + viewCache.invalidate(source); TableInput viewTableInput = getViewTableInput( target.getTableName(), - existingView.getViewOriginalText(), - existingView.getOwner(), + existingView.viewOriginalText(), + existingView.owner(), createViewProperties(session, trinoVersion, TRINO_CREATED_BY_VALUE)); createTable(target.getSchemaName(), viewTableInput); newTableCreated = true; @@ -963,10 +1033,10 @@ public void dropView(ConnectorSession session, SchemaTableName schemaViewName) } try { - viewCache.remove(schemaViewName); + viewCache.invalidate(schemaViewName); deleteTable(schemaViewName.getSchemaName(), schemaViewName.getTableName()); } - catch (AmazonServiceException e) { + catch (SdkException e) { throw new TrinoException(HIVE_METASTORE_ERROR, e); } } @@ -974,51 +1044,35 @@ public void dropView(ConnectorSession session, SchemaTableName schemaViewName) @Override public List listViews(ConnectorSession session, Optional namespace) { - ImmutableList.Builder views = ImmutableList.builder(); - try { - List namespaces = listNamespaces(session, namespace); - for (String glueNamespace : namespaces) { - try { - views.addAll(getGlueTables(glueNamespace) - .filter(table -> isPrestoView(getTableParameters(table)) && !isTrinoMaterializedView(getTableType(table), getTableParameters(table))) // TODO isTrinoMaterializedView should not be needed, isPrestoView should not return true for materialized views - .map(table -> new SchemaTableName(glueNamespace, table.getName())) - .collect(toImmutableList())); - } - catch (EntityNotFoundException | AccessDeniedException e) { - // Namespace may have been deleted or permission denied - } - } - } - catch (AmazonServiceException e) { - throw new TrinoException(ICEBERG_CATALOG_ERROR, e); - } - return views.build(); + return listRelations(session, namespace) + .filter(entry -> entry.getValue() == RelationType.VIEW) + .map(Entry::getKey) + .collect(toImmutableList()); } @Override public Optional getView(ConnectorSession session, SchemaTableName viewName) { - ConnectorViewDefinition cachedView = viewCache.get(viewName); + ConnectorViewDefinition cachedView = viewCache.getIfPresent(viewName); if (cachedView != null) { return Optional.of(cachedView); } - if (tableMetadataCache.containsKey(viewName) || materializedViewCache.containsKey(viewName)) { + if (tableMetadataCache.asMap().containsKey(viewName) || materializedViewCache.asMap().containsKey(viewName)) { // Entries in these caches are not views return Optional.empty(); } - Optional table = getTableAndCacheMetadata(session, viewName); + Optional
table = getTableAndCacheMetadata(session, viewName); if (table.isEmpty()) { return Optional.empty(); } - com.amazonaws.services.glue.model.Table viewDefinition = table.get(); + Table viewDefinition = table.get(); return TrinoViewUtil.getView( - viewName, - Optional.ofNullable(viewDefinition.getViewOriginalText()), + Optional.ofNullable(viewDefinition.viewOriginalText()), getTableType(viewDefinition), - getTableParameters(viewDefinition), - Optional.ofNullable(viewDefinition.getOwner())); + viewDefinition.parameters(), + Optional.ofNullable(viewDefinition.owner())); } @Override @@ -1068,47 +1122,24 @@ private void updateView(ConnectorSession session, SchemaTableName viewName, Conn try { updateTable(viewName.getSchemaName(), viewTableInput); } - catch (AmazonServiceException e) { + catch (SdkException e) { throw new TrinoException(ICEBERG_CATALOG_ERROR, e); } } - @Override - public List listMaterializedViews(ConnectorSession session, Optional namespace) - { - ImmutableList.Builder materializedViews = ImmutableList.builder(); - try { - List namespaces = listNamespaces(session, namespace); - for (String glueNamespace : namespaces) { - try { - materializedViews.addAll(getGlueTables(glueNamespace) - .filter(table -> isTrinoMaterializedView(getTableType(table), getTableParameters(table))) - .map(table -> new SchemaTableName(glueNamespace, table.getName())) - .collect(toImmutableList())); - } - catch (EntityNotFoundException | AccessDeniedException e) { - // Namespace may have been deleted or permission denied - } - } - } - catch (AmazonServiceException e) { - throw new TrinoException(ICEBERG_CATALOG_ERROR, e); - } - return materializedViews.build(); - } - @Override public void createMaterializedView( ConnectorSession session, SchemaTableName viewName, ConnectorMaterializedViewDefinition definition, + Map materializedViewProperties, boolean replace, boolean ignoreExisting) { - Optional existing = getTableAndCacheMetadata(session, viewName); + Optional
existing = getTableAndCacheMetadata(session, viewName); if (existing.isPresent()) { - if (!isTrinoMaterializedView(getTableType(existing.get()), getTableParameters(existing.get()))) { + if (!isTrinoMaterializedView(getTableType(existing.get()), existing.get().parameters())) { throw new TrinoException(UNSUPPORTED_TABLE_TYPE, "Existing table is not a Materialized View: " + viewName); } if (!replace) { @@ -1119,13 +1150,39 @@ public void createMaterializedView( } } + if (hideMaterializedViewStorageTable) { + Location storageMetadataLocation = createMaterializedViewStorage(session, viewName, definition, Map.of()); + TableInput materializedViewTableInput = getMaterializedViewTableInput( + viewName.getTableName(), + encodeMaterializedViewData(fromConnectorMaterializedViewDefinition(definition)), + isUsingSystemSecurity ? null : session.getUser(), + createMaterializedViewProperties(session, storageMetadataLocation)); + if (existing.isPresent()) { + updateTable(viewName.getSchemaName(), materializedViewTableInput); + } + else { + createTable(viewName.getSchemaName(), materializedViewTableInput); + } + } + else { + createMaterializedViewWithStorageTable(session, viewName, definition, Map.of(), existing); + } + } + + private void createMaterializedViewWithStorageTable( + ConnectorSession session, + SchemaTableName viewName, + ConnectorMaterializedViewDefinition definition, + Map materializedViewProperties, + Optional
existing) + { // Create the storage table - SchemaTableName storageTable = createMaterializedViewStorageTable(session, viewName, definition); + SchemaTableName storageTable = createMaterializedViewStorageTable(session, viewName, definition, materializedViewProperties); // Create a view indicating the storage table TableInput materializedViewTableInput = getMaterializedViewTableInput( viewName.getTableName(), encodeMaterializedViewData(fromConnectorMaterializedViewDefinition(definition)), - session.getUser(), + isUsingSystemSecurity ? null : session.getUser(), createMaterializedViewProperties(session, storageTable)); if (existing.isPresent()) { @@ -1144,7 +1201,7 @@ public void createMaterializedView( } } } - dropStorageTable(session, existing.get()); + dropMaterializedViewStorage(session, existing.get()); } else { createTable(viewName.getSchemaName(), materializedViewTableInput); @@ -1169,21 +1226,21 @@ public void updateMaterializedViewColumnComment(ConnectorSession session, Schema definition.getOwner(), definition.getProperties()); - updateMaterializedView(session, viewName, newDefinition); + updateMaterializedView(viewName, newDefinition); } - private void updateMaterializedView(ConnectorSession session, SchemaTableName viewName, ConnectorMaterializedViewDefinition newDefinition) + private void updateMaterializedView(SchemaTableName viewName, ConnectorMaterializedViewDefinition newDefinition) { + Table table = getTable(viewName, false); TableInput materializedViewTableInput = getMaterializedViewTableInput( viewName.getTableName(), encodeMaterializedViewData(fromConnectorMaterializedViewDefinition(newDefinition)), - session.getUser(), - createMaterializedViewProperties(session, newDefinition.getStorageTable().orElseThrow().getSchemaTableName())); - + table.owner(), + table.parameters()); try { updateTable(viewName.getSchemaName(), materializedViewTableInput); } - catch (AmazonServiceException e) { + catch (SdkException e) { throw new TrinoException(ICEBERG_CATALOG_ERROR, e); } } @@ -1191,29 +1248,39 @@ private void updateMaterializedView(ConnectorSession session, SchemaTableName vi @Override public void dropMaterializedView(ConnectorSession session, SchemaTableName viewName) { - com.amazonaws.services.glue.model.Table view = getTableAndCacheMetadata(session, viewName) + Table view = getTableAndCacheMetadata(session, viewName) .orElseThrow(() -> new MaterializedViewNotFoundException(viewName)); - if (!isTrinoMaterializedView(getTableType(view), getTableParameters(view))) { - throw new TrinoException(UNSUPPORTED_TABLE_TYPE, "Not a Materialized View: " + view.getDatabaseName() + "." + view.getName()); + if (!isTrinoMaterializedView(getTableType(view), view.parameters())) { + throw new TrinoException(UNSUPPORTED_TABLE_TYPE, "Not a Materialized View: " + view.databaseName() + "." + view.name()); } - materializedViewCache.remove(viewName); - dropStorageTable(session, view); - deleteTable(view.getDatabaseName(), view.getName()); + materializedViewCache.invalidate(viewName); + dropMaterializedViewStorage(session, view); + deleteTable(view.databaseName(), view.name()); } - private void dropStorageTable(ConnectorSession session, com.amazonaws.services.glue.model.Table view) + private void dropMaterializedViewStorage(ConnectorSession session, Table view) { - Map parameters = getTableParameters(view); + Map parameters = view.parameters(); String storageTableName = parameters.get(STORAGE_TABLE); if (storageTableName != null) { String storageSchema = Optional.ofNullable(parameters.get(STORAGE_SCHEMA)) - .orElse(view.getDatabaseName()); + .orElse(view.databaseName()); try { dropTable(session, new SchemaTableName(storageSchema, storageTableName)); } catch (TrinoException e) { - LOG.warn(e, "Failed to drop storage table '%s.%s' for materialized view '%s'", storageSchema, storageTableName, view.getName()); + LOG.warn(e, "Failed to drop storage table '%s.%s' for materialized view '%s'", storageSchema, storageTableName, view.name()); + } + } + else { + String storageMetadataLocation = parameters.get(METADATA_LOCATION_PROP); + checkState(storageMetadataLocation != null, "Storage location missing in definition of materialized view " + view.name()); + try { + dropMaterializedViewStorage(fileSystemFactory.create(session), storageMetadataLocation); + } + catch (IOException e) { + LOG.warn(e, "Failed to delete storage table metadata '%s' for materialized view '%s'", storageMetadataLocation, view.name()); } } } @@ -1221,63 +1288,114 @@ private void dropStorageTable(ConnectorSession session, com.amazonaws.services.g @Override protected Optional doGetMaterializedView(ConnectorSession session, SchemaTableName viewName) { - ConnectorMaterializedViewDefinition materializedViewDefinition = materializedViewCache.get(viewName); - if (materializedViewDefinition != null) { - return Optional.of(materializedViewDefinition); + MaterializedViewData materializedViewData = materializedViewCache.getIfPresent(viewName); + if (materializedViewData != null) { + return Optional.of(materializedViewData.connectorMaterializedViewDefinition); } - if (tableMetadataCache.containsKey(viewName) || viewCache.containsKey(viewName)) { + if (tableMetadataCache.asMap().containsKey(viewName) || viewCache.asMap().containsKey(viewName)) { // Entries in these caches are not materialized views. return Optional.empty(); } - Optional maybeTable = getTableAndCacheMetadata(session, viewName); + Optional
maybeTable = getTableAndCacheMetadata(session, viewName); if (maybeTable.isEmpty()) { return Optional.empty(); } - com.amazonaws.services.glue.model.Table table = maybeTable.get(); - if (!isTrinoMaterializedView(getTableType(table), getTableParameters(table))) { + Table table = maybeTable.get(); + if (!isTrinoMaterializedView(getTableType(table), table.parameters())) { return Optional.empty(); } - return createMaterializedViewDefinition(session, viewName, table); + return Optional.of(createMaterializedViewDefinition(session, viewName, table)); } - private Optional createMaterializedViewDefinition( + private ConnectorMaterializedViewDefinition createMaterializedViewDefinition( ConnectorSession session, SchemaTableName viewName, - com.amazonaws.services.glue.model.Table table) + Table table) { - Map materializedViewParameters = getTableParameters(table); + Map materializedViewParameters = table.parameters(); String storageTable = materializedViewParameters.get(STORAGE_TABLE); - checkState(storageTable != null, "Storage table missing in definition of materialized view " + viewName); - String storageSchema = Optional.ofNullable(materializedViewParameters.get(STORAGE_SCHEMA)) - .orElse(viewName.getSchemaName()); - SchemaTableName storageTableName = new SchemaTableName(storageSchema, storageTable); + String storageMetadataLocation = materializedViewParameters.get(METADATA_LOCATION_PROP); + if ((storageTable == null) == (storageMetadataLocation == null)) { + throw new TrinoException(ICEBERG_BAD_DATA, "Materialized view should have exactly one of the %s properties set: %s".formatted( + ImmutableList.of(STORAGE_TABLE, METADATA_LOCATION_PROP), + materializedViewParameters)); + } - Table icebergTable; - try { - icebergTable = loadTable(session, storageTableName); + SchemaTableName storageTableName; + if (storageTable != null) { + String storageSchema = Optional.ofNullable(materializedViewParameters.get(STORAGE_SCHEMA)) + .orElse(viewName.getSchemaName()); + storageTableName = new SchemaTableName(storageSchema, storageTable); + + if (table.viewOriginalText() == null) { + throw new TrinoException(ICEBERG_BAD_DATA, "Materialized view did not have original text " + viewName); + } } - catch (RuntimeException e) { - // The materialized view could be removed concurrently. This may manifest in a number of ways, e.g. - // - io.trino.spi.connector.TableNotFoundException - // - org.apache.iceberg.exceptions.NotFoundException when accessing manifest file - // - other failures when reading storage table's metadata files - // Retry, as we're catching broadly. - throw new MaterializedViewMayBeBeingRemovedException(e); + else { + storageTableName = new SchemaTableName(viewName.getSchemaName(), tableNameWithType(viewName.getTableName(), MATERIALIZED_VIEW_STORAGE)); } - String viewOriginalText = table.getViewOriginalText(); - if (viewOriginalText == null) { - throw new TrinoException(ICEBERG_BAD_DATA, "Materialized view did not have original text " + viewName); + return getMaterializedViewDefinition( + Optional.ofNullable(table.owner()), + table.viewOriginalText(), + storageTableName); + } + + @Override + public Optional getMaterializedViewStorageTable(ConnectorSession session, SchemaTableName viewName) + { + String storageMetadataLocation; + MaterializedViewData materializedViewData = materializedViewCache.getIfPresent(viewName); + if (materializedViewData == null) { + Optional
maybeTable = getTableAndCacheMetadata(session, viewName); + if (maybeTable.isEmpty()) { + return Optional.empty(); + } + Table materializedView = maybeTable.get(); + verify(isTrinoMaterializedView(getTableType(materializedView), materializedView.parameters()), + "getMaterializedViewStorageTable received a table, not a materialized view"); + + // TODO getTableAndCacheMetadata saved the value in materializedViewCache, so we could just use that, except when conversion fails + storageMetadataLocation = materializedView.parameters().get(METADATA_LOCATION_PROP); + checkState(storageMetadataLocation != null, "Storage location missing in definition of materialized view " + materializedView.name()); + } + else { + storageMetadataLocation = materializedViewData.storageMetadataLocation + .orElseThrow(() -> new IllegalStateException("Storage location not defined for materialized view " + viewName)); + } + + SchemaTableName storageTableName = new SchemaTableName(viewName.getSchemaName(), tableNameWithType(viewName.getTableName(), MATERIALIZED_VIEW_STORAGE)); + IcebergTableOperations operations = tableOperationsProvider.createTableOperations( + this, + session, + storageTableName.getSchemaName(), + storageTableName.getTableName(), + Optional.empty(), + Optional.empty()); + + try { + TableMetadata metadata = getMaterializedViewTableMetadata(session, storageTableName, storageMetadataLocation); + operations.initializeFromMetadata(metadata); + return Optional.of(new BaseTable(operations, quotedTableName(storageTableName), TRINO_METRICS_REPORTER)); + } + catch (NotFoundException e) { + // Removed during reading + return Optional.empty(); } - return Optional.of(getMaterializedViewDefinition( - icebergTable, - Optional.ofNullable(table.getOwner()), - viewOriginalText, - storageTableName)); + } + + private TableMetadata getMaterializedViewTableMetadata(ConnectorSession session, SchemaTableName storageTableName, String storageMetadataLocation) + { + requireNonNull(storageTableName, "storageTableName is null"); + requireNonNull(storageMetadataLocation, "storageMetadataLocation is null"); + return uncheckedCacheGet(tableMetadataCache, storageTableName, () -> { + TrinoFileSystem fileSystem = fileSystemFactory.create(session); + return TableMetadataParser.read(new ForwardingFileIo(fileSystem), storageMetadataLocation); + }); } @Override @@ -1285,14 +1403,14 @@ public void renameMaterializedView(ConnectorSession session, SchemaTableName sou { boolean newTableCreated = false; try { - com.amazonaws.services.glue.model.Table glueTable = getTableAndCacheMetadata(session, source) + Table glueTable = getTableAndCacheMetadata(session, source) .orElseThrow(() -> new TableNotFoundException(source)); - materializedViewCache.remove(source); - Map tableParameters = getTableParameters(glueTable); + materializedViewCache.invalidate(source); + Map tableParameters = glueTable.parameters(); if (!isTrinoMaterializedView(getTableType(glueTable), tableParameters)) { throw new TrinoException(UNSUPPORTED_TABLE_TYPE, "Not a Materialized View: " + source); } - TableInput tableInput = getMaterializedViewTableInput(target.getTableName(), glueTable.getViewOriginalText(), glueTable.getOwner(), tableParameters); + TableInput tableInput = getMaterializedViewTableInput(target.getTableName(), glueTable.viewOriginalText(), glueTable.owner(), tableParameters); createTable(target.getSchemaName(), tableInput); newTableCreated = true; deleteTable(source.getSchemaName(), source.getTableName()); @@ -1329,19 +1447,25 @@ public Optional redirectTable(ConnectorSession session, tableName.getSchemaName(), tableName.getTableName().substring(0, metadataMarkerIndex)); - Optional table = getTableAndCacheMetadata(session, new SchemaTableName(tableNameBase.getSchemaName(), tableNameBase.getTableName())); + Optional
table = getTableAndCacheMetadata(session, new SchemaTableName(tableNameBase.getSchemaName(), tableNameBase.getTableName())); if (table.isEmpty() || VIRTUAL_VIEW.name().equals(getTableTypeNullable(table.get()))) { return Optional.empty(); } - if (!isIcebergTable(getTableParameters(table.get()))) { + if (!isIcebergTable(table.get().parameters())) { // After redirecting, use the original table name, with "$partitions" and similar suffixes return Optional.of(new CatalogSchemaTableName(hiveCatalogName, tableName)); } return Optional.empty(); } - com.amazonaws.services.glue.model.Table getTable(SchemaTableName tableName, boolean invalidateCaches) + @Override + protected void invalidateTableCache(SchemaTableName schemaTableName) + { + tableMetadataCache.invalidate(schemaTableName); + } + + Table getTable(SchemaTableName tableName, boolean invalidateCaches) { if (invalidateCaches) { glueTableCache.invalidate(tableName); @@ -1350,14 +1474,18 @@ com.amazonaws.services.glue.model.Table getTable(SchemaTableName tableName, bool try { return uncheckedCacheGet(glueTableCache, tableName, () -> { try { - GetTableRequest getTableRequest = new GetTableRequest() - .withDatabaseName(tableName.getSchemaName()) - .withName(tableName.getTableName()); - return stats.getGetTable().call(() -> glueClient.getTable(getTableRequest).getTable()); + return stats.getGetTable().call(() -> + glueClient.getTable(x -> x + .databaseName(tableName.getSchemaName()) + .name(tableName.getTableName())) + .table()); } catch (EntityNotFoundException e) { throw new TableNotFoundException(tableName, e); } + catch (SdkException e) { + throw new TrinoException(ICEBERG_CATALOG_ERROR, e); + } }); } catch (UncheckedExecutionException e) { @@ -1366,42 +1494,112 @@ com.amazonaws.services.glue.model.Table getTable(SchemaTableName tableName, bool } } - private Stream getGlueTables(String glueNamespace) + private Stream> listRelations(ConnectorSession session, Optional namespace) + { + List namespaces = listNamespaces(session, namespace); + return namespaces.stream() + .flatMap(glueNamespace -> getGlueTablesWithExceptionHandling(glueNamespace) + .map(table -> { + String tableType = getTableType(table); + Map tableParameters = table.parameters(); + RelationType relationType; + if (isTrinoView(tableType, tableParameters)) { + relationType = RelationType.VIEW; + } + else if (isTrinoMaterializedView(tableType, tableParameters)) { + relationType = RelationType.MATERIALIZED_VIEW; + } + else { + relationType = RelationType.TABLE; + } + + return entry(new SchemaTableName(glueNamespace, table.name()), relationType); + })); + } + + private Stream
getGlueTablesWithExceptionHandling(String glueNamespace) { - return getPaginatedResults( - glueClient::getTables, - new GetTablesRequest().withDatabaseName(glueNamespace), - GetTablesRequest::setNextToken, - GetTablesResult::getNextToken, - stats.getGetTables()) - .map(GetTablesResult::getTableList) - .flatMap(List::stream); + return stream(new AbstractIterator<>() + { + private Iterator
delegate; + + @Override + protected Table computeNext() + { + boolean firstCall = (delegate == null); + try { + if (delegate == null) { + delegate = getGlueTables(glueNamespace) + .iterator(); + } + + if (!delegate.hasNext()) { + return endOfData(); + } + return delegate.next(); + } + catch (EntityNotFoundException e) { + // database does not exist or deleted during iteration + return endOfData(); + } + catch (AccessDeniedException e) { + // permission denied may actually mean "does not exist" + if (!firstCall) { + LOG.warn(e, "Permission denied when getting next batch of tables from namespace %s", glueNamespace); + } + return endOfData(); + } + catch (SdkException e) { + throw new TrinoException(ICEBERG_CATALOG_ERROR, e); + } + } + }); + } + + private Stream
getGlueTables(String glueNamespace) + { + return stats.getGetTables().call(() -> + glueClient.getTablesPaginator(x -> x.databaseName(glueNamespace)) + .stream() + .map(GetTablesResponse::tableList) + .flatMap(List::stream)); } private void createTable(String schemaName, TableInput tableInput) { glueTableCache.invalidateAll(); stats.getCreateTable().call(() -> - glueClient.createTable(new CreateTableRequest() - .withDatabaseName(schemaName) - .withTableInput(tableInput))); + glueClient.createTable(x -> x + .databaseName(schemaName) + .tableInput(tableInput))); } private void updateTable(String schemaName, TableInput tableInput) { glueTableCache.invalidateAll(); stats.getUpdateTable().call(() -> - glueClient.updateTable(new UpdateTableRequest() - .withDatabaseName(schemaName) - .withTableInput(tableInput))); + glueClient.updateTable(x -> x + .databaseName(schemaName) + .tableInput(tableInput))); } private void deleteTable(String schema, String table) { glueTableCache.invalidateAll(); stats.getDeleteTable().call(() -> - glueClient.deleteTable(new DeleteTableRequest() - .withDatabaseName(schema) - .withName(table))); + glueClient.deleteTable(x -> x + .databaseName(schema) + .name(table))); + } + + private record MaterializedViewData( + ConnectorMaterializedViewDefinition connectorMaterializedViewDefinition, + Optional storageMetadataLocation) + { + private MaterializedViewData + { + requireNonNull(connectorMaterializedViewDefinition, "connectorMaterializedViewDefinition is null"); + requireNonNull(storageMetadataLocation, "storageMetadataLocation is null"); + } } } diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/glue/TrinoGlueCatalogFactory.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/glue/TrinoGlueCatalogFactory.java index 054c1bbc1c79..017cbeb44771 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/glue/TrinoGlueCatalogFactory.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/glue/TrinoGlueCatalogFactory.java @@ -13,14 +13,16 @@ */ package io.trino.plugin.iceberg.catalog.glue; -import com.amazonaws.services.glue.AWSGlueAsync; import com.google.inject.Inject; +import io.airlift.concurrent.BoundedExecutor; import io.trino.filesystem.TrinoFileSystemFactory; import io.trino.plugin.base.CatalogName; import io.trino.plugin.hive.NodeVersion; import io.trino.plugin.hive.metastore.glue.GlueHiveMetastoreConfig; import io.trino.plugin.hive.metastore.glue.GlueMetastoreStats; +import io.trino.plugin.iceberg.ForIcebergMetadata; import io.trino.plugin.iceberg.IcebergConfig; +import io.trino.plugin.iceberg.IcebergSecurityConfig; import io.trino.plugin.iceberg.catalog.IcebergTableOperationsProvider; import io.trino.plugin.iceberg.catalog.TrinoCatalog; import io.trino.plugin.iceberg.catalog.TrinoCatalogFactory; @@ -28,9 +30,14 @@ import io.trino.spi.type.TypeManager; import org.weakref.jmx.Flatten; import org.weakref.jmx.Managed; +import software.amazon.awssdk.services.glue.GlueClient; import java.util.Optional; +import java.util.concurrent.Executor; +import java.util.concurrent.ExecutorService; +import static com.google.common.util.concurrent.MoreExecutors.directExecutor; +import static io.trino.plugin.iceberg.IcebergSecurityConfig.IcebergSecurity.SYSTEM; import static java.util.Objects.requireNonNull; public class TrinoGlueCatalogFactory @@ -43,9 +50,12 @@ public class TrinoGlueCatalogFactory private final IcebergTableOperationsProvider tableOperationsProvider; private final String trinoVersion; private final Optional defaultSchemaLocation; - private final AWSGlueAsync glueClient; + private final GlueClient glueClient; private final boolean isUniqueTableLocation; + private final boolean hideMaterializedViewStorageTable; private final GlueMetastoreStats stats; + private final boolean isUsingSystemSecurity; + private final Executor metadataFetchingExecutor; @Inject public TrinoGlueCatalogFactory( @@ -57,8 +67,10 @@ public TrinoGlueCatalogFactory( GlueHiveMetastoreConfig glueConfig, IcebergConfig icebergConfig, IcebergGlueCatalogConfig catalogConfig, + IcebergSecurityConfig securityConfig, GlueMetastoreStats stats, - AWSGlueAsync glueClient) + GlueClient glueClient, + @ForIcebergMetadata ExecutorService metadataExecutorService) { this.catalogName = requireNonNull(catalogName, "catalogName is null"); this.fileSystemFactory = requireNonNull(fileSystemFactory, "fileSystemFactory is null"); @@ -69,7 +81,15 @@ public TrinoGlueCatalogFactory( this.defaultSchemaLocation = glueConfig.getDefaultWarehouseDir(); this.glueClient = requireNonNull(glueClient, "glueClient is null"); this.isUniqueTableLocation = icebergConfig.isUniqueTableLocation(); + this.hideMaterializedViewStorageTable = icebergConfig.isHideMaterializedViewStorageTable(); this.stats = requireNonNull(stats, "stats is null"); + this.isUsingSystemSecurity = securityConfig.getSecuritySystem() == SYSTEM; + if (icebergConfig.getMetadataParallelism() == 1) { + this.metadataFetchingExecutor = directExecutor(); + } + else { + this.metadataFetchingExecutor = new BoundedExecutor(metadataExecutorService, icebergConfig.getMetadataParallelism()); + } } @Managed @@ -91,7 +111,10 @@ public TrinoCatalog create(ConnectorIdentity identity) trinoVersion, glueClient, stats, + isUsingSystemSecurity, defaultSchemaLocation, - isUniqueTableLocation); + isUniqueTableLocation, + hideMaterializedViewStorageTable, + metadataFetchingExecutor); } } diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/hms/AbstractMetastoreTableOperations.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/hms/AbstractMetastoreTableOperations.java index 0774d7103c2e..c531a5c4b343 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/hms/AbstractMetastoreTableOperations.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/hms/AbstractMetastoreTableOperations.java @@ -14,7 +14,6 @@ package io.trino.plugin.iceberg.catalog.hms; import io.trino.annotation.NotThreadSafe; -import io.trino.plugin.hive.TableAlreadyExistsException; import io.trino.plugin.hive.metastore.MetastoreUtil; import io.trino.plugin.hive.metastore.PrincipalPrivileges; import io.trino.plugin.hive.metastore.Table; @@ -23,7 +22,6 @@ import io.trino.plugin.iceberg.catalog.AbstractIcebergTableOperations; import io.trino.spi.TrinoException; import io.trino.spi.connector.ConnectorSession; -import io.trino.spi.connector.SchemaNotFoundException; import io.trino.spi.connector.TableNotFoundException; import org.apache.iceberg.TableMetadata; import org.apache.iceberg.io.FileIO; @@ -31,13 +29,15 @@ import java.util.Optional; import static com.google.common.base.Verify.verify; -import static io.trino.plugin.hive.HiveMetadata.TABLE_COMMENT; import static io.trino.plugin.hive.TableType.EXTERNAL_TABLE; -import static io.trino.plugin.hive.ViewReaderUtil.isHiveOrPrestoView; -import static io.trino.plugin.hive.ViewReaderUtil.isPrestoView; +import static io.trino.plugin.hive.ViewReaderUtil.isTrinoMaterializedView; +import static io.trino.plugin.hive.ViewReaderUtil.isTrinoView; import static io.trino.plugin.hive.metastore.PrincipalPrivileges.NO_PRIVILEGES; +import static io.trino.plugin.hive.metastore.Table.TABLE_COMMENT; import static io.trino.plugin.hive.util.HiveUtil.isIcebergTable; import static io.trino.plugin.iceberg.IcebergErrorCode.ICEBERG_INVALID_METADATA; +import static io.trino.plugin.iceberg.IcebergTableName.isMaterializedViewStorage; +import static io.trino.plugin.iceberg.IcebergTableName.tableNameFrom; import static java.lang.String.format; import static java.util.Locale.ENGLISH; import static java.util.Objects.requireNonNull; @@ -71,13 +71,22 @@ protected final String getRefreshedLocation(boolean invalidateCaches) if (invalidateCaches) { metastore.invalidateTable(database, tableName); } - Table table = getTable(); - if (isPrestoView(table) && isHiveOrPrestoView(table)) { + boolean isMaterializedViewStorageTable = isMaterializedViewStorage(tableName); + + Table table; + if (isMaterializedViewStorageTable) { + table = getTable(database, tableNameFrom(tableName)); + } + else { + table = getTable(); + } + + if (!isMaterializedViewStorageTable && (isTrinoView(table) || isTrinoMaterializedView(table))) { // this is a Hive view, hence not a table throw new TableNotFoundException(getSchemaTableName()); } - if (!isIcebergTable(table)) { + if (!isMaterializedViewStorageTable && !isIcebergTable(table)) { throw new UnknownTableTypeException(getSchemaTableName()); } @@ -112,8 +121,7 @@ protected final void commitNewTable(TableMetadata metadata) try { metastore.createTable(table, privileges); } - catch (SchemaNotFoundException - | TableAlreadyExistsException e) { + catch (Exception e) { // clean up metadata files corresponding to the current transaction fileIo.deleteFile(newMetadataLocation); throw e; @@ -131,6 +139,11 @@ protected Table.Builder updateMetastoreTable(Table.Builder builder, TableMetadat } protected Table getTable() + { + return getTable(database, tableName); + } + + protected Table getTable(String database, String tableName) { return metastore.getTable(database, tableName) .orElseThrow(() -> new TableNotFoundException(getSchemaTableName())); diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/hms/HiveMetastoreTableOperations.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/hms/HiveMetastoreTableOperations.java index 5a5d4dead2b1..9894494c09c8 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/hms/HiveMetastoreTableOperations.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/hms/HiveMetastoreTableOperations.java @@ -13,6 +13,7 @@ */ package io.trino.plugin.iceberg.catalog.hms; +import com.google.common.collect.ImmutableMap; import io.airlift.log.Logger; import io.trino.annotation.NotThreadSafe; import io.trino.plugin.hive.metastore.AcidTransactionOwner; @@ -28,14 +29,20 @@ import org.apache.iceberg.exceptions.CommitStateUnknownException; import org.apache.iceberg.io.FileIO; +import java.util.Map; import java.util.Optional; +import java.util.function.BiFunction; import static com.google.common.base.Preconditions.checkState; import static io.trino.plugin.hive.metastore.PrincipalPrivileges.NO_PRIVILEGES; import static io.trino.plugin.hive.metastore.thrift.ThriftMetastoreUtil.fromMetastoreApiTable; +import static io.trino.plugin.iceberg.IcebergTableName.tableNameFrom; import static io.trino.plugin.iceberg.IcebergUtil.fixBrokenMetadataLocation; +import static java.lang.Boolean.parseBoolean; import static java.util.Objects.requireNonNull; import static org.apache.iceberg.BaseMetastoreTableOperations.METADATA_LOCATION_PROP; +import static org.apache.iceberg.BaseMetastoreTableOperations.PREVIOUS_METADATA_LOCATION_PROP; +import static org.apache.iceberg.TableProperties.HIVE_LOCK_ENABLED; @NotThreadSafe public class HiveMetastoreTableOperations @@ -43,11 +50,13 @@ public class HiveMetastoreTableOperations { private static final Logger log = Logger.get(HiveMetastoreTableOperations.class); private final ThriftMetastore thriftMetastore; + private final boolean lockingEnabled; public HiveMetastoreTableOperations( FileIO fileIo, CachingHiveMetastore metastore, ThriftMetastore thriftMetastore, + boolean lockingEnabled, ConnectorSession session, String database, String table, @@ -56,20 +65,39 @@ public HiveMetastoreTableOperations( { super(fileIo, metastore, session, database, table, owner, location); this.thriftMetastore = requireNonNull(thriftMetastore, "thriftMetastore is null"); + this.lockingEnabled = lockingEnabled; } @Override protected void commitToExistingTable(TableMetadata base, TableMetadata metadata) + { + Table currentTable = getTable(); + commitTableUpdate(currentTable, metadata, (table, newMetadataLocation) -> Table.builder(table) + .apply(builder -> updateMetastoreTable(builder, metadata, newMetadataLocation, Optional.of(currentMetadataLocation))) + .build()); + } + + @Override + protected final void commitMaterializedViewRefresh(TableMetadata base, TableMetadata metadata) + { + Table materializedView = getTable(database, tableNameFrom(tableName)); + commitTableUpdate(materializedView, metadata, (table, newMetadataLocation) -> Table.builder(table) + .apply(builder -> builder + .setParameter(METADATA_LOCATION_PROP, newMetadataLocation) + .setParameter(PREVIOUS_METADATA_LOCATION_PROP, currentMetadataLocation)) + .build()); + } + + private void commitTableUpdate(Table table, TableMetadata metadata, BiFunction tableUpdateFunction) { String newMetadataLocation = writeNewMetadata(metadata, version.orElseThrow() + 1); - long lockId = thriftMetastore.acquireTableExclusiveLock( - new AcidTransactionOwner(session.getUser()), - session.getQueryId(), - database, - tableName); + boolean lockingEnabled = parseBoolean(table.getParameters().getOrDefault(HIVE_LOCK_ENABLED, Boolean.toString(this.lockingEnabled))); + HiveLock hiveLock = lockingEnabled ? new ThriftMetastoreLock(table) : new NoLock(); + hiveLock.acquire(); + try { - Table currentTable = fromMetastoreApiTable(thriftMetastore.getTable(database, tableName) + Table currentTable = fromMetastoreApiTable(thriftMetastore.getTable(database, table.getTableName()) .orElseThrow(() -> new TableNotFoundException(getSchemaTableName()))); checkState(currentMetadataLocation != null, "No current metadata location for existing table"); @@ -79,14 +107,12 @@ protected void commitToExistingTable(TableMetadata base, TableMetadata metadata) currentMetadataLocation, metadataLocation, getSchemaTableName()); } - Table table = Table.builder(currentTable) - .apply(builder -> updateMetastoreTable(builder, metadata, newMetadataLocation, Optional.of(currentMetadataLocation))) - .build(); + Table updatedTable = tableUpdateFunction.apply(table, newMetadataLocation); // todo privileges should not be replaced for an alter PrincipalPrivileges privileges = table.getOwner().map(MetastoreUtil::buildInitialPrivilegeSet).orElse(NO_PRIVILEGES); try { - metastore.replaceTable(database, tableName, table, privileges); + metastore.replaceTable(table.getDatabaseName(), table.getTableName(), updatedTable, privileges, environmentContext(metadataLocation)); } catch (RuntimeException e) { // Cannot determine whether the `replaceTable` operation was successful, @@ -95,6 +121,48 @@ protected void commitToExistingTable(TableMetadata base, TableMetadata metadata) } } finally { + hiveLock.release(); + } + + shouldRefresh = true; + } + + private static Map environmentContext(String metadataLocation) + { + if (metadataLocation == null) { + return ImmutableMap.of(); + } + return ImmutableMap.builder() + .put("expected_parameter_key", "metadata_location") + .put("expected_parameter_value", metadataLocation) + .buildOrThrow(); + } + + private class ThriftMetastoreLock + implements HiveLock + { + private long lockId; + + private final Table table; + + public ThriftMetastoreLock(Table table) + { + this.table = requireNonNull(table, "table is null"); + } + + @Override + public void acquire() + { + lockId = thriftMetastore.acquireTableExclusiveLock( + new AcidTransactionOwner(session.getUser()), + session.getQueryId(), + table.getDatabaseName(), + table.getTableName()); + } + + @Override + public void release() + { try { thriftMetastore.releaseTableLock(lockId); } @@ -103,10 +171,26 @@ protected void commitToExistingTable(TableMetadata base, TableMetadata metadata) // So, that underlying iceberg API will not do the metadata cleanup, otherwise table will be in unusable state. // If configured and supported, the unreleased lock will be automatically released by the metastore after not hearing a heartbeat for a while, // or otherwise it might need to be manually deleted from the metastore backend storage. - log.error(e, "Failed to release lock %s when committing to table %s", lockId, tableName); + log.error(e, "Failed to release lock %s when committing to table %s", lockId, table.getTableName()); } } + } - shouldRefresh = true; + // HIVE-26882 requires HMS client 2 or later. Our HMS client is based on version 3. + private static class NoLock + implements HiveLock + { + @Override + public void acquire() {} + + @Override + public void release() {} + } + + private interface HiveLock + { + void acquire(); + + void release(); } } diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/hms/HiveMetastoreTableOperationsProvider.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/hms/HiveMetastoreTableOperationsProvider.java index b7a31bdbe817..ea6ef865ae9e 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/hms/HiveMetastoreTableOperationsProvider.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/hms/HiveMetastoreTableOperationsProvider.java @@ -31,12 +31,17 @@ public class HiveMetastoreTableOperationsProvider { private final TrinoFileSystemFactory fileSystemFactory; private final ThriftMetastoreFactory thriftMetastoreFactory; + private final boolean lockingEnabled; @Inject - public HiveMetastoreTableOperationsProvider(TrinoFileSystemFactory fileSystemFactory, ThriftMetastoreFactory thriftMetastoreFactory) + public HiveMetastoreTableOperationsProvider( + TrinoFileSystemFactory fileSystemFactory, + ThriftMetastoreFactory thriftMetastoreFactory, + IcebergHiveCatalogConfig metastoreConfig) { this.fileSystemFactory = requireNonNull(fileSystemFactory, "fileSystemFactory is null"); this.thriftMetastoreFactory = requireNonNull(thriftMetastoreFactory, "thriftMetastoreFactory is null"); + this.lockingEnabled = metastoreConfig.getLockingEnabled(); } @Override @@ -52,6 +57,7 @@ public IcebergTableOperations createTableOperations( new ForwardingFileIo(fileSystemFactory.create(session)), ((TrinoHiveCatalog) catalog).getMetastore(), thriftMetastoreFactory.createMetastore(Optional.of(session.getIdentity())), + lockingEnabled, session, database, table, diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/aws/athena/projection/InjectedProjectionFactory.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/hms/IcebergHiveCatalogConfig.java similarity index 50% rename from plugin/trino-hive/src/main/java/io/trino/plugin/hive/aws/athena/projection/InjectedProjectionFactory.java rename to plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/hms/IcebergHiveCatalogConfig.java index b90d0acc6ddb..90e315617631 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/aws/athena/projection/InjectedProjectionFactory.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/hms/IcebergHiveCatalogConfig.java @@ -11,26 +11,25 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package io.trino.plugin.hive.aws.athena.projection; +package io.trino.plugin.iceberg.catalog.hms; -import io.trino.spi.type.Type; +import io.airlift.configuration.Config; +import io.airlift.configuration.ConfigDescription; -import java.util.Map; - -import static io.trino.plugin.hive.metastore.MetastoreUtil.canConvertSqlTypeToStringForParts; - -public class InjectedProjectionFactory - implements ProjectionFactory +public class IcebergHiveCatalogConfig { - @Override - public boolean isSupportedColumnType(Type columnType) + private boolean lockingEnabled = true; + + public boolean getLockingEnabled() { - return canConvertSqlTypeToStringForParts(columnType, true); + return lockingEnabled; } - @Override - public Projection create(String columnName, Type columnType, Map columnProperties) + @Config("iceberg.hive-catalog.locking-enabled") + @ConfigDescription("Acquire locks when updating tables") + public IcebergHiveCatalogConfig setLockingEnabled(boolean lockingEnabled) { - return new InjectedProjection(columnName); + this.lockingEnabled = lockingEnabled; + return this; } } diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/hms/TrinoHiveCatalog.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/hms/TrinoHiveCatalog.java index 65c89fe3f517..40b2cac40ed5 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/hms/TrinoHiveCatalog.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/hms/TrinoHiveCatalog.java @@ -13,11 +13,15 @@ */ package io.trino.plugin.iceberg.catalog.hms; +import com.google.common.cache.Cache; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; +import com.google.common.util.concurrent.UncheckedExecutionException; import io.airlift.log.Logger; +import io.trino.cache.EvictableCacheBuilder; import io.trino.filesystem.Location; +import io.trino.filesystem.TrinoFileSystem; import io.trino.filesystem.TrinoFileSystemFactory; import io.trino.plugin.base.CatalogName; import io.trino.plugin.hive.HiveSchemaProperties; @@ -27,11 +31,15 @@ import io.trino.plugin.hive.metastore.HivePrincipal; import io.trino.plugin.hive.metastore.MetastoreUtil; import io.trino.plugin.hive.metastore.PrincipalPrivileges; +import io.trino.plugin.hive.metastore.TableInfo; import io.trino.plugin.hive.metastore.cache.CachingHiveMetastore; import io.trino.plugin.hive.util.HiveUtil; +import io.trino.plugin.iceberg.IcebergTableName; import io.trino.plugin.iceberg.UnknownTableTypeException; import io.trino.plugin.iceberg.catalog.AbstractTrinoCatalog; +import io.trino.plugin.iceberg.catalog.IcebergTableOperations; import io.trino.plugin.iceberg.catalog.IcebergTableOperationsProvider; +import io.trino.plugin.iceberg.fileio.ForwardingFileIo; import io.trino.spi.TrinoException; import io.trino.spi.connector.CatalogSchemaTableName; import io.trino.spi.connector.ColumnMetadata; @@ -53,46 +61,58 @@ import org.apache.iceberg.SortOrder; import org.apache.iceberg.Table; import org.apache.iceberg.TableMetadata; +import org.apache.iceberg.TableMetadataParser; import org.apache.iceberg.Transaction; +import org.apache.iceberg.exceptions.NotFoundException; import java.io.IOException; +import java.util.Collection; import java.util.Iterator; import java.util.List; import java.util.Map; import java.util.Objects; import java.util.Optional; import java.util.Set; -import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.Callable; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.Executor; import java.util.function.Predicate; import java.util.function.UnaryOperator; import java.util.stream.Collectors; import static com.google.common.base.Preconditions.checkState; +import static com.google.common.base.Throwables.throwIfUnchecked; +import static com.google.common.base.Verify.verify; import static com.google.common.collect.ImmutableList.toImmutableList; +import static io.trino.cache.CacheUtils.uncheckedCacheGet; import static io.trino.filesystem.Locations.appendPath; +import static io.trino.plugin.base.util.ExecutorUtil.processWithAdditionalThreads; import static io.trino.plugin.hive.HiveErrorCode.HIVE_DATABASE_LOCATION_ERROR; import static io.trino.plugin.hive.HiveErrorCode.HIVE_INVALID_METADATA; import static io.trino.plugin.hive.HiveMetadata.STORAGE_TABLE; -import static io.trino.plugin.hive.HiveMetadata.TABLE_COMMENT; import static io.trino.plugin.hive.HiveType.HIVE_STRING; import static io.trino.plugin.hive.TableType.EXTERNAL_TABLE; import static io.trino.plugin.hive.TableType.VIRTUAL_VIEW; -import static io.trino.plugin.hive.ViewReaderUtil.ICEBERG_MATERIALIZED_VIEW_COMMENT; -import static io.trino.plugin.hive.ViewReaderUtil.encodeViewData; import static io.trino.plugin.hive.ViewReaderUtil.isHiveOrPrestoView; import static io.trino.plugin.hive.ViewReaderUtil.isTrinoMaterializedView; import static io.trino.plugin.hive.metastore.MetastoreUtil.buildInitialPrivilegeSet; import static io.trino.plugin.hive.metastore.PrincipalPrivileges.NO_PRIVILEGES; import static io.trino.plugin.hive.metastore.StorageFormat.VIEW_STORAGE_FORMAT; +import static io.trino.plugin.hive.metastore.TableInfo.ICEBERG_MATERIALIZED_VIEW_COMMENT; import static io.trino.plugin.hive.util.HiveUtil.isHiveSystemSchema; import static io.trino.plugin.hive.util.HiveUtil.isIcebergTable; -import static io.trino.plugin.iceberg.IcebergMaterializedViewAdditionalProperties.STORAGE_SCHEMA; +import static io.trino.plugin.iceberg.IcebergErrorCode.ICEBERG_BAD_DATA; import static io.trino.plugin.iceberg.IcebergMaterializedViewDefinition.encodeMaterializedViewData; import static io.trino.plugin.iceberg.IcebergMaterializedViewDefinition.fromConnectorMaterializedViewDefinition; +import static io.trino.plugin.iceberg.IcebergMaterializedViewProperties.STORAGE_SCHEMA; import static io.trino.plugin.iceberg.IcebergSchemaProperties.LOCATION_PROPERTY; +import static io.trino.plugin.iceberg.IcebergUtil.TRINO_QUERY_ID_NAME; import static io.trino.plugin.iceberg.IcebergUtil.getIcebergTableWithMetadata; import static io.trino.plugin.iceberg.IcebergUtil.loadIcebergTable; +import static io.trino.plugin.iceberg.IcebergUtil.quotedTableName; import static io.trino.plugin.iceberg.IcebergUtil.validateTableCanBeDropped; +import static io.trino.plugin.iceberg.TableType.MATERIALIZED_VIEW_STORAGE; +import static io.trino.plugin.iceberg.TrinoMetricsReporter.TRINO_METRICS_REPORTER; import static io.trino.plugin.iceberg.catalog.AbstractIcebergTableOperations.ICEBERG_METASTORE_STORAGE_FORMAT; import static io.trino.plugin.iceberg.catalog.AbstractIcebergTableOperations.toHiveColumns; import static io.trino.spi.StandardErrorCode.ALREADY_EXISTS; @@ -113,6 +133,7 @@ public class TrinoHiveCatalog extends AbstractTrinoCatalog { private static final Logger log = Logger.get(TrinoHiveCatalog.class); + private static final int PER_QUERY_CACHE_SIZE = 1000; public static final String DEPENDS_ON_TABLES = "dependsOnTables"; // Value should be ISO-8601 formatted time instant public static final String TRINO_QUERY_START_TIME = "trino-query-start-time"; @@ -122,8 +143,12 @@ public class TrinoHiveCatalog private final TrinoFileSystemFactory fileSystemFactory; private final boolean isUsingSystemSecurity; private final boolean deleteSchemaLocationsFallback; + private final boolean hideMaterializedViewStorageTable; + private final Executor metadataFetchingExecutor; - private final Map tableMetadataCache = new ConcurrentHashMap<>(); + private final Cache tableMetadataCache = EvictableCacheBuilder.newBuilder() + .maximumSize(PER_QUERY_CACHE_SIZE) + .build(); public TrinoHiveCatalog( CatalogName catalogName, @@ -134,14 +159,18 @@ public TrinoHiveCatalog( IcebergTableOperationsProvider tableOperationsProvider, boolean useUniqueTableLocation, boolean isUsingSystemSecurity, - boolean deleteSchemaLocationsFallback) + boolean deleteSchemaLocationsFallback, + boolean hideMaterializedViewStorageTable, + Executor metadataFetchingExecutor) { - super(catalogName, typeManager, tableOperationsProvider, useUniqueTableLocation); + super(catalogName, typeManager, tableOperationsProvider, fileSystemFactory, useUniqueTableLocation); this.metastore = requireNonNull(metastore, "metastore is null"); this.trinoViewHiveMetastore = requireNonNull(trinoViewHiveMetastore, "trinoViewHiveMetastore is null"); this.fileSystemFactory = requireNonNull(fileSystemFactory, "fileSystemFactory is null"); this.isUsingSystemSecurity = isUsingSystemSecurity; this.deleteSchemaLocationsFallback = deleteSchemaLocationsFallback; + this.hideMaterializedViewStorageTable = hideMaterializedViewStorageTable; + this.metadataFetchingExecutor = requireNonNull(metadataFetchingExecutor, "metadataFetchingExecutor is null"); } public CachingHiveMetastore getMetastore() @@ -266,7 +295,7 @@ public Transaction newCreateTableTransaction( Schema schema, PartitionSpec partitionSpec, SortOrder sortOrder, - String location, + Optional location, Map properties) { return newCreateTableTransaction( @@ -280,6 +309,26 @@ public Transaction newCreateTableTransaction( isUsingSystemSecurity ? Optional.empty() : Optional.of(session.getUser())); } + @Override + public Transaction newCreateOrReplaceTableTransaction( + ConnectorSession session, + SchemaTableName schemaTableName, + Schema schema, PartitionSpec partitionSpec, + SortOrder sortOrder, + String location, + Map properties) + { + return newCreateOrReplaceTableTransaction( + session, + schemaTableName, + schema, + partitionSpec, + sortOrder, + location, + properties, + isUsingSystemSecurity ? Optional.empty() : Optional.of(session.getUser())); + } + @Override public void registerTable(ConnectorSession session, SchemaTableName schemaTableName, TableMetadata tableMetadata) throws TrinoException @@ -297,6 +346,7 @@ public void registerTable(ConnectorSession session, SchemaTableName schemaTableN .withStorage(storage -> storage.setStorageFormat(ICEBERG_METASTORE_STORAGE_FORMAT)) // This is a must-have property for the EXTERNAL_TABLE table type .setParameter("EXTERNAL", "TRUE") + .setParameter(TRINO_QUERY_ID_NAME, session.getQueryId()) .setParameter(TABLE_TYPE_PROP, ICEBERG_TABLE_TYPE_VALUE.toUpperCase(ENGLISH)) .setParameter(METADATA_LOCATION_PROP, tableMetadata.metadataFileLocation()); @@ -308,16 +358,45 @@ public void registerTable(ConnectorSession session, SchemaTableName schemaTableN public void unregisterTable(ConnectorSession session, SchemaTableName schemaTableName) { dropTableFromMetastore(schemaTableName); + invalidateTableCache(schemaTableName); } @Override - public List listTables(ConnectorSession session, Optional namespace) + public List listTables(ConnectorSession session, Optional namespace) { - ImmutableSet.Builder tablesListBuilder = ImmutableSet.builder(); - for (String schemaName : listNamespaces(session, namespace)) { - metastore.getAllTables(schemaName).forEach(tableName -> tablesListBuilder.add(new SchemaTableName(schemaName, tableName))); + List>> tasks = listNamespaces(session, namespace).stream() + .map(schema -> (Callable>) () -> metastore.getTables(schema)) + .collect(toImmutableList()); + try { + return processWithAdditionalThreads(tasks, metadataFetchingExecutor).stream() + .flatMap(Collection::stream) + .collect(toImmutableList()); + } + catch (ExecutionException e) { + throw new RuntimeException(e.getCause()); + } + } + + @Override + public List listIcebergTables(ConnectorSession session, Optional namespace) + { + List>> tasks = listNamespaces(session, namespace).stream() + .map(schema -> (Callable>) () -> metastore.getTableNamesWithParameters(schema, TABLE_TYPE_PROP, ImmutableSet.of( + // Get tables with parameter table_type set to "ICEBERG" or "iceberg". This is required because + // Trino uses lowercase value whereas Spark and Flink use uppercase. + ICEBERG_TABLE_TYPE_VALUE.toLowerCase(ENGLISH), + ICEBERG_TABLE_TYPE_VALUE.toUpperCase(ENGLISH))).stream() + .map(tableName -> new SchemaTableName(schema, tableName)) + .collect(toImmutableList())) + .collect(toImmutableList()); + try { + return processWithAdditionalThreads(tasks, metadataFetchingExecutor).stream() + .flatMap(Collection::stream) + .collect(toImmutableList()); + } + catch (ExecutionException e) { + throw new RuntimeException(e.getCause()); } - return tablesListBuilder.build().asList(); } @Override @@ -364,6 +443,7 @@ public void dropTable(ConnectorSession session, SchemaTableName schemaTableName) log.warn(e, "Failed to delete table data referenced by metadata"); } deleteTableDirectory(fileSystemFactory.create(session), schemaTableName, metastoreTable.getStorage().getLocation()); + invalidateTableCache(schemaTableName); } @Override @@ -371,6 +451,7 @@ public void dropCorruptedTable(ConnectorSession session, SchemaTableName schemaT { io.trino.plugin.hive.metastore.Table table = dropTableFromMetastore(schemaTableName); deleteTableDirectory(fileSystemFactory.create(session), schemaTableName, table.getStorage().getLocation()); + invalidateTableCache(schemaTableName); } private io.trino.plugin.hive.metastore.Table dropTableFromMetastore(SchemaTableName schemaTableName) @@ -392,14 +473,23 @@ private io.trino.plugin.hive.metastore.Table dropTableFromMetastore(SchemaTableN public void renameTable(ConnectorSession session, SchemaTableName from, SchemaTableName to) { metastore.renameTable(from.getSchemaName(), from.getTableName(), to.getSchemaName(), to.getTableName()); + invalidateTableCache(from); } @Override - public Table loadTable(ConnectorSession session, SchemaTableName schemaTableName) + public BaseTable loadTable(ConnectorSession session, SchemaTableName schemaTableName) { - TableMetadata metadata = tableMetadataCache.computeIfAbsent( - schemaTableName, - ignore -> ((BaseTable) loadIcebergTable(this, tableOperationsProvider, session, schemaTableName)).operations().current()); + TableMetadata metadata; + try { + metadata = uncheckedCacheGet( + tableMetadataCache, + schemaTableName, + () -> loadIcebergTable(this, tableOperationsProvider, session, schemaTableName).operations().current()); + } + catch (UncheckedExecutionException e) { + throwIfUnchecked(e.getCause()); + throw e; + } return getIcebergTableWithMetadata(this, tableOperationsProvider, session, schemaTableName, metadata); } @@ -422,16 +512,6 @@ public void updateViewColumnComment(ConnectorSession session, SchemaTableName vi trinoViewHiveMetastore.updateViewColumnComment(session, viewName, columnName, comment); } - private void replaceView(ConnectorSession session, SchemaTableName viewName, io.trino.plugin.hive.metastore.Table view, ConnectorViewDefinition newDefinition) - { - io.trino.plugin.hive.metastore.Table.Builder viewBuilder = io.trino.plugin.hive.metastore.Table.builder(view) - .setViewOriginalText(Optional.of(encodeViewData(newDefinition))); - - PrincipalPrivileges principalPrivileges = isUsingSystemSecurity ? NO_PRIVILEGES : buildInitialPrivilegeSet(session.getUser()); - - metastore.replaceTable(viewName.getSchemaName(), viewName.getTableName(), viewBuilder.build(), principalPrivileges); - } - @Override public String defaultTableLocation(ConnectorSession session, SchemaTableName schemaTableName) { @@ -487,21 +567,12 @@ public Optional getView(ConnectorSession session, Schem return trinoViewHiveMetastore.getView(viewName); } - @Override - public List listMaterializedViews(ConnectorSession session, Optional namespace) - { - // Filter on ICEBERG_MATERIALIZED_VIEW_COMMENT is used to avoid listing hive views in case of a shared HMS and to distinguish from standard views - return listNamespaces(session, namespace).stream() - .flatMap(schema -> metastore.getTablesWithParameter(schema, TABLE_COMMENT, ICEBERG_MATERIALIZED_VIEW_COMMENT).stream() - .map(table -> new SchemaTableName(schema, table))) - .collect(toImmutableList()); - } - @Override public void createMaterializedView( ConnectorSession session, SchemaTableName viewName, ConnectorMaterializedViewDefinition definition, + Map materializedViewProperties, boolean replace, boolean ignoreExisting) { @@ -519,11 +590,51 @@ public void createMaterializedView( } } - SchemaTableName storageTable = createMaterializedViewStorageTable(session, viewName, definition); + if (hideMaterializedViewStorageTable) { + Location storageMetadataLocation = createMaterializedViewStorage(session, viewName, definition, materializedViewProperties); + + Map viewProperties = createMaterializedViewProperties(session, storageMetadataLocation); + Column dummyColumn = new Column("dummy", HIVE_STRING, Optional.empty(), ImmutableMap.of()); + io.trino.plugin.hive.metastore.Table.Builder tableBuilder = io.trino.plugin.hive.metastore.Table.builder() + .setDatabaseName(viewName.getSchemaName()) + .setTableName(viewName.getTableName()) + .setOwner(isUsingSystemSecurity ? Optional.empty() : Optional.of(session.getUser())) + .setTableType(VIRTUAL_VIEW.name()) + .setDataColumns(ImmutableList.of(dummyColumn)) + .setPartitionColumns(ImmutableList.of()) + .setParameters(viewProperties) + .withStorage(storage -> storage.setStorageFormat(VIEW_STORAGE_FORMAT)) + .withStorage(storage -> storage.setLocation("")) + .setViewOriginalText(Optional.of( + encodeMaterializedViewData(fromConnectorMaterializedViewDefinition(definition)))) + .setViewExpandedText(Optional.of("/* " + ICEBERG_MATERIALIZED_VIEW_COMMENT + " */")); + io.trino.plugin.hive.metastore.Table table = tableBuilder.build(); + PrincipalPrivileges principalPrivileges = isUsingSystemSecurity ? NO_PRIVILEGES : buildInitialPrivilegeSet(session.getUser()); + + if (existing.isPresent()) { + metastore.replaceTable(viewName.getSchemaName(), viewName.getTableName(), table, principalPrivileges, ImmutableMap.of()); + } + else { + metastore.createTable(table, principalPrivileges); + } + } + else { + createMaterializedViewWithStorageTable(session, viewName, definition, materializedViewProperties, existing); + } + } + + private void createMaterializedViewWithStorageTable( + ConnectorSession session, + SchemaTableName viewName, + ConnectorMaterializedViewDefinition definition, + Map materializedViewProperties, + Optional existing) + { + SchemaTableName storageTable = createMaterializedViewStorageTable(session, viewName, definition, materializedViewProperties); // Create a view indicating the storage table Map viewProperties = createMaterializedViewProperties(session, storageTable); - Column dummyColumn = new Column("dummy", HIVE_STRING, Optional.empty()); + Column dummyColumn = new Column("dummy", HIVE_STRING, Optional.empty(), Map.of()); io.trino.plugin.hive.metastore.Table.Builder tableBuilder = io.trino.plugin.hive.metastore.Table.builder() .setDatabaseName(viewName.getSchemaName()) @@ -549,7 +660,7 @@ public void createMaterializedView( metastore.dropTable(storageSchema, oldStorageTable, true); } // Replace the existing view definition - metastore.replaceTable(viewName.getSchemaName(), viewName.getTableName(), table, principalPrivileges); + metastore.replaceTable(viewName.getSchemaName(), viewName.getTableName(), table, principalPrivileges, ImmutableMap.of()); return; } // create the view definition @@ -594,7 +705,7 @@ private void replaceMaterializedView(ConnectorSession session, SchemaTableName v PrincipalPrivileges principalPrivileges = isUsingSystemSecurity ? NO_PRIVILEGES : buildInitialPrivilegeSet(session.getUser()); - metastore.replaceTable(viewName.getSchemaName(), viewName.getTableName(), viewBuilder.build(), principalPrivileges); + metastore.replaceTable(viewName.getSchemaName(), viewName.getTableName(), viewBuilder.build(), principalPrivileges, ImmutableMap.of()); } @Override @@ -618,6 +729,20 @@ public void dropMaterializedView(ConnectorSession session, SchemaTableName viewN log.warn(e, "Failed to drop storage table '%s.%s' for materialized view '%s'", storageSchema, storageTableName, viewName); } } + + String storageMetadataLocation = view.getParameters().get(METADATA_LOCATION_PROP); + checkState(storageMetadataLocation != null, "Storage location missing in definition of materialized view " + viewName); + + TrinoFileSystem fileSystem = fileSystemFactory.create(session); + TableMetadata metadata = TableMetadataParser.read(new ForwardingFileIo(fileSystem), storageMetadataLocation); + String storageLocation = metadata.location(); + try { + fileSystem.deleteDirectory(Location.of(storageLocation)); + } + catch (IOException e) { + log.warn(e, "Failed to delete storage location '%s' for materialized view '%s'", storageLocation, viewName); + } + metastore.dropTable(viewName.getSchemaName(), viewName.getTableName(), true); } @@ -629,38 +754,114 @@ protected Optional doGetMaterializedView(Co return Optional.empty(); } - io.trino.plugin.hive.metastore.Table table = tableOptional.get(); - if (!isTrinoMaterializedView(table.getTableType(), table.getParameters())) { + io.trino.plugin.hive.metastore.Table materializedView = tableOptional.get(); + if (!isTrinoMaterializedView(materializedView.getTableType(), materializedView.getParameters())) { return Optional.empty(); } - io.trino.plugin.hive.metastore.Table materializedView = tableOptional.get(); String storageTable = materializedView.getParameters().get(STORAGE_TABLE); - checkState(storageTable != null, "Storage table missing in definition of materialized view " + viewName); - String storageSchema = Optional.ofNullable(materializedView.getParameters().get(STORAGE_SCHEMA)) - .orElse(viewName.getSchemaName()); - SchemaTableName storageTableName = new SchemaTableName(storageSchema, storageTable); + String storageMetadataLocation = materializedView.getParameters().get(METADATA_LOCATION_PROP); + if ((storageTable == null) == (storageMetadataLocation == null)) { + throw new TrinoException(ICEBERG_BAD_DATA, "Materialized view should have exactly one of the %s properties set: %s".formatted( + ImmutableList.of(STORAGE_TABLE, METADATA_LOCATION_PROP), + materializedView.getParameters())); + } + + if (storageTable != null) { + String storageSchema = Optional.ofNullable(materializedView.getParameters().get(STORAGE_SCHEMA)) + .orElse(viewName.getSchemaName()); + SchemaTableName storageTableName = new SchemaTableName(storageSchema, storageTable); + + Table icebergTable; + try { + icebergTable = loadTable(session, storageTableName); + } + catch (RuntimeException e) { + // The materialized view could be removed concurrently. This may manifest in a number of ways, e.g. + // - io.trino.spi.connector.TableNotFoundException + // - org.apache.iceberg.exceptions.NotFoundException when accessing manifest file + // - other failures when reading storage table's metadata files + // Retry, as we're catching broadly. + metastore.invalidateTable(viewName.getSchemaName(), viewName.getTableName()); + metastore.invalidateTable(storageSchema, storageTable); + throw new MaterializedViewMayBeBeingRemovedException(e); + } + return Optional.of(getMaterializedViewDefinition( + materializedView.getOwner(), + materializedView.getViewOriginalText() + .orElseThrow(() -> new TrinoException(HIVE_INVALID_METADATA, "No view original text: " + viewName)), + storageTableName)); + } - Table icebergTable; + SchemaTableName storageTableName = new SchemaTableName(viewName.getSchemaName(), IcebergTableName.tableNameWithType(viewName.getTableName(), MATERIALIZED_VIEW_STORAGE)); + IcebergTableOperations operations = tableOperationsProvider.createTableOperations( + this, + session, + storageTableName.getSchemaName(), + storageTableName.getTableName(), + Optional.empty(), + Optional.empty()); try { - icebergTable = loadTable(session, storageTableName); + TableMetadata metadata = getMaterializedViewTableMetadata(session, storageTableName, materializedView); + operations.initializeFromMetadata(metadata); + Table icebergTable = new BaseTable(operations, quotedTableName(storageTableName), TRINO_METRICS_REPORTER); + + return Optional.of(getMaterializedViewDefinition( + materializedView.getOwner(), + materializedView.getViewOriginalText() + .orElseThrow(() -> new TrinoException(HIVE_INVALID_METADATA, "No view original text: " + viewName)), + storageTableName)); } catch (RuntimeException e) { // The materialized view could be removed concurrently. This may manifest in a number of ways, e.g. - // - io.trino.spi.connector.TableNotFoundException // - org.apache.iceberg.exceptions.NotFoundException when accessing manifest file // - other failures when reading storage table's metadata files // Retry, as we're catching broadly. metastore.invalidateTable(viewName.getSchemaName(), viewName.getTableName()); - metastore.invalidateTable(storageSchema, storageTable); throw new MaterializedViewMayBeBeingRemovedException(e); } - return Optional.of(getMaterializedViewDefinition( - icebergTable, - table.getOwner(), - materializedView.getViewOriginalText() - .orElseThrow(() -> new TrinoException(HIVE_INVALID_METADATA, "No view original text: " + viewName)), - storageTableName)); + } + + @Override + public Optional getMaterializedViewStorageTable(ConnectorSession session, SchemaTableName viewName) + { + Optional tableOptional = metastore.getTable(viewName.getSchemaName(), viewName.getTableName()); + if (tableOptional.isEmpty()) { + return Optional.empty(); + } + + io.trino.plugin.hive.metastore.Table materializedView = tableOptional.get(); + verify(isTrinoMaterializedView(materializedView.getTableType(), materializedView.getParameters()), + "getMaterializedViewStorageTable received a table, not a materialized view"); + + SchemaTableName storageTableName = new SchemaTableName(viewName.getSchemaName(), IcebergTableName.tableNameWithType(viewName.getTableName(), MATERIALIZED_VIEW_STORAGE)); + IcebergTableOperations operations = tableOperationsProvider.createTableOperations( + this, + session, + storageTableName.getSchemaName(), + storageTableName.getTableName(), + Optional.empty(), + Optional.empty()); + + try { + TableMetadata metadata = getMaterializedViewTableMetadata(session, storageTableName, materializedView); + operations.initializeFromMetadata(metadata); + return Optional.of(new BaseTable(operations, quotedTableName(storageTableName), TRINO_METRICS_REPORTER)); + } + catch (NotFoundException e) { + // Removed during reading + return Optional.empty(); + } + } + + private TableMetadata getMaterializedViewTableMetadata(ConnectorSession session, SchemaTableName storageTableName, io.trino.plugin.hive.metastore.Table materializedView) + { + return uncheckedCacheGet(tableMetadataCache, storageTableName, () -> { + String storageMetadataLocation = materializedView.getParameters().get(METADATA_LOCATION_PROP); + checkState(storageMetadataLocation != null, "Storage location missing in definition of materialized view " + materializedView.getTableName()); + TrinoFileSystem fileSystem = fileSystemFactory.create(session); + return TableMetadataParser.read(new ForwardingFileIo(fileSystem), storageMetadataLocation); + }); } @Override @@ -708,4 +909,10 @@ public Optional redirectTable(ConnectorSession session, } return Optional.empty(); } + + @Override + protected void invalidateTableCache(SchemaTableName schemaTableName) + { + tableMetadataCache.invalidate(schemaTableName); + } } diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/hms/TrinoHiveCatalogFactory.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/hms/TrinoHiveCatalogFactory.java index cc3d405a65cd..42f0f760e1c2 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/hms/TrinoHiveCatalogFactory.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/hms/TrinoHiveCatalogFactory.java @@ -14,14 +14,16 @@ package io.trino.plugin.iceberg.catalog.hms; import com.google.inject.Inject; +import io.airlift.concurrent.BoundedExecutor; import io.trino.filesystem.TrinoFileSystemFactory; import io.trino.plugin.base.CatalogName; import io.trino.plugin.hive.NodeVersion; import io.trino.plugin.hive.TrinoViewHiveMetastore; import io.trino.plugin.hive.metastore.HiveMetastoreFactory; import io.trino.plugin.hive.metastore.cache.CachingHiveMetastore; +import io.trino.plugin.hive.security.UsingSystemSecurity; +import io.trino.plugin.iceberg.ForIcebergMetadata; import io.trino.plugin.iceberg.IcebergConfig; -import io.trino.plugin.iceberg.IcebergSecurityConfig; import io.trino.plugin.iceberg.catalog.IcebergTableOperationsProvider; import io.trino.plugin.iceberg.catalog.TrinoCatalog; import io.trino.plugin.iceberg.catalog.TrinoCatalogFactory; @@ -29,9 +31,11 @@ import io.trino.spi.type.TypeManager; import java.util.Optional; +import java.util.concurrent.Executor; +import java.util.concurrent.ExecutorService; -import static io.trino.plugin.hive.metastore.cache.CachingHiveMetastore.memoizeMetastore; -import static io.trino.plugin.iceberg.IcebergSecurityConfig.IcebergSecurity.SYSTEM; +import static com.google.common.util.concurrent.MoreExecutors.directExecutor; +import static io.trino.plugin.hive.metastore.cache.CachingHiveMetastore.createPerTransactionCache; import static io.trino.plugin.iceberg.catalog.AbstractTrinoCatalog.TRINO_CREATED_BY_VALUE; import static java.util.Objects.requireNonNull; @@ -47,6 +51,8 @@ public class TrinoHiveCatalogFactory private final boolean isUniqueTableLocation; private final boolean isUsingSystemSecurity; private final boolean deleteSchemaLocationsFallback; + private final boolean hideMaterializedViewStorageTable; + private final Executor metadataFetchingExecutor; @Inject public TrinoHiveCatalogFactory( @@ -57,7 +63,8 @@ public TrinoHiveCatalogFactory( TypeManager typeManager, IcebergTableOperationsProvider tableOperationsProvider, NodeVersion nodeVersion, - IcebergSecurityConfig securityConfig) + @UsingSystemSecurity boolean isUsingSystemSecurity, + @ForIcebergMetadata ExecutorService metadataExecutorService) { this.catalogName = requireNonNull(catalogName, "catalogName is null"); this.metastoreFactory = requireNonNull(metastoreFactory, "metastoreFactory is null"); @@ -66,14 +73,21 @@ public TrinoHiveCatalogFactory( this.tableOperationsProvider = requireNonNull(tableOperationsProvider, "tableOperationProvider is null"); this.trinoVersion = nodeVersion.toString(); this.isUniqueTableLocation = config.isUniqueTableLocation(); - this.isUsingSystemSecurity = securityConfig.getSecuritySystem() == SYSTEM; + this.isUsingSystemSecurity = isUsingSystemSecurity; this.deleteSchemaLocationsFallback = config.isDeleteSchemaLocationsFallback(); + this.hideMaterializedViewStorageTable = config.isHideMaterializedViewStorageTable(); + if (config.getMetadataParallelism() == 1) { + this.metadataFetchingExecutor = directExecutor(); + } + else { + this.metadataFetchingExecutor = new BoundedExecutor(metadataExecutorService, config.getMetadataParallelism()); + } } @Override public TrinoCatalog create(ConnectorIdentity identity) { - CachingHiveMetastore metastore = memoizeMetastore(metastoreFactory.createMetastore(Optional.of(identity)), 1000); + CachingHiveMetastore metastore = createPerTransactionCache(metastoreFactory.createMetastore(Optional.of(identity)), 1000); return new TrinoHiveCatalog( catalogName, metastore, @@ -83,6 +97,8 @@ public TrinoCatalog create(ConnectorIdentity identity) tableOperationsProvider, isUniqueTableLocation, isUsingSystemSecurity, - deleteSchemaLocationsFallback); + deleteSchemaLocationsFallback, + hideMaterializedViewStorageTable, + metadataFetchingExecutor); } } diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/jdbc/IcebergJdbcTableOperations.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/jdbc/IcebergJdbcTableOperations.java index 1b7e3bda1afa..458435828543 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/jdbc/IcebergJdbcTableOperations.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/jdbc/IcebergJdbcTableOperations.java @@ -67,4 +67,10 @@ protected void commitToExistingTable(TableMetadata base, TableMetadata metadata) jdbcClient.alterTable(database, tableName, newMetadataLocation, currentMetadataLocation); shouldRefresh = true; } + + @Override + protected void commitMaterializedViewRefresh(TableMetadata base, TableMetadata metadata) + { + throw new UnsupportedOperationException(); + } } diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/jdbc/TrinoJdbcCatalog.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/jdbc/TrinoJdbcCatalog.java index 092aef4da5b7..de8d5329f500 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/jdbc/TrinoJdbcCatalog.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/jdbc/TrinoJdbcCatalog.java @@ -13,12 +13,15 @@ */ package io.trino.plugin.iceberg.catalog.jdbc; +import com.google.common.cache.Cache; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; -import com.google.common.collect.ImmutableSet; +import com.google.common.util.concurrent.UncheckedExecutionException; import io.airlift.log.Logger; +import io.trino.cache.EvictableCacheBuilder; import io.trino.filesystem.TrinoFileSystemFactory; import io.trino.plugin.base.CatalogName; +import io.trino.plugin.hive.metastore.TableInfo; import io.trino.plugin.iceberg.catalog.AbstractTrinoCatalog; import io.trino.plugin.iceberg.catalog.IcebergTableOperationsProvider; import io.trino.spi.TrinoException; @@ -38,26 +41,31 @@ import org.apache.iceberg.PartitionSpec; import org.apache.iceberg.Schema; import org.apache.iceberg.SortOrder; -import org.apache.iceberg.Table; import org.apache.iceberg.TableMetadata; import org.apache.iceberg.Transaction; import org.apache.iceberg.catalog.Namespace; import org.apache.iceberg.catalog.TableIdentifier; import org.apache.iceberg.exceptions.NoSuchNamespaceException; import org.apache.iceberg.jdbc.JdbcCatalog; +import org.apache.iceberg.jdbc.UncheckedInterruptedException; +import org.apache.iceberg.jdbc.UncheckedSQLException; +import java.util.HashMap; +import java.util.HashSet; import java.util.Iterator; import java.util.List; import java.util.Map; import java.util.Optional; import java.util.Set; -import java.util.concurrent.ConcurrentHashMap; import java.util.function.Predicate; +import java.util.function.Supplier; import java.util.function.UnaryOperator; +import static com.google.common.base.Throwables.throwIfUnchecked; import static com.google.common.collect.ImmutableList.toImmutableList; import static com.google.common.collect.ImmutableMap.toImmutableMap; import static com.google.common.collect.Maps.transformValues; +import static io.trino.cache.CacheUtils.uncheckedCacheGet; import static io.trino.filesystem.Locations.appendPath; import static io.trino.plugin.iceberg.IcebergErrorCode.ICEBERG_CATALOG_ERROR; import static io.trino.plugin.iceberg.IcebergErrorCode.ICEBERG_INVALID_METADATA; @@ -76,11 +84,15 @@ public class TrinoJdbcCatalog { private static final Logger LOG = Logger.get(TrinoJdbcCatalog.class); + private static final int PER_QUERY_CACHE_SIZE = 1000; + private final JdbcCatalog jdbcCatalog; private final IcebergJdbcClient jdbcClient; private final TrinoFileSystemFactory fileSystemFactory; private final String defaultWarehouseDir; - private final Map tableMetadataCache = new ConcurrentHashMap<>(); + private final Cache tableMetadataCache = EvictableCacheBuilder.newBuilder() + .maximumSize(PER_QUERY_CACHE_SIZE) + .build(); public TrinoJdbcCatalog( CatalogName catalogName, @@ -92,7 +104,7 @@ public TrinoJdbcCatalog( boolean useUniqueTableLocation, String defaultWarehouseDir) { - super(catalogName, typeManager, tableOperationsProvider, useUniqueTableLocation); + super(catalogName, typeManager, tableOperationsProvider, fileSystemFactory, useUniqueTableLocation); this.jdbcCatalog = requireNonNull(jdbcCatalog, "jdbcCatalog is null"); this.jdbcClient = requireNonNull(jdbcClient, "jdbcClient is null"); this.fileSystemFactory = requireNonNull(fileSystemFactory, "fileSystemFactory is null"); @@ -156,20 +168,60 @@ public void renameNamespace(ConnectorSession session, String source, String targ } @Override - public List listTables(ConnectorSession session, Optional namespace) + public List listTables(ConnectorSession session, Optional namespace) + { + List namespaces = listNamespaces(session, namespace); + + // Build as a map and convert to list for removing duplicate entries due to case difference + Map tablesListBuilder = new HashMap<>(); + for (String schemaName : namespaces) { + try { + listTableIdentifiers(schemaName, () -> jdbcCatalog.listTables(Namespace.of(schemaName))).stream() + .map(tableId -> SchemaTableName.schemaTableName(schemaName, tableId.name())) + .forEach(schemaTableName -> tablesListBuilder.put(schemaTableName, new TableInfo(schemaTableName, TableInfo.ExtendedRelationType.TABLE))); + listTableIdentifiers(schemaName, () -> jdbcCatalog.listViews(Namespace.of(schemaName))).stream() + .map(tableId -> SchemaTableName.schemaTableName(schemaName, tableId.name())) + .forEach(schemaTableName -> tablesListBuilder.put(schemaTableName, new TableInfo(schemaTableName, TableInfo.ExtendedRelationType.OTHER_VIEW))); + } + catch (NoSuchNamespaceException e) { + // Namespace may have been deleted + } + } + return ImmutableList.copyOf(tablesListBuilder.values()); + } + + @Override + public List listIcebergTables(ConnectorSession session, Optional namespace) { List namespaces = listNamespaces(session, namespace); + // Build as a set and convert to list for removing duplicate entries due to case difference - ImmutableSet.Builder tablesListBuilder = ImmutableSet.builder(); + Set tablesListBuilder = new HashSet<>(); for (String schemaName : namespaces) { try { - jdbcCatalog.listTables(Namespace.of(schemaName)).forEach(table -> tablesListBuilder.add(new SchemaTableName(schemaName, table.name()))); + listTableIdentifiers(schemaName, () -> jdbcCatalog.listTables(Namespace.of(schemaName))).stream() + .map(tableId -> SchemaTableName.schemaTableName(schemaName, tableId.name())) + .forEach(tablesListBuilder::add); } catch (NoSuchNamespaceException e) { // Namespace may have been deleted } } - return tablesListBuilder.build().asList(); + return ImmutableList.copyOf(tablesListBuilder); + } + + private static List listTableIdentifiers(String namespace, Supplier> tableIdentifiersProvider) + { + try { + return tableIdentifiersProvider.get(); + } + catch (NoSuchNamespaceException e) { + // Namespace may have been deleted during listing + } + catch (UncheckedSQLException | UncheckedInterruptedException e) { + throw new TrinoException(ICEBERG_CATALOG_ERROR, "Failed to list tables from namespace: " + namespace, e); + } + return ImmutableList.of(); } @Override @@ -207,7 +259,7 @@ public Transaction newCreateTableTransaction( Schema schema, PartitionSpec partitionSpec, SortOrder sortOrder, - String location, + Optional location, Map properties) { if (!listNamespaces(session, Optional.of(schemaTableName.getSchemaName())).contains(schemaTableName.getSchemaName())) { @@ -224,6 +276,27 @@ public Transaction newCreateTableTransaction( Optional.of(session.getUser())); } + @Override + public Transaction newCreateOrReplaceTableTransaction( + ConnectorSession session, + SchemaTableName schemaTableName, + Schema schema, + PartitionSpec partitionSpec, + SortOrder sortOrder, + String location, + Map properties) + { + return newCreateOrReplaceTableTransaction( + session, + schemaTableName, + schema, + partitionSpec, + sortOrder, + location, + properties, + Optional.of(session.getUser())); + } + @Override public void registerTable(ConnectorSession session, SchemaTableName tableName, TableMetadata tableMetadata) { @@ -256,6 +329,7 @@ public void dropTable(ConnectorSession session, SchemaTableName schemaTableName) LOG.warn(e, "Failed to delete table data referenced by metadata"); } deleteTableDirectory(fileSystemFactory.create(session), schemaTableName, table.location()); + invalidateTableCache(schemaTableName); } @Override @@ -270,6 +344,7 @@ public void dropCorruptedTable(ConnectorSession session, SchemaTableName schemaT } String tableLocation = metadataLocation.get().replaceFirst("/metadata/[^/]*$", ""); deleteTableDirectory(fileSystemFactory.create(session), schemaTableName, tableLocation); + invalidateTableCache(schemaTableName); } @Override @@ -281,14 +356,23 @@ public void renameTable(ConnectorSession session, SchemaTableName from, SchemaTa catch (RuntimeException e) { throw new TrinoException(ICEBERG_CATALOG_ERROR, "Failed to rename table from %s to %s".formatted(from, to), e); } + invalidateTableCache(from); } @Override - public Table loadTable(ConnectorSession session, SchemaTableName schemaTableName) + public BaseTable loadTable(ConnectorSession session, SchemaTableName schemaTableName) { - TableMetadata metadata = tableMetadataCache.computeIfAbsent( - schemaTableName, - ignore -> ((BaseTable) loadIcebergTable(this, tableOperationsProvider, session, schemaTableName)).operations().current()); + TableMetadata metadata; + try { + metadata = uncheckedCacheGet( + tableMetadataCache, + schemaTableName, + () -> ((BaseTable) loadIcebergTable(this, tableOperationsProvider, session, schemaTableName)).operations().current()); + } + catch (UncheckedExecutionException e) { + throwIfUnchecked(e.getCause()); + throw e; + } return getIcebergTableWithMetadata(this, tableOperationsProvider, session, schemaTableName, metadata); } @@ -383,19 +467,19 @@ public Optional getView(ConnectorSession session, Schem } @Override - public List listMaterializedViews(ConnectorSession session, Optional namespace) + protected Optional doGetMaterializedView(ConnectorSession session, SchemaTableName schemaViewName) { - return ImmutableList.of(); + return Optional.empty(); } @Override - protected Optional doGetMaterializedView(ConnectorSession session, SchemaTableName schemaViewName) + public Optional getMaterializedViewStorageTable(ConnectorSession session, SchemaTableName viewName) { - return Optional.empty(); + throw new TrinoException(NOT_SUPPORTED, "The Iceberg JDBC catalog does not support materialized views"); } @Override - public void createMaterializedView(ConnectorSession session, SchemaTableName schemaViewName, ConnectorMaterializedViewDefinition definition, boolean replace, boolean ignoreExisting) + public void createMaterializedView(ConnectorSession session, SchemaTableName schemaViewName, ConnectorMaterializedViewDefinition definition, Map materializedViewProperties, boolean replace, boolean ignoreExisting) { throw new TrinoException(NOT_SUPPORTED, "createMaterializedView is not supported for Iceberg JDBC catalogs"); } @@ -424,6 +508,12 @@ public Optional redirectTable(ConnectorSession session, return Optional.empty(); } + @Override + protected void invalidateTableCache(SchemaTableName schemaTableName) + { + tableMetadataCache.invalidate(schemaTableName); + } + private static TableIdentifier toIdentifier(SchemaTableName table) { return TableIdentifier.of(table.getSchemaName(), table.getTableName()); diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/nessie/IcebergNessieCatalogModule.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/nessie/IcebergNessieCatalogModule.java index 809c39589bd5..00c79a347dc8 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/nessie/IcebergNessieCatalogModule.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/nessie/IcebergNessieCatalogModule.java @@ -22,8 +22,8 @@ import io.trino.plugin.iceberg.catalog.IcebergTableOperationsProvider; import io.trino.plugin.iceberg.catalog.TrinoCatalogFactory; import org.apache.iceberg.nessie.NessieIcebergClient; +import org.projectnessie.client.NessieClientBuilder; import org.projectnessie.client.api.NessieApiV1; -import org.projectnessie.client.http.HttpClientBuilder; import static io.airlift.configuration.ConfigBinder.configBinder; import static org.weakref.jmx.guice.ExportBinder.newExporter; @@ -46,9 +46,9 @@ protected void setup(Binder binder) public static NessieIcebergClient createNessieIcebergClient(IcebergNessieCatalogConfig icebergNessieCatalogConfig) { return new NessieIcebergClient( - HttpClientBuilder.builder() + NessieClientBuilder.createClientBuilder("http", null) .withUri(icebergNessieCatalogConfig.getServerUri()) - .withEnableApiCompatibilityCheck(false) + //.withEnableApiCompatibilityCheck(false) .build(NessieApiV1.class), icebergNessieCatalogConfig.getDefaultReferenceName(), null, diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/nessie/IcebergNessieTableOperations.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/nessie/IcebergNessieTableOperations.java index 75be83f8293d..236af5851e85 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/nessie/IcebergNessieTableOperations.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/nessie/IcebergNessieTableOperations.java @@ -97,7 +97,8 @@ protected void commitNewTable(TableMetadata metadata) { verify(version.isEmpty(), "commitNewTable called on a table which already exists"); try { - nessieClient.commitTable(null, metadata, writeNewMetadata(metadata, 0), table, toKey(new SchemaTableName(database, this.tableName))); + String contentId = table == null ? null : table.getId(); + nessieClient.commitTable(null, metadata, writeNewMetadata(metadata, 0), contentId, toKey(new SchemaTableName(database, this.tableName))); } catch (NessieNotFoundException e) { throw new TrinoException(ICEBERG_COMMIT_ERROR, format("Cannot commit: ref '%s' no longer exists", nessieClient.refName()), e); @@ -114,7 +115,7 @@ protected void commitToExistingTable(TableMetadata base, TableMetadata metadata) { verify(version.orElseThrow() >= 0, "commitToExistingTable called on a new table"); try { - nessieClient.commitTable(base, metadata, writeNewMetadata(metadata, version.getAsInt() + 1), table, toKey(new SchemaTableName(database, this.tableName))); + nessieClient.commitTable(base, metadata, writeNewMetadata(metadata, version.getAsInt() + 1), table.getId(), toKey(new SchemaTableName(database, this.tableName))); } catch (NessieNotFoundException e) { throw new TrinoException(ICEBERG_COMMIT_ERROR, format("Cannot commit: ref '%s' no longer exists", nessieClient.refName()), e); @@ -126,6 +127,12 @@ protected void commitToExistingTable(TableMetadata base, TableMetadata metadata) shouldRefresh = true; } + @Override + protected void commitMaterializedViewRefresh(TableMetadata base, TableMetadata metadata) + { + throw new UnsupportedOperationException(); + } + private static ContentKey toKey(SchemaTableName tableName) { return ContentKey.of(Namespace.parse(tableName.getSchemaName()), tableName.getTableName()); diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/nessie/TrinoNessieCatalog.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/nessie/TrinoNessieCatalog.java index 36ae7c6fac00..434d333ea3de 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/nessie/TrinoNessieCatalog.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/nessie/TrinoNessieCatalog.java @@ -13,11 +13,15 @@ */ package io.trino.plugin.iceberg.catalog.nessie; +import com.google.common.cache.Cache; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.common.collect.Maps; +import com.google.common.util.concurrent.UncheckedExecutionException; +import io.trino.cache.EvictableCacheBuilder; import io.trino.filesystem.TrinoFileSystemFactory; import io.trino.plugin.base.CatalogName; +import io.trino.plugin.hive.metastore.TableInfo; import io.trino.plugin.iceberg.catalog.AbstractTrinoCatalog; import io.trino.plugin.iceberg.catalog.IcebergTableOperationsProvider; import io.trino.spi.TrinoException; @@ -36,7 +40,6 @@ import org.apache.iceberg.PartitionSpec; import org.apache.iceberg.Schema; import org.apache.iceberg.SortOrder; -import org.apache.iceberg.Table; import org.apache.iceberg.TableMetadata; import org.apache.iceberg.TableOperations; import org.apache.iceberg.Transaction; @@ -49,11 +52,12 @@ import java.util.Map; import java.util.Optional; import java.util.Set; -import java.util.concurrent.ConcurrentHashMap; import java.util.function.Predicate; import java.util.function.UnaryOperator; +import static com.google.common.base.Throwables.throwIfUnchecked; import static com.google.common.collect.ImmutableList.toImmutableList; +import static io.trino.cache.CacheUtils.uncheckedCacheGet; import static io.trino.filesystem.Locations.appendPath; import static io.trino.plugin.iceberg.IcebergSchemaProperties.LOCATION_PROPERTY; import static io.trino.plugin.iceberg.IcebergUtil.getIcebergTableWithMetadata; @@ -67,11 +71,16 @@ public class TrinoNessieCatalog extends AbstractTrinoCatalog { + private static final int PER_QUERY_CACHE_SIZE = 1000; + private final String warehouseLocation; private final NessieIcebergClient nessieClient; - private final Map tableMetadataCache = new ConcurrentHashMap<>(); private final TrinoFileSystemFactory fileSystemFactory; + private final Cache tableMetadataCache = EvictableCacheBuilder.newBuilder() + .maximumSize(PER_QUERY_CACHE_SIZE) + .build(); + public TrinoNessieCatalog( CatalogName catalogName, TypeManager typeManager, @@ -81,7 +90,7 @@ public TrinoNessieCatalog( String warehouseLocation, boolean useUniqueTableLocation) { - super(catalogName, typeManager, tableOperationsProvider, useUniqueTableLocation); + super(catalogName, typeManager, tableOperationsProvider, fileSystemFactory, useUniqueTableLocation); this.fileSystemFactory = requireNonNull(fileSystemFactory, "fileSystemFactory is null"); this.warehouseLocation = requireNonNull(warehouseLocation, "warehouseLocation is null"); this.nessieClient = requireNonNull(nessieClient, "nessieClient is null"); @@ -153,11 +162,20 @@ public void renameNamespace(ConnectorSession session, String source, String targ } @Override - public List listTables(ConnectorSession session, Optional namespace) + public List listTables(ConnectorSession session, Optional namespace) { + // views and materialized views are currently not supported, so everything is a table return nessieClient.listTables(namespace.isEmpty() ? Namespace.empty() : Namespace.of(namespace.get())) .stream() - .map(id -> schemaTableName(id.namespace().toString(), id.name())) + .map(id -> new TableInfo(schemaTableName(id.namespace().toString(), id.name()), TableInfo.ExtendedRelationType.TABLE)) + .collect(toImmutableList()); + } + + @Override + public List listIcebergTables(ConnectorSession session, Optional namespace) + { + return listTables(session, namespace).stream() + .map(TableInfo::tableName) .collect(toImmutableList()); } @@ -182,21 +200,28 @@ public Optional> streamRelationComments( } @Override - public Table loadTable(ConnectorSession session, SchemaTableName table) + public BaseTable loadTable(ConnectorSession session, SchemaTableName table) { - TableMetadata metadata = tableMetadataCache.computeIfAbsent( - table, - ignore -> { - TableOperations operations = tableOperationsProvider.createTableOperations( - this, - session, - table.getSchemaName(), - table.getTableName(), - Optional.empty(), - Optional.empty()); - return new BaseTable(operations, quotedTableName(table)).operations().current(); - }); - + TableMetadata metadata; + try { + metadata = uncheckedCacheGet( + tableMetadataCache, + table, + () -> { + TableOperations operations = tableOperationsProvider.createTableOperations( + this, + session, + table.getSchemaName(), + table.getTableName(), + Optional.empty(), + Optional.empty()); + return new BaseTable(operations, quotedTableName(table)).operations().current(); + }); + } + catch (UncheckedExecutionException e) { + throwIfUnchecked(e.getCause()); + throw e; + } return getIcebergTableWithMetadata( this, tableOperationsProvider, @@ -218,6 +243,7 @@ public void dropTable(ConnectorSession session, SchemaTableName schemaTableName) validateTableCanBeDropped(table); nessieClient.dropTable(toIdentifier(schemaTableName), true); deleteTableDirectory(fileSystemFactory.create(session), schemaTableName, table.location()); + invalidateTableCache(schemaTableName); } @Override @@ -230,6 +256,7 @@ public void dropCorruptedTable(ConnectorSession session, SchemaTableName schemaT public void renameTable(ConnectorSession session, SchemaTableName from, SchemaTableName to) { nessieClient.renameTable(toIdentifier(from), toIdentifier(to)); + invalidateTableCache(from); } @Override @@ -239,7 +266,7 @@ public Transaction newCreateTableTransaction( Schema schema, PartitionSpec partitionSpec, SortOrder sortOrder, - String location, + Optional location, Map properties) { return newCreateTableTransaction( @@ -253,6 +280,27 @@ public Transaction newCreateTableTransaction( Optional.of(session.getUser())); } + @Override + public Transaction newCreateOrReplaceTableTransaction( + ConnectorSession session, + SchemaTableName schemaTableName, + Schema schema, + PartitionSpec partitionSpec, + SortOrder sortOrder, + String location, + Map properties) + { + return newCreateOrReplaceTableTransaction( + session, + schemaTableName, + schema, + partitionSpec, + sortOrder, + location, + properties, + Optional.of(session.getUser())); + } + @Override public void registerTable(ConnectorSession session, SchemaTableName tableName, TableMetadata tableMetadata) { @@ -346,17 +394,12 @@ public Optional getView(ConnectorSession session, Schem return Optional.empty(); } - @Override - public List listMaterializedViews(ConnectorSession session, Optional namespace) - { - return ImmutableList.of(); - } - @Override public void createMaterializedView( ConnectorSession session, SchemaTableName schemaViewName, ConnectorMaterializedViewDefinition definition, + Map materializedViewProperties, boolean replace, boolean ignoreExisting) { @@ -375,6 +418,12 @@ public Optional getMaterializedView(Connect return Optional.empty(); } + @Override + public Optional getMaterializedViewStorageTable(ConnectorSession session, SchemaTableName viewName) + { + throw new TrinoException(NOT_SUPPORTED, "The Iceberg Nessie catalog does not support materialized views"); + } + @Override protected Optional doGetMaterializedView(ConnectorSession session, SchemaTableName schemaViewName) { @@ -392,4 +441,10 @@ public Optional redirectTable(ConnectorSession session, { return Optional.empty(); } + + @Override + protected void invalidateTableCache(SchemaTableName schemaTableName) + { + tableMetadataCache.invalidate(schemaTableName); + } } diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/glue/AWSGlueAsyncAdapterProvider.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/rest/AwsProperties.java similarity index 73% rename from plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/glue/AWSGlueAsyncAdapterProvider.java rename to plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/rest/AwsProperties.java index e4084971e6eb..c28e7e35e5b0 100644 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/glue/AWSGlueAsyncAdapterProvider.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/rest/AwsProperties.java @@ -11,11 +11,12 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package io.trino.plugin.iceberg.catalog.glue; +package io.trino.plugin.iceberg.catalog.rest; -import com.amazonaws.services.glue.AWSGlueAsync; +import java.util.Map; -public interface AWSGlueAsyncAdapterProvider +@FunctionalInterface +public interface AwsProperties { - AWSGlueAsync createAWSGlueAsyncAdapter(AWSGlueAsync delegate); + Map get(); } diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/rest/DefaultIcebergFileSystemFactory.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/rest/DefaultIcebergFileSystemFactory.java new file mode 100644 index 000000000000..e61601e71ea8 --- /dev/null +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/rest/DefaultIcebergFileSystemFactory.java @@ -0,0 +1,42 @@ +/* + * 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.iceberg.catalog.rest; + +import com.google.inject.Inject; +import io.trino.filesystem.TrinoFileSystem; +import io.trino.filesystem.TrinoFileSystemFactory; +import io.trino.plugin.iceberg.IcebergFileSystemFactory; +import io.trino.spi.security.ConnectorIdentity; + +import java.util.Map; + +import static java.util.Objects.requireNonNull; + +public class DefaultIcebergFileSystemFactory + implements IcebergFileSystemFactory +{ + private final TrinoFileSystemFactory fileSystemFactory; + + @Inject + public DefaultIcebergFileSystemFactory(TrinoFileSystemFactory fileSystemFactory) + { + this.fileSystemFactory = requireNonNull(fileSystemFactory, "fileSystemFactory is null"); + } + + @Override + public TrinoFileSystem create(ConnectorIdentity identity, Map fileIoProperties) + { + return fileSystemFactory.create(identity); + } +} diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/rest/IcebergRestCatalogConfig.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/rest/IcebergRestCatalogConfig.java index af8e3a5fe0b4..7f7268f7682d 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/rest/IcebergRestCatalogConfig.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/rest/IcebergRestCatalogConfig.java @@ -15,11 +15,17 @@ import io.airlift.configuration.Config; import io.airlift.configuration.ConfigDescription; +import io.airlift.units.Duration; +import io.airlift.units.MinDuration; import jakarta.validation.constraints.NotNull; +import org.apache.iceberg.CatalogProperties; import java.net.URI; import java.util.Optional; +import static java.util.concurrent.TimeUnit.MILLISECONDS; +import static java.util.concurrent.TimeUnit.MINUTES; + public class IcebergRestCatalogConfig { public enum Security @@ -35,9 +41,17 @@ public enum SessionType } private URI restUri; + private Optional prefix = Optional.empty(); private Optional warehouse = Optional.empty(); + private boolean nestedNamespaceEnabled; private Security security = Security.NONE; private SessionType sessionType = SessionType.NONE; + private Duration sessionTimeout = new Duration(CatalogProperties.AUTH_SESSION_TIMEOUT_MS_DEFAULT, MILLISECONDS); + private boolean vendedCredentialsEnabled; + private boolean viewEndpointsEnabled = true; + private boolean sigV4Enabled; + private boolean caseInsensitiveNameMatching; + private Duration caseInsensitiveNameMatchingCacheTtl = new Duration(1, MINUTES); @NotNull public URI getBaseUri() @@ -55,6 +69,45 @@ public IcebergRestCatalogConfig setBaseUri(String uri) return this; } + public Optional getPrefix() + { + return prefix; + } + + @Config("iceberg.rest-catalog.prefix") + @ConfigDescription("The prefix for the resource path to use with the REST catalog server") + public IcebergRestCatalogConfig setPrefix(String prefix) + { + this.prefix = Optional.ofNullable(prefix); + return this; + } + + public Optional getWarehouse() + { + return warehouse; + } + + @Config("iceberg.rest-catalog.warehouse") + @ConfigDescription("The warehouse location/identifier to use with the REST catalog server") + public IcebergRestCatalogConfig setWarehouse(String warehouse) + { + this.warehouse = Optional.ofNullable(warehouse); + return this; + } + + public boolean isNestedNamespaceEnabled() + { + return nestedNamespaceEnabled; + } + + @Config("iceberg.rest-catalog.nested-namespace-enabled") + @ConfigDescription("Support querying objects under nested namespace") + public IcebergRestCatalogConfig setNestedNamespaceEnabled(boolean nestedNamespaceEnabled) + { + this.nestedNamespaceEnabled = nestedNamespaceEnabled; + return this; + } + @NotNull public Security getSecurity() { @@ -83,16 +136,85 @@ public IcebergRestCatalogConfig setSessionType(SessionType sessionType) return this; } - public Optional getWarehouse() + @NotNull + @MinDuration("0ms") + public Duration getSessionTimeout() { - return warehouse; + return sessionTimeout; } - @Config("iceberg.rest-catalog.warehouse") - @ConfigDescription("The warehouse location/identifier to use with the REST catalog server") - public IcebergRestCatalogConfig setWarehouse(String warehouse) + @Config("iceberg.rest-catalog.session-timeout") + @ConfigDescription("Duration to keep authentication session in cache") + public IcebergRestCatalogConfig setSessionTimeout(Duration sessionTimeout) { - this.warehouse = Optional.ofNullable(warehouse); + this.sessionTimeout = sessionTimeout; + return this; + } + + public boolean isVendedCredentialsEnabled() + { + return vendedCredentialsEnabled; + } + + @Config("iceberg.rest-catalog.vended-credentials-enabled") + @ConfigDescription("Use credentials provided by the REST backend for file system access") + public IcebergRestCatalogConfig setVendedCredentialsEnabled(boolean vendedCredentialsEnabled) + { + this.vendedCredentialsEnabled = vendedCredentialsEnabled; + return this; + } + + public boolean isViewEndpointsEnabled() + { + return viewEndpointsEnabled; + } + + @Config("iceberg.rest-catalog.view-endpoints-enabled") + @ConfigDescription("Enable view endpoints") + public IcebergRestCatalogConfig setViewEndpointsEnabled(boolean viewEndpointsEnabled) + { + this.viewEndpointsEnabled = viewEndpointsEnabled; + return this; + } + + public boolean isSigV4Enabled() + { + return sigV4Enabled; + } + + @Config("iceberg.rest-catalog.sigv4-enabled") + @ConfigDescription("Enable AWS Signature version 4 (SigV4)") + public IcebergRestCatalogConfig setSigV4Enabled(boolean sigV4Enabled) + { + this.sigV4Enabled = sigV4Enabled; + return this; + } + + public boolean isCaseInsensitiveNameMatching() + { + return caseInsensitiveNameMatching; + } + + @Config("iceberg.rest-catalog.case-insensitive-name-matching") + @ConfigDescription("Match object names case-insensitively") + public IcebergRestCatalogConfig setCaseInsensitiveNameMatching(boolean caseInsensitiveNameMatching) + { + this.caseInsensitiveNameMatching = caseInsensitiveNameMatching; + return this; + } + + @NotNull + @MinDuration("0ms") + public Duration getCaseInsensitiveNameMatchingCacheTtl() + { + return caseInsensitiveNameMatchingCacheTtl; + } + + @Config("iceberg.rest-catalog.case-insensitive-name-matching.cache-ttl") + @ConfigDescription("Duration to keep case insensitive object mapping prior to eviction") + public IcebergRestCatalogConfig setCaseInsensitiveNameMatchingCacheTtl(Duration caseInsensitiveNameMatchingCacheTtl) + { + this.caseInsensitiveNameMatchingCacheTtl = caseInsensitiveNameMatchingCacheTtl; return this; } } diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/rest/IcebergRestCatalogFileSystemFactory.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/rest/IcebergRestCatalogFileSystemFactory.java new file mode 100644 index 000000000000..b5e777c7ee15 --- /dev/null +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/rest/IcebergRestCatalogFileSystemFactory.java @@ -0,0 +1,71 @@ +/* + * 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.iceberg.catalog.rest; + +import com.google.common.collect.ImmutableMap; +import com.google.inject.Inject; +import io.trino.filesystem.TrinoFileSystem; +import io.trino.filesystem.TrinoFileSystemFactory; +import io.trino.plugin.iceberg.IcebergFileSystemFactory; +import io.trino.spi.security.ConnectorIdentity; + +import java.util.Map; + +import static io.trino.filesystem.s3.S3FileSystemConstants.EXTRA_CREDENTIALS_ACCESS_KEY_PROPERTY; +import static io.trino.filesystem.s3.S3FileSystemConstants.EXTRA_CREDENTIALS_SECRET_KEY_PROPERTY; +import static io.trino.filesystem.s3.S3FileSystemConstants.EXTRA_CREDENTIALS_SESSION_TOKEN_PROPERTY; +import static java.util.Objects.requireNonNull; + +public class IcebergRestCatalogFileSystemFactory + implements IcebergFileSystemFactory +{ + private static final String VENDED_S3_ACCESS_KEY = "s3.access-key-id"; + private static final String VENDED_S3_SECRET_KEY = "s3.secret-access-key"; + private static final String VENDED_S3_SESSION_TOKEN = "s3.session-token"; + + private final TrinoFileSystemFactory fileSystemFactory; + private final boolean vendedCredentialsEnabled; + + @Inject + public IcebergRestCatalogFileSystemFactory(TrinoFileSystemFactory fileSystemFactory, IcebergRestCatalogConfig config) + { + this.fileSystemFactory = requireNonNull(fileSystemFactory, "fileSystemFactory is null"); + this.vendedCredentialsEnabled = config.isVendedCredentialsEnabled(); + } + + @Override + public TrinoFileSystem create(ConnectorIdentity identity, Map fileIoProperties) + { + if (vendedCredentialsEnabled && + fileIoProperties.containsKey(VENDED_S3_ACCESS_KEY) && + fileIoProperties.containsKey(VENDED_S3_SECRET_KEY) && + fileIoProperties.containsKey(VENDED_S3_SESSION_TOKEN)) { + // Do not include original credentials as they should not be used in vended mode + ConnectorIdentity identityWithExtraCredentials = ConnectorIdentity.forUser(identity.getUser()) + .withGroups(identity.getGroups()) + .withPrincipal(identity.getPrincipal()) + .withEnabledSystemRoles(identity.getEnabledSystemRoles()) + .withConnectorRole(identity.getConnectorRole()) + .withExtraCredentials(ImmutableMap.builder() + .put(EXTRA_CREDENTIALS_ACCESS_KEY_PROPERTY, fileIoProperties.get(VENDED_S3_ACCESS_KEY)) + .put(EXTRA_CREDENTIALS_SECRET_KEY_PROPERTY, fileIoProperties.get(VENDED_S3_SECRET_KEY)) + .put(EXTRA_CREDENTIALS_SESSION_TOKEN_PROPERTY, fileIoProperties.get(VENDED_S3_SESSION_TOKEN)) + .buildOrThrow()) + .build(); + return fileSystemFactory.create(identityWithExtraCredentials); + } + + return fileSystemFactory.create(identity); + } +} diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/rest/IcebergRestCatalogModule.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/rest/IcebergRestCatalogModule.java index 8950212bbbaf..9efc4465c567 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/rest/IcebergRestCatalogModule.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/rest/IcebergRestCatalogModule.java @@ -13,14 +13,20 @@ */ package io.trino.plugin.iceberg.catalog.rest; +import com.google.common.collect.ImmutableMap; import com.google.inject.Binder; import com.google.inject.Scopes; import io.airlift.configuration.AbstractConfigurationAwareModule; +import io.trino.plugin.iceberg.IcebergConfig; +import io.trino.plugin.iceberg.IcebergFileSystemFactory; import io.trino.plugin.iceberg.catalog.TrinoCatalogFactory; import io.trino.plugin.iceberg.catalog.rest.IcebergRestCatalogConfig.Security; +import io.trino.spi.TrinoException; +import static com.google.inject.multibindings.OptionalBinder.newOptionalBinder; import static io.airlift.configuration.ConditionalModule.conditionalModule; import static io.airlift.configuration.ConfigBinder.configBinder; +import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED; public class IcebergRestCatalogModule extends AbstractConfigurationAwareModule @@ -34,7 +40,22 @@ protected void setup(Binder binder) config -> config.getSecurity() == Security.OAUTH2, new OAuth2SecurityModule(), new NoneSecurityModule())); + install(conditionalModule( + IcebergRestCatalogConfig.class, + IcebergRestCatalogConfig::isSigV4Enabled, + internalBinder -> { + configBinder(internalBinder).bindConfig(IcebergRestCatalogSigV4Config.class); + internalBinder.bind(AwsProperties.class).to(SigV4AwsProperties.class).in(Scopes.SINGLETON); + }, + internalBinder -> internalBinder.bind(AwsProperties.class).toInstance(ImmutableMap::of))); binder.bind(TrinoCatalogFactory.class).to(TrinoIcebergRestCatalogFactory.class).in(Scopes.SINGLETON); + newOptionalBinder(binder, IcebergFileSystemFactory.class).setBinding().to(IcebergRestCatalogFileSystemFactory.class).in(Scopes.SINGLETON); + + IcebergConfig icebergConfig = buildConfigObject(IcebergConfig.class); + IcebergRestCatalogConfig restCatalogConfig = buildConfigObject(IcebergRestCatalogConfig.class); + if (restCatalogConfig.isVendedCredentialsEnabled() && icebergConfig.isRegisterTableProcedureEnabled()) { + throw new TrinoException(NOT_SUPPORTED, "Using the `register_table` procedure with vended credentials is currently not supported"); + } } } diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/rest/IcebergRestCatalogSigV4Config.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/rest/IcebergRestCatalogSigV4Config.java new file mode 100644 index 000000000000..e7c80e628762 --- /dev/null +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/rest/IcebergRestCatalogSigV4Config.java @@ -0,0 +1,38 @@ +/* + * 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.iceberg.catalog.rest; + +import io.airlift.configuration.Config; +import io.airlift.configuration.ConfigDescription; +import jakarta.validation.constraints.NotNull; +import org.apache.iceberg.aws.AwsProperties; + +public class IcebergRestCatalogSigV4Config +{ + private String signingName = AwsProperties.REST_SIGNING_NAME_DEFAULT; + + @NotNull + public String getSigningName() + { + return signingName; + } + + @Config("iceberg.rest-catalog.signing-name") + @ConfigDescription("AWS SigV4 signing service name") + public IcebergRestCatalogSigV4Config setSigningName(String signingName) + { + this.signingName = signingName; + return this; + } +} diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/rest/OAuth2SecurityConfig.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/rest/OAuth2SecurityConfig.java index 20a9920430a4..ccb5675d0e66 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/rest/OAuth2SecurityConfig.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/rest/OAuth2SecurityConfig.java @@ -17,13 +17,18 @@ import io.airlift.configuration.ConfigDescription; import io.airlift.configuration.ConfigSecuritySensitive; import jakarta.validation.constraints.AssertTrue; +import org.apache.iceberg.rest.auth.OAuth2Properties; +import java.net.URI; import java.util.Optional; public class OAuth2SecurityConfig { private String credential; + private String scope; private String token; + private URI serverUri; + private boolean tokenRefreshEnabled = OAuth2Properties.TOKEN_REFRESH_ENABLED_DEFAULT; public Optional getCredential() { @@ -39,6 +44,19 @@ public OAuth2SecurityConfig setCredential(String credential) return this; } + public Optional getScope() + { + return Optional.ofNullable(scope); + } + + @Config("iceberg.rest-catalog.oauth2.scope") + @ConfigDescription("The scope which will be used for interactions with the server") + public OAuth2SecurityConfig setScope(String scope) + { + this.scope = scope; + return this; + } + public Optional getToken() { return Optional.ofNullable(token); @@ -53,9 +71,41 @@ public OAuth2SecurityConfig setToken(String token) return this; } + public Optional getServerUri() + { + return Optional.ofNullable(serverUri); + } + + @Config("iceberg.rest-catalog.oauth2.server-uri") + @ConfigDescription("The endpoint to retrieve access token from OAuth2 Server") + public OAuth2SecurityConfig setServerUri(URI serverUri) + { + this.serverUri = serverUri; + return this; + } + + public boolean isTokenRefreshEnabled() + { + return tokenRefreshEnabled; + } + + @Config("iceberg.rest-catalog.oauth2.token-refresh-enabled") + @ConfigDescription("Controls whether a token should be refreshed if information about its expiration time is available") + public OAuth2SecurityConfig setTokenRefreshEnabled(boolean tokenRefreshEnabled) + { + this.tokenRefreshEnabled = tokenRefreshEnabled; + return this; + } + @AssertTrue(message = "OAuth2 requires a credential or token") public boolean credentialOrTokenPresent() { return credential != null || token != null; } + + @AssertTrue(message = "Scope is applicable only when using credential") + public boolean scopePresentOnlyWithCredential() + { + return !(token != null && scope != null); + } } diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/rest/OAuth2SecurityProperties.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/rest/OAuth2SecurityProperties.java index 6a85cb3f80ac..3df0c5d27340 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/rest/OAuth2SecurityProperties.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/rest/OAuth2SecurityProperties.java @@ -33,9 +33,16 @@ public OAuth2SecurityProperties(OAuth2SecurityConfig securityConfig) ImmutableMap.Builder propertiesBuilder = ImmutableMap.builder(); securityConfig.getCredential().ifPresent( - value -> propertiesBuilder.put(OAuth2Properties.CREDENTIAL, value)); + credential -> { + propertiesBuilder.put(OAuth2Properties.CREDENTIAL, credential); + securityConfig.getScope() + .ifPresent(scope -> propertiesBuilder.put(OAuth2Properties.SCOPE, scope)); + }); securityConfig.getToken().ifPresent( value -> propertiesBuilder.put(OAuth2Properties.TOKEN, value)); + securityConfig.getServerUri().ifPresent( + value -> propertiesBuilder.put(OAuth2Properties.OAUTH2_SERVER_URI, value.toString())); + propertiesBuilder.put(OAuth2Properties.TOKEN_REFRESH_ENABLED, String.valueOf(securityConfig.isTokenRefreshEnabled())); this.securityProperties = propertiesBuilder.buildOrThrow(); } diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/rest/SigV4AwsCredentialProvider.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/rest/SigV4AwsCredentialProvider.java new file mode 100644 index 000000000000..67f3f3f7adf4 --- /dev/null +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/rest/SigV4AwsCredentialProvider.java @@ -0,0 +1,130 @@ +/* + * 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.iceberg.catalog.rest; + +import io.trino.spi.TrinoException; +import software.amazon.awssdk.auth.credentials.AwsBasicCredentials; +import software.amazon.awssdk.auth.credentials.AwsCredentials; +import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider; +import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider; +import software.amazon.awssdk.identity.spi.AwsCredentialsIdentity; +import software.amazon.awssdk.identity.spi.ResolveIdentityRequest; +import software.amazon.awssdk.regions.Region; +import software.amazon.awssdk.services.sts.StsClient; +import software.amazon.awssdk.services.sts.StsClientBuilder; +import software.amazon.awssdk.services.sts.auth.StsAssumeRoleCredentialsProvider; + +import java.net.URI; +import java.util.Map; +import java.util.Optional; +import java.util.concurrent.CompletableFuture; +import java.util.function.Consumer; + +import static io.trino.plugin.iceberg.IcebergErrorCode.ICEBERG_CATALOG_ERROR; +import static java.util.Objects.requireNonNull; + +public class SigV4AwsCredentialProvider + implements AwsCredentialsProvider +{ + static final String AWS_STS_ACCESS_KEY_ID = "aws_sts_access_key_id"; + static final String AWS_STS_SECRET_ACCESS_KEY = "aws_sts_secret_access_key"; + static final String AWS_STS_SIGNER_REGION = "aws_sts_signer_region"; + static final String AWS_STS_REGION = "aws_sts_region"; + static final String AWS_STS_ENDPOINT = "aws_sts_endpoint"; + + static final String AWS_IAM_ROLE = "aws_iam_role"; + static final String AWS_ROLE_EXTERNAL_ID = "aws_external_id"; + static final String AWS_IAM_ROLE_SESSION_NAME = "aws_iam_role_session_name"; + + private final AwsCredentialsProvider delegate; + + public SigV4AwsCredentialProvider(AwsCredentialsProvider delegate) + { + this.delegate = requireNonNull(delegate, "delegate is null"); + } + + public static SigV4AwsCredentialProvider create(Map properties) + { + if (properties.containsKey(AWS_IAM_ROLE)) { + String accessKey = properties.get(AWS_STS_ACCESS_KEY_ID); + String secretAccessKey = properties.get(AWS_STS_SECRET_ACCESS_KEY); + + Optional staticCredentialsProvider = createStaticCredentialsProvider(accessKey, secretAccessKey); + return new SigV4AwsCredentialProvider(StsAssumeRoleCredentialsProvider.builder() + .refreshRequest(request -> request + .roleArn(properties.get(AWS_IAM_ROLE)) + .roleSessionName(AWS_IAM_ROLE_SESSION_NAME) + .externalId(properties.get(AWS_ROLE_EXTERNAL_ID))) + .stsClient(createStsClient( + properties.get(AWS_STS_ENDPOINT), + properties.get(AWS_STS_REGION), + properties.get(AWS_STS_SIGNER_REGION), + staticCredentialsProvider)) + .asyncCredentialUpdateEnabled(true) + .build()); + } + + throw new TrinoException(ICEBERG_CATALOG_ERROR, "IAM role configs are not configured"); + } + + @Override + public CompletableFuture resolveIdentity(Consumer consumer) + { + return delegate.resolveIdentity(consumer); + } + + @Override + public CompletableFuture resolveIdentity() + { + return delegate.resolveIdentity(); + } + + @Override + public AwsCredentials resolveCredentials() + { + return delegate.resolveCredentials(); + } + + @Override + public Class identityType() + { + return delegate.identityType(); + } + + @Override + public CompletableFuture resolveIdentity(ResolveIdentityRequest request) + { + return delegate.resolveIdentity(request); + } + + private static Optional createStaticCredentialsProvider(String accessKey, String secretKey) + { + if (accessKey != null || secretKey != null) { + return Optional.of(StaticCredentialsProvider.create( + AwsBasicCredentials.create(accessKey, secretKey))); + } + return Optional.empty(); + } + + private static StsClient createStsClient(String stsEndpoint, String stsRegion, String region, Optional credentialsProvider) + { + StsClientBuilder sts = StsClient.builder(); + Optional.ofNullable(stsEndpoint).map(URI::create).ifPresent(sts::endpointOverride); + Optional.ofNullable(stsRegion) + .or(() -> Optional.ofNullable(region)) + .map(Region::of).ifPresent(sts::region); + credentialsProvider.ifPresent(sts::credentialsProvider); + return sts.build(); + } +} diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/rest/SigV4AwsProperties.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/rest/SigV4AwsProperties.java new file mode 100644 index 000000000000..c19a2e285e4e --- /dev/null +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/rest/SigV4AwsProperties.java @@ -0,0 +1,92 @@ +/* + * 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.iceberg.catalog.rest; + +import com.google.common.collect.ImmutableMap; +import com.google.inject.Inject; +import io.trino.filesystem.s3.S3FileSystemConfig; + +import java.util.Map; +import java.util.Optional; + +import static io.trino.plugin.iceberg.catalog.rest.SigV4AwsCredentialProvider.AWS_IAM_ROLE; +import static io.trino.plugin.iceberg.catalog.rest.SigV4AwsCredentialProvider.AWS_IAM_ROLE_SESSION_NAME; +import static io.trino.plugin.iceberg.catalog.rest.SigV4AwsCredentialProvider.AWS_ROLE_EXTERNAL_ID; +import static io.trino.plugin.iceberg.catalog.rest.SigV4AwsCredentialProvider.AWS_STS_ACCESS_KEY_ID; +import static io.trino.plugin.iceberg.catalog.rest.SigV4AwsCredentialProvider.AWS_STS_ENDPOINT; +import static io.trino.plugin.iceberg.catalog.rest.SigV4AwsCredentialProvider.AWS_STS_REGION; +import static io.trino.plugin.iceberg.catalog.rest.SigV4AwsCredentialProvider.AWS_STS_SECRET_ACCESS_KEY; +import static io.trino.plugin.iceberg.catalog.rest.SigV4AwsCredentialProvider.AWS_STS_SIGNER_REGION; +import static java.util.Objects.requireNonNull; +import static org.apache.iceberg.aws.AwsClientProperties.CLIENT_CREDENTIALS_PROVIDER; +import static org.apache.iceberg.aws.AwsProperties.REST_ACCESS_KEY_ID; +import static org.apache.iceberg.aws.AwsProperties.REST_SECRET_ACCESS_KEY; +import static org.apache.iceberg.aws.AwsProperties.REST_SIGNER_REGION; +import static org.apache.iceberg.aws.AwsProperties.REST_SIGNING_NAME; + +public class SigV4AwsProperties + implements AwsProperties +{ + // Copy of `org.apache.iceberg.aws.AwsClientProperties.CLIENT_CREDENTIAL_PROVIDER_PREFIX` https://github.com/apache/iceberg/blob/ab6fc83ec0269736355a0a89c51e44e822264da8/aws/src/main/java/org/apache/iceberg/aws/AwsClientProperties.java#L69 + private static final String CLIENT_CREDENTIAL_PROVIDER_PREFIX = "client.credentials-provider."; + + private static final String CLIENT_CREDENTIAL_AWS_ACCESS_KEY_ID = CLIENT_CREDENTIAL_PROVIDER_PREFIX + AWS_STS_ACCESS_KEY_ID; + private static final String CLIENT_CREDENTIAL_AWS_SECRET_ACCESS_KEY = CLIENT_CREDENTIAL_PROVIDER_PREFIX + AWS_STS_SECRET_ACCESS_KEY; + private static final String CLIENT_CREDENTIAL_AWS_SIGNER_REGION = CLIENT_CREDENTIAL_PROVIDER_PREFIX + AWS_STS_SIGNER_REGION; + + private static final String CLIENT_CREDENTIAL_AWS_STS_REGION = CLIENT_CREDENTIAL_PROVIDER_PREFIX + AWS_STS_REGION; + private static final String CLIENT_CREDENTIAL_AWS_STS_ENDPOINT = CLIENT_CREDENTIAL_PROVIDER_PREFIX + AWS_STS_ENDPOINT; + private static final String CLIENT_CREDENTIAL_AWS_IAM_ROLE = CLIENT_CREDENTIAL_PROVIDER_PREFIX + AWS_IAM_ROLE; + private static final String CLIENT_CREDENTIAL_AWS_ROLE_EXTERNAL_ID = CLIENT_CREDENTIAL_PROVIDER_PREFIX + AWS_ROLE_EXTERNAL_ID; + private static final String CLIENT_CREDENTIAL_AWS_IAM_ROLE_SESSION_NAME = CLIENT_CREDENTIAL_PROVIDER_PREFIX + AWS_IAM_ROLE_SESSION_NAME; + + private final Map properties; + + @Inject + public SigV4AwsProperties(IcebergRestCatalogSigV4Config sigV4Config, S3FileSystemConfig s3Config) + { + ImmutableMap.Builder builder = ImmutableMap.builder() + .put("rest.sigv4-enabled", "true") + .put(REST_SIGNING_NAME, sigV4Config.getSigningName()) + .put(REST_SIGNER_REGION, requireNonNull(s3Config.getRegion(), "s3.region is null")) + .put("rest-metrics-reporting-enabled", "false"); + + if (s3Config.getIamRole() != null) { + builder + .put(CLIENT_CREDENTIALS_PROVIDER, SigV4AwsCredentialProvider.class.getName()) + .put(CLIENT_CREDENTIAL_AWS_IAM_ROLE, s3Config.getIamRole()) + .put(CLIENT_CREDENTIAL_AWS_IAM_ROLE_SESSION_NAME, "trino-iceberg-rest-catalog") + .put(CLIENT_CREDENTIAL_AWS_SIGNER_REGION, s3Config.getRegion()); + Optional.ofNullable(s3Config.getExternalId()).ifPresent(externalId -> builder.put(CLIENT_CREDENTIAL_AWS_ROLE_EXTERNAL_ID, externalId)); + + Optional.ofNullable(s3Config.getStsRegion()).ifPresent(stsRegion -> builder.put(CLIENT_CREDENTIAL_AWS_STS_REGION, stsRegion)); + Optional.ofNullable(s3Config.getAwsAccessKey()).ifPresent(accessKey -> builder.put(CLIENT_CREDENTIAL_AWS_ACCESS_KEY_ID, accessKey)); + Optional.ofNullable(s3Config.getAwsSecretKey()).ifPresent(secretAccessKey -> builder.put(CLIENT_CREDENTIAL_AWS_SECRET_ACCESS_KEY, secretAccessKey)); + Optional.ofNullable(s3Config.getStsEndpoint()).ifPresent(endpoint -> builder.put(CLIENT_CREDENTIAL_AWS_STS_ENDPOINT, endpoint)); + } + else { + builder + .put(REST_ACCESS_KEY_ID, requireNonNull(s3Config.getAwsAccessKey(), "s3.aws-access-key is null")) + .put(REST_SECRET_ACCESS_KEY, requireNonNull(s3Config.getAwsSecretKey(), "s3.aws-secret-key is null")); + } + + properties = builder.buildOrThrow(); + } + + @Override + public Map get() + { + return properties; + } +} diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/rest/TrinoIcebergRestCatalogFactory.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/rest/TrinoIcebergRestCatalogFactory.java index e883cc3beb4b..cbca0981143e 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/rest/TrinoIcebergRestCatalogFactory.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/rest/TrinoIcebergRestCatalogFactory.java @@ -13,49 +13,74 @@ */ package io.trino.plugin.iceberg.catalog.rest; +import com.google.common.cache.Cache; import com.google.common.collect.ImmutableMap; +import com.google.common.collect.Maps; import com.google.errorprone.annotations.concurrent.GuardedBy; import com.google.inject.Inject; -import io.trino.filesystem.TrinoFileSystemFactory; +import io.airlift.units.Duration; +import io.trino.cache.EvictableCacheBuilder; import io.trino.plugin.base.CatalogName; import io.trino.plugin.hive.NodeVersion; import io.trino.plugin.iceberg.IcebergConfig; +import io.trino.plugin.iceberg.IcebergFileSystemFactory; import io.trino.plugin.iceberg.catalog.TrinoCatalog; import io.trino.plugin.iceberg.catalog.TrinoCatalogFactory; import io.trino.plugin.iceberg.catalog.rest.IcebergRestCatalogConfig.SessionType; import io.trino.plugin.iceberg.fileio.ForwardingFileIo; import io.trino.spi.security.ConnectorIdentity; +import io.trino.spi.type.TypeManager; import org.apache.iceberg.CatalogProperties; +import org.apache.iceberg.catalog.Namespace; +import org.apache.iceberg.catalog.TableIdentifier; import org.apache.iceberg.rest.HTTPClient; import org.apache.iceberg.rest.RESTSessionCatalog; import java.net.URI; +import java.util.Map; import java.util.Optional; +import java.util.Set; import static java.util.Objects.requireNonNull; +import static java.util.concurrent.TimeUnit.MILLISECONDS; +import static org.apache.iceberg.CatalogProperties.AUTH_SESSION_TIMEOUT_MS; +import static org.apache.iceberg.rest.auth.OAuth2Properties.CREDENTIAL; +import static org.apache.iceberg.rest.auth.OAuth2Properties.TOKEN; public class TrinoIcebergRestCatalogFactory implements TrinoCatalogFactory { - private final TrinoFileSystemFactory fileSystemFactory; + private final IcebergFileSystemFactory fileSystemFactory; private final CatalogName catalogName; private final String trinoVersion; private final URI serverUri; + private final Optional prefix; private final Optional warehouse; + private final boolean nestedNamespaceEnabled; private final SessionType sessionType; + private final Duration sessionTimeout; + private final boolean vendedCredentialsEnabled; + private final boolean viewEndpointsEnabled; private final SecurityProperties securityProperties; + private final AwsProperties awsProperties; private final boolean uniqueTableLocation; + private final TypeManager typeManager; + private final boolean caseInsensitiveNameMatching; + private final Cache remoteNamespaceMappingCache; + private final Cache remoteTableMappingCache; @GuardedBy("this") private RESTSessionCatalog icebergCatalog; @Inject public TrinoIcebergRestCatalogFactory( - TrinoFileSystemFactory fileSystemFactory, + IcebergFileSystemFactory fileSystemFactory, CatalogName catalogName, IcebergRestCatalogConfig restConfig, SecurityProperties securityProperties, + AwsProperties awsProperties, IcebergConfig icebergConfig, + TypeManager typeManager, NodeVersion nodeVersion) { this.fileSystemFactory = requireNonNull(fileSystemFactory, "fileSystemFactory is null"); @@ -63,11 +88,27 @@ public TrinoIcebergRestCatalogFactory( this.trinoVersion = requireNonNull(nodeVersion, "nodeVersion is null").toString(); requireNonNull(restConfig, "restConfig is null"); this.serverUri = restConfig.getBaseUri(); + this.prefix = restConfig.getPrefix(); this.warehouse = restConfig.getWarehouse(); + this.nestedNamespaceEnabled = restConfig.isNestedNamespaceEnabled(); this.sessionType = restConfig.getSessionType(); + this.sessionTimeout = restConfig.getSessionTimeout(); + this.vendedCredentialsEnabled = restConfig.isVendedCredentialsEnabled(); + this.viewEndpointsEnabled = restConfig.isViewEndpointsEnabled(); this.securityProperties = requireNonNull(securityProperties, "securityProperties is null"); + this.awsProperties = requireNonNull(awsProperties, "awsProperties is null"); requireNonNull(icebergConfig, "icebergConfig is null"); this.uniqueTableLocation = icebergConfig.isUniqueTableLocation(); + this.typeManager = requireNonNull(typeManager, "typeManager is null"); + this.caseInsensitiveNameMatching = restConfig.isCaseInsensitiveNameMatching(); + this.remoteNamespaceMappingCache = EvictableCacheBuilder.newBuilder() + .expireAfterWrite(restConfig.getCaseInsensitiveNameMatchingCacheTtl().toMillis(), MILLISECONDS) + .shareNothingWhenDisabled() + .build(); + this.remoteTableMappingCache = EvictableCacheBuilder.newBuilder() + .expireAfterWrite(restConfig.getCaseInsensitiveNameMatchingCacheTtl().toMillis(), MILLISECONDS) + .shareNothingWhenDisabled() + .build(); } @Override @@ -79,21 +120,47 @@ public synchronized TrinoCatalog create(ConnectorIdentity identity) ImmutableMap.Builder properties = ImmutableMap.builder(); properties.put(CatalogProperties.URI, serverUri.toString()); warehouse.ifPresent(location -> properties.put(CatalogProperties.WAREHOUSE_LOCATION, location)); + prefix.ifPresent(prefix -> properties.put("prefix", prefix)); + properties.put("view-endpoints-supported", Boolean.toString(viewEndpointsEnabled)); properties.put("trino-version", trinoVersion); + properties.put(AUTH_SESSION_TIMEOUT_MS, String.valueOf(sessionTimeout.toMillis())); properties.putAll(securityProperties.get()); + properties.putAll(awsProperties.get()); + + if (vendedCredentialsEnabled) { + properties.put("header.X-Iceberg-Access-Delegation", "vended-credentials"); + } + RESTSessionCatalog icebergCatalogInstance = new RESTSessionCatalog( - config -> HTTPClient.builder(config).uri(config.get(CatalogProperties.URI)).build(), + config -> HTTPClient.builder(config) + .uri(config.get(CatalogProperties.URI)) + .build(), (context, config) -> { ConnectorIdentity currentIdentity = (context.wrappedIdentity() != null) ? ((ConnectorIdentity) context.wrappedIdentity()) : ConnectorIdentity.ofUser("fake"); - return new ForwardingFileIo(fileSystemFactory.create(currentIdentity)); + return new ForwardingFileIo(fileSystemFactory.create(currentIdentity, config), config); }); icebergCatalogInstance.initialize(catalogName.toString(), properties.buildOrThrow()); icebergCatalog = icebergCatalogInstance; } - return new TrinoRestCatalog(icebergCatalog, catalogName, sessionType, trinoVersion, uniqueTableLocation); + // `OAuth2Properties.SCOPE` is not set as scope passed through credentials is unused in + // https://github.com/apache/iceberg/blob/229d8f6fcd109e6c8943ea7cbb41dab746c6d0ed/core/src/main/java/org/apache/iceberg/rest/auth/OAuth2Util.java#L714-L721 + Map credentials = Maps.filterKeys(securityProperties.get(), key -> Set.of(TOKEN, CREDENTIAL).contains(key)); + + return new TrinoRestCatalog( + icebergCatalog, + catalogName, + sessionType, + credentials, + nestedNamespaceEnabled, + trinoVersion, + typeManager, + uniqueTableLocation, + caseInsensitiveNameMatching, + remoteNamespaceMappingCache, + remoteTableMappingCache); } } diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/rest/TrinoRestCatalog.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/rest/TrinoRestCatalog.java index 1b474bb5c989..2764a042ac02 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/rest/TrinoRestCatalog.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/rest/TrinoRestCatalog.java @@ -13,14 +13,22 @@ */ package io.trino.plugin.iceberg.catalog.rest; +import com.google.common.base.Joiner; +import com.google.common.base.Splitter; +import com.google.common.cache.Cache; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.common.collect.Maps; +import com.google.common.util.concurrent.UncheckedExecutionException; +import io.airlift.log.Logger; import io.jsonwebtoken.impl.DefaultJwtBuilder; import io.jsonwebtoken.jackson.io.JacksonSerializer; +import io.trino.cache.EvictableCacheBuilder; import io.trino.plugin.base.CatalogName; +import io.trino.plugin.hive.metastore.TableInfo; import io.trino.plugin.iceberg.ColumnIdentity; import io.trino.plugin.iceberg.IcebergSchemaProperties; +import io.trino.plugin.iceberg.IcebergUtil; import io.trino.plugin.iceberg.catalog.TrinoCatalog; import io.trino.plugin.iceberg.catalog.rest.IcebergRestCatalogConfig.SessionType; import io.trino.spi.TrinoException; @@ -34,7 +42,9 @@ import io.trino.spi.connector.SchemaNotFoundException; import io.trino.spi.connector.SchemaTableName; import io.trino.spi.connector.TableNotFoundException; +import io.trino.spi.connector.ViewNotFoundException; import io.trino.spi.security.TrinoPrincipal; +import io.trino.spi.type.TypeManager; import org.apache.iceberg.BaseTable; import org.apache.iceberg.PartitionSpec; import org.apache.iceberg.Schema; @@ -42,85 +52,165 @@ import org.apache.iceberg.Table; import org.apache.iceberg.TableMetadata; import org.apache.iceberg.Transaction; +import org.apache.iceberg.catalog.Catalog; import org.apache.iceberg.catalog.Namespace; import org.apache.iceberg.catalog.SessionCatalog; import org.apache.iceberg.catalog.SessionCatalog.SessionContext; import org.apache.iceberg.catalog.TableIdentifier; +import org.apache.iceberg.exceptions.ForbiddenException; import org.apache.iceberg.exceptions.NoSuchNamespaceException; import org.apache.iceberg.exceptions.NoSuchTableException; +import org.apache.iceberg.exceptions.NoSuchViewException; import org.apache.iceberg.exceptions.RESTException; import org.apache.iceberg.rest.RESTSessionCatalog; import org.apache.iceberg.rest.auth.OAuth2Properties; - +import org.apache.iceberg.view.ReplaceViewVersion; +import org.apache.iceberg.view.SQLViewRepresentation; +import org.apache.iceberg.view.UpdateViewProperties; +import org.apache.iceberg.view.View; +import org.apache.iceberg.view.ViewBuilder; +import org.apache.iceberg.view.ViewRepresentation; +import org.apache.iceberg.view.ViewVersion; + +import java.util.Arrays; import java.util.Date; import java.util.Iterator; import java.util.List; import java.util.Map; import java.util.Optional; import java.util.Set; -import java.util.concurrent.ConcurrentHashMap; import java.util.function.Predicate; +import java.util.function.Supplier; import java.util.function.UnaryOperator; +import java.util.stream.Stream; -import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.collect.ImmutableList.toImmutableList; +import static io.trino.cache.CacheUtils.uncheckedCacheGet; import static io.trino.filesystem.Locations.appendPath; -import static io.trino.plugin.hive.HiveMetadata.TABLE_COMMENT; +import static io.trino.plugin.hive.metastore.Table.TABLE_COMMENT; import static io.trino.plugin.iceberg.IcebergErrorCode.ICEBERG_CATALOG_ERROR; +import static io.trino.plugin.iceberg.IcebergErrorCode.ICEBERG_UNSUPPORTED_VIEW_DIALECT; import static io.trino.plugin.iceberg.IcebergUtil.quotedTableName; +import static io.trino.plugin.iceberg.catalog.AbstractTrinoCatalog.ICEBERG_VIEW_RUN_AS_OWNER; import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED; import static java.lang.String.format; +import static java.util.Locale.ENGLISH; import static java.util.Objects.requireNonNull; import static java.util.UUID.randomUUID; +import static org.apache.iceberg.view.ViewProperties.COMMENT; public class TrinoRestCatalog implements TrinoCatalog { + private static final Logger log = Logger.get(TrinoRestCatalog.class); + + private static final int PER_QUERY_CACHE_SIZE = 1000; + private static final String NAMESPACE_SEPARATOR = "."; + private final RESTSessionCatalog restSessionCatalog; private final CatalogName catalogName; + private final TypeManager typeManager; private final SessionType sessionType; + private final Map credentials; + private final boolean nestedNamespaceEnabled; private final String trinoVersion; private final boolean useUniqueTableLocation; + private final boolean caseInsensitiveNameMatching; + private final Cache remoteNamespaceMappingCache; + private final Cache remoteTableMappingCache; - private final Map tableCache = new ConcurrentHashMap<>(); + private final Cache tableCache = EvictableCacheBuilder.newBuilder() + .maximumSize(PER_QUERY_CACHE_SIZE) + .build(); public TrinoRestCatalog( RESTSessionCatalog restSessionCatalog, CatalogName catalogName, SessionType sessionType, + Map credentials, + boolean nestedNamespaceEnabled, String trinoVersion, - boolean useUniqueTableLocation) + TypeManager typeManager, + boolean useUniqueTableLocation, + boolean caseInsensitiveNameMatching, + Cache remoteNamespaceMappingCache, + Cache remoteTableMappingCache) { this.restSessionCatalog = requireNonNull(restSessionCatalog, "restSessionCatalog is null"); this.catalogName = requireNonNull(catalogName, "catalogName is null"); this.sessionType = requireNonNull(sessionType, "sessionType is null"); + this.credentials = ImmutableMap.copyOf(requireNonNull(credentials, "credentials is null")); + this.nestedNamespaceEnabled = nestedNamespaceEnabled; this.trinoVersion = requireNonNull(trinoVersion, "trinoVersion is null"); + this.typeManager = requireNonNull(typeManager, "typeManager is null"); this.useUniqueTableLocation = useUniqueTableLocation; + this.caseInsensitiveNameMatching = caseInsensitiveNameMatching; + this.remoteNamespaceMappingCache = requireNonNull(remoteNamespaceMappingCache, "remoteNamespaceMappingCache is null"); + this.remoteTableMappingCache = requireNonNull(remoteTableMappingCache, "remoteTableMappingCache is null"); + } + + @Override + public Optional getNamespaceSeparator() + { + return Optional.of(NAMESPACE_SEPARATOR); } @Override public boolean namespaceExists(ConnectorSession session, String namespace) { - return restSessionCatalog.namespaceExists(convert(session), Namespace.of(namespace)); + try { + return restSessionCatalog.namespaceExists(convert(session), toRemoteNamespace(session, toNamespace(namespace))); + } + catch (RESTException e) { + throw new TrinoException(ICEBERG_CATALOG_ERROR, "Failed to check namespace '%s'".formatted(namespace), e); + } } @Override public List listNamespaces(ConnectorSession session) { - return restSessionCatalog.listNamespaces(convert(session)).stream() - .map(Namespace::toString) - .collect(toImmutableList()); + if (nestedNamespaceEnabled) { + return collectNamespaces(session, Namespace.empty()); + } + try { + return restSessionCatalog.listNamespaces(convert(session)).stream() + .map(this::toSchemaName) + .collect(toImmutableList()); + } + catch (RESTException e) { + throw new TrinoException(ICEBERG_CATALOG_ERROR, "Failed to list namespaces", e); + } + } + + private List collectNamespaces(ConnectorSession session, Namespace parentNamespace) + { + try { + return restSessionCatalog.listNamespaces(convert(session), parentNamespace).stream() + .flatMap(childNamespace -> Stream.concat( + Stream.of(childNamespace.toString()), + collectNamespaces(session, childNamespace).stream())) + .collect(toImmutableList()); + } + catch (RESTException e) { + throw new TrinoException(ICEBERG_CATALOG_ERROR, "Failed to list namespaces", e); + } } @Override public void dropNamespace(ConnectorSession session, String namespace) { try { - restSessionCatalog.dropNamespace(convert(session), Namespace.of(namespace)); + restSessionCatalog.dropNamespace(convert(session), toRemoteNamespace(session, toNamespace(namespace))); } catch (NoSuchNamespaceException e) { throw new SchemaNotFoundException(namespace); } + catch (RESTException e) { + throw new TrinoException(ICEBERG_CATALOG_ERROR, "Failed to drop namespace '%s'".formatted(namespace), e); + } + if (caseInsensitiveNameMatching) { + remoteNamespaceMappingCache.invalidate(toNamespace(namespace)); + } } @Override @@ -128,11 +218,14 @@ public Map loadNamespaceMetadata(ConnectorSession session, Strin { try { // Return immutable metadata as direct modifications will not be reflected on the namespace - return ImmutableMap.copyOf(restSessionCatalog.loadNamespaceMetadata(convert(session), Namespace.of(namespace))); + return ImmutableMap.copyOf(restSessionCatalog.loadNamespaceMetadata(convert(session), toRemoteNamespace(session, toNamespace(namespace)))); } catch (NoSuchNamespaceException e) { throw new SchemaNotFoundException(namespace); } + catch (RESTException e) { + throw new TrinoException(ICEBERG_CATALOG_ERROR, "Failed to load metadata for namespace '%s'".formatted(namespace), e); + } } @Override @@ -145,15 +238,20 @@ public Optional getNamespacePrincipal(ConnectorSession session, @Override public void createNamespace(ConnectorSession session, String namespace, Map properties, TrinoPrincipal owner) { - restSessionCatalog.createNamespace( - convert(session), - Namespace.of(namespace), - Maps.transformValues(properties, property -> { - if (property instanceof String stringProperty) { - return stringProperty; - } - throw new TrinoException(NOT_SUPPORTED, "Non-string properties are not support for Iceberg REST catalog"); - })); + try { + restSessionCatalog.createNamespace( + convert(session), + toNamespace(namespace), + Maps.transformValues(properties, property -> { + if (property instanceof String stringProperty) { + return stringProperty; + } + throw new TrinoException(NOT_SUPPORTED, "Non-string properties are not support for Iceberg REST catalog"); + })); + } + catch (RESTException e) { + throw new TrinoException(ICEBERG_CATALOG_ERROR, "Failed to create namespace '%s'".formatted(namespace), e); + } } @Override @@ -169,35 +267,98 @@ public void renameNamespace(ConnectorSession session, String source, String targ } @Override - public List listTables(ConnectorSession session, Optional namespace) + public List listTables(ConnectorSession session, Optional namespace) { SessionContext sessionContext = convert(session); - List namespaces; + List namespaces = listNamespaces(session, namespace); - if (namespace.isPresent() && namespaceExists(session, namespace.get())) { - namespaces = ImmutableList.of(Namespace.of(namespace.get())); - } - else { - namespaces = listNamespaces(session).stream() - .map(Namespace::of) - .collect(toImmutableList()); + ImmutableList.Builder tables = ImmutableList.builder(); + for (Namespace restNamespace : namespaces) { + listTableIdentifiers(restNamespace, () -> { + try { + return restSessionCatalog.listTables(sessionContext, toRemoteNamespace(session, restNamespace)); + } + catch (RESTException e) { + throw new TrinoException(ICEBERG_CATALOG_ERROR, "Failed to list tables", e); + } + }).stream() + .map(id -> new TableInfo(SchemaTableName.schemaTableName(toSchemaName(id.namespace()), id.name()), TableInfo.ExtendedRelationType.TABLE)) + .forEach(tables::add); + listTableIdentifiers(restNamespace, () -> { + try { + return restSessionCatalog.listViews(sessionContext, toRemoteNamespace(session, restNamespace)); + } + catch (RESTException e) { + throw new TrinoException(ICEBERG_CATALOG_ERROR, "Failed to list views", e); + } + }).stream() + .map(id -> new TableInfo(SchemaTableName.schemaTableName(toSchemaName(id.namespace()), id.name()), TableInfo.ExtendedRelationType.OTHER_VIEW)) + .forEach(tables::add); } + return tables.build(); + } + + @Override + public List listIcebergTables(ConnectorSession session, Optional namespace) + { + SessionContext sessionContext = convert(session); + List namespaces = listNamespaces(session, namespace); ImmutableList.Builder tables = ImmutableList.builder(); for (Namespace restNamespace : namespaces) { - try { - tables.addAll( - restSessionCatalog.listTables(sessionContext, restNamespace).stream() - .map(id -> SchemaTableName.schemaTableName(id.namespace().toString(), id.name())) - .collect(toImmutableList())); - } - catch (NoSuchNamespaceException e) { - // Namespace may have been deleted during listing - } + listTableIdentifiers(restNamespace, () -> { + try { + return restSessionCatalog.listTables(sessionContext, toRemoteNamespace(session, restNamespace)); + } + catch (RESTException e) { + throw new TrinoException(ICEBERG_CATALOG_ERROR, "Failed to list tables", e); + } + }).stream() + .map(id -> SchemaTableName.schemaTableName(toSchemaName(id.namespace()), id.name())) + .forEach(tables::add); } return tables.build(); } + @Override + public List listViews(ConnectorSession session, Optional namespace) + { + SessionContext sessionContext = convert(session); + List namespaces = listNamespaces(session, namespace); + + ImmutableList.Builder viewNames = ImmutableList.builder(); + for (Namespace restNamespace : namespaces) { + listTableIdentifiers(restNamespace, () -> { + try { + return restSessionCatalog.listViews(sessionContext, toRemoteNamespace(session, restNamespace)); + } + catch (RESTException e) { + throw new TrinoException(ICEBERG_CATALOG_ERROR, "Failed to list views", e); + } + }).stream() + .map(id -> SchemaTableName.schemaTableName(id.namespace().toString(), id.name())) + .forEach(viewNames::add); + } + return viewNames.build(); + } + + private static List listTableIdentifiers(Namespace restNamespace, Supplier> tableIdentifiersProvider) + { + try { + return tableIdentifiersProvider.get(); + } + catch (NoSuchNamespaceException e) { + // Namespace may have been deleted during listing + } + catch (ForbiddenException e) { + log.debug(e, "Failed to list tables from %s namespace because of insufficient permissions", restNamespace); + } + catch (RESTException e) { + throw new TrinoException(ICEBERG_CATALOG_ERROR, format("Failed to list tables from namespace: %s", restNamespace), e); + } + return ImmutableList.of(); + } + @Override public Optional> streamRelationColumns( ConnectorSession session, @@ -220,6 +381,32 @@ public Optional> streamRelationComments( @Override public Transaction newCreateTableTransaction( + ConnectorSession session, + SchemaTableName schemaTableName, + Schema schema, + PartitionSpec partitionSpec, + SortOrder sortOrder, + Optional location, + Map properties) + { + try { + Catalog.TableBuilder tableBuilder = restSessionCatalog.buildTable(convert(session), toRemoteTable(session, schemaTableName, true), schema) + .withPartitionSpec(partitionSpec) + .withSortOrder(sortOrder) + .withProperties(properties); + if (location.isEmpty()) { + // TODO Replace with createTransaction once S3 Tables supports stage-create option + return tableBuilder.create().newTransaction(); + } + return tableBuilder.withLocation(location.get()).createTransaction(); + } + catch (RESTException e) { + throw new TrinoException(ICEBERG_CATALOG_ERROR, "Failed to create transaction", e); + } + } + + @Override + public Transaction newCreateOrReplaceTableTransaction( ConnectorSession session, SchemaTableName schemaTableName, Schema schema, @@ -228,32 +415,59 @@ public Transaction newCreateTableTransaction( String location, Map properties) { - return restSessionCatalog.buildTable(convert(session), toIdentifier(schemaTableName), schema) - .withPartitionSpec(partitionSpec) - .withSortOrder(sortOrder) - .withLocation(location) - .withProperties(properties) - .createTransaction(); + try { + return restSessionCatalog.buildTable(convert(session), toRemoteTable(session, schemaTableName, true), schema) + .withPartitionSpec(partitionSpec) + .withSortOrder(sortOrder) + .withLocation(location) + .withProperties(properties) + .createOrReplaceTransaction(); + } + catch (RESTException e) { + throw new TrinoException(ICEBERG_CATALOG_ERROR, "Failed to create transaction", e); + } } @Override public void registerTable(ConnectorSession session, SchemaTableName tableName, TableMetadata tableMetadata) { - throw new TrinoException(NOT_SUPPORTED, "registerTable is not supported for Iceberg REST catalog"); + TableIdentifier tableIdentifier = TableIdentifier.of(toRemoteNamespace(session, toNamespace(tableName.getSchemaName())), tableName.getTableName()); + try { + restSessionCatalog.registerTable(convert(session), tableIdentifier, tableMetadata.metadataFileLocation()); + } + catch (RESTException e) { + throw new TrinoException(ICEBERG_CATALOG_ERROR, "Failed to register table '%s'".formatted(tableName.getTableName()), e); + } } @Override public void unregisterTable(ConnectorSession session, SchemaTableName tableName) { - throw new TrinoException(NOT_SUPPORTED, "unregisterTable is not supported for Iceberg REST catalogs"); + try { + if (!restSessionCatalog.dropTable(convert(session), toRemoteTable(session, tableName, true))) { + throw new TableNotFoundException(tableName); + } + } + catch (RESTException e) { + throw new TrinoException(ICEBERG_CATALOG_ERROR, "Failed to unregister table '%s'".formatted(tableName.getTableName()), e); + } + invalidateTableCache(tableName); + invalidateTableMappingCache(tableName); } @Override public void dropTable(ConnectorSession session, SchemaTableName schemaTableName) { - if (!restSessionCatalog.purgeTable(convert(session), toIdentifier(schemaTableName))) { - throw new TrinoException(ICEBERG_CATALOG_ERROR, format("Failed to drop table: %s", schemaTableName)); + try { + if (!restSessionCatalog.purgeTable(convert(session), toRemoteTable(session, schemaTableName, true))) { + throw new TrinoException(ICEBERG_CATALOG_ERROR, "Failed to drop table '%s'".formatted(schemaTableName)); + } + } + catch (RESTException e) { + throw new TrinoException(ICEBERG_CATALOG_ERROR, "Failed to drop table '%s'".formatted(schemaTableName.getTableName()), e); } + invalidateTableCache(schemaTableName); + invalidateTableMappingCache(schemaTableName); } @Override @@ -268,31 +482,57 @@ public void dropCorruptedTable(ConnectorSession session, SchemaTableName schemaT public void renameTable(ConnectorSession session, SchemaTableName from, SchemaTableName to) { try { - restSessionCatalog.renameTable(convert(session), toIdentifier(from), toIdentifier(to)); + restSessionCatalog.renameTable(convert(session), toRemoteTable(session, from, true), toRemoteTable(session, to, true)); } catch (RESTException e) { throw new TrinoException(ICEBERG_CATALOG_ERROR, format("Failed to rename table %s to %s", from, to), e); } + invalidateTableCache(from); + invalidateTableMappingCache(from); } @Override - public Table loadTable(ConnectorSession session, SchemaTableName schemaTableName) + public BaseTable loadTable(ConnectorSession session, SchemaTableName schemaTableName) { try { - return tableCache.computeIfAbsent( - schemaTableName.toString(), - key -> { - BaseTable baseTable = (BaseTable) restSessionCatalog.loadTable(convert(session), toIdentifier(schemaTableName)); + return uncheckedCacheGet( + tableCache, + schemaTableName, + () -> { + BaseTable baseTable; + try { + baseTable = (BaseTable) restSessionCatalog.loadTable(convert(session), toRemoteObject(session, schemaTableName)); + } + catch (RESTException e) { + throw new TrinoException(ICEBERG_CATALOG_ERROR, "Failed to load table '%s'".formatted(schemaTableName.getTableName()), e); + } // Creating a new base table is necessary to adhere to Trino's expectations for quoted table names return new BaseTable(baseTable.operations(), quotedTableName(schemaTableName)); }); } - catch (NoSuchTableException e) { - throw new TableNotFoundException(schemaTableName, e); + catch (UncheckedExecutionException e) { + if (e.getCause() instanceof NoSuchTableException) { + throw new TableNotFoundException(schemaTableName, e.getCause()); + } + throw new TrinoException(ICEBERG_CATALOG_ERROR, format("Failed to load table: %s", schemaTableName), e.getCause()); + } + } + + private TableIdentifier toRemoteObject(ConnectorSession session, SchemaTableName schemaTableName) + { + TableIdentifier remoteTable = toRemoteTable(session, schemaTableName, false); + if (!remoteTable.name().equals(schemaTableName.getTableName())) { + return remoteTable; + } + + TableIdentifier remoteView = toRemoteView(session, schemaTableName, false); + if (!remoteView.name().equals(schemaTableName.getTableName())) { + return remoteView; } - catch (RuntimeException e) { - throw new TrinoException(ICEBERG_CATALOG_ERROR, format("Failed to load table: %s", schemaTableName), e); + if (remoteView.name().equals(schemaTableName.getTableName()) && remoteTable.name().equals(schemaTableName.getTableName())) { + return remoteTable; } + throw new RuntimeException("Unable to find remote object"); } @Override @@ -304,13 +544,20 @@ public Map> tryGetColumnMetadata(Connector @Override public void updateTableComment(ConnectorSession session, SchemaTableName schemaTableName, Optional comment) { - Table icebergTable = restSessionCatalog.loadTable(convert(session), toIdentifier(schemaTableName)); + Table icebergTable; + try { + icebergTable = restSessionCatalog.loadTable(convert(session), toRemoteTable(session, schemaTableName, true)); + } + catch (RESTException e) { + throw new TrinoException(ICEBERG_CATALOG_ERROR, "Failed to load table '%s'".formatted(schemaTableName.getTableName()), e); + } if (comment.isEmpty()) { icebergTable.updateProperties().remove(TABLE_COMMENT).commit(); } else { icebergTable.updateProperties().set(TABLE_COMMENT, comment.get()).commit(); } + invalidateTableCache(schemaTableName); } @Override @@ -320,11 +567,12 @@ public String defaultTableLocation(ConnectorSession session, SchemaTableName sch Map properties = loadNamespaceMetadata(session, schemaTableName.getSchemaName()); String databaseLocation = (String) properties.get(IcebergSchemaProperties.LOCATION_PROPERTY); - checkArgument(databaseLocation != null, "location must be set for %s", schemaTableName.getSchemaName()); - - if (databaseLocation.endsWith("/")) { - return databaseLocation + tableName; + if (databaseLocation == null) { + // Iceberg REST catalog doesn't require location property. + // S3 Tables doesn't return the property. + return null; } + return appendPath(databaseLocation, tableName); } @@ -346,13 +594,40 @@ public void setTablePrincipal(ConnectorSession session, SchemaTableName schemaTa @Override public void createView(ConnectorSession session, SchemaTableName schemaViewName, ConnectorViewDefinition definition, boolean replace) { - throw new TrinoException(NOT_SUPPORTED, "createView is not supported for Iceberg REST catalog"); + ImmutableMap.Builder properties = ImmutableMap.builder(); + definition.getOwner().ifPresent(owner -> properties.put(ICEBERG_VIEW_RUN_AS_OWNER, owner)); + definition.getComment().ifPresent(comment -> properties.put(COMMENT, comment)); + Schema schema = IcebergUtil.schemaFromViewColumns(typeManager, definition.getColumns()); + ViewBuilder viewBuilder = restSessionCatalog.buildView(convert(session), toRemoteView(session, schemaViewName, true)); + viewBuilder = viewBuilder.withSchema(schema) + .withQuery("trino", definition.getOriginalSql()) + .withDefaultNamespace(toRemoteNamespace(session, toNamespace(schemaViewName.getSchemaName()))) + .withDefaultCatalog(definition.getCatalog().orElse(null)) + .withProperties(properties.buildOrThrow()) + .withLocation(defaultTableLocation(session, schemaViewName)); + try { + if (replace) { + viewBuilder.createOrReplace(); + } + else { + viewBuilder.create(); + } + } + catch (RESTException e) { + throw new TrinoException(ICEBERG_CATALOG_ERROR, "Failed to create view '%s'".formatted(schemaViewName.getTableName()), e); + } } @Override public void renameView(ConnectorSession session, SchemaTableName source, SchemaTableName target) { - throw new TrinoException(NOT_SUPPORTED, "renameView is not supported for Iceberg REST catalog"); + try { + restSessionCatalog.renameView(convert(session), toRemoteView(session, source, true), toRemoteView(session, target, true)); + } + catch (RESTException e) { + throw new TrinoException(ICEBERG_CATALOG_ERROR, "Failed to rename view '%s' to '%s'".formatted(source, target), e); + } + invalidateTableMappingCache(source); } @Override @@ -364,37 +639,98 @@ public void setViewPrincipal(ConnectorSession session, SchemaTableName schemaVie @Override public void dropView(ConnectorSession session, SchemaTableName schemaViewName) { - throw new TrinoException(NOT_SUPPORTED, "dropView is not supported for Iceberg REST catalog"); + try { + restSessionCatalog.dropView(convert(session), toRemoteView(session, schemaViewName, true)); + } + catch (RESTException e) { + throw new TrinoException(ICEBERG_CATALOG_ERROR, "Failed to drop view '%s'".formatted(schemaViewName.getTableName()), e); + } + invalidateTableMappingCache(schemaViewName); } @Override - public List listViews(ConnectorSession session, Optional namespace) + public Map getViews(ConnectorSession session, Optional namespace) { - return ImmutableList.of(); + SessionContext sessionContext = convert(session); + ImmutableMap.Builder views = ImmutableMap.builder(); + for (Namespace restNamespace : listNamespaces(session, namespace)) { + List restViews; + try { + restViews = restSessionCatalog.listViews(sessionContext, toRemoteNamespace(session, restNamespace)); + } + catch (RESTException e) { + throw new TrinoException(ICEBERG_CATALOG_ERROR, "Failed to list views", e); + } + for (TableIdentifier restView : restViews) { + SchemaTableName schemaTableName = SchemaTableName.schemaTableName(restView.namespace().toString(), restView.name()); + try { + getView(session, schemaTableName).ifPresent(view -> views.put(schemaTableName, view)); + } + catch (TrinoException e) { + if (e.getErrorCode().equals(ICEBERG_UNSUPPORTED_VIEW_DIALECT.toErrorCode())) { + log.debug(e, "Skip unsupported view dialect: %s", schemaTableName); + continue; + } + throw e; + } + } + } + + return views.buildOrThrow(); } @Override - public Map getViews(ConnectorSession session, Optional namespace) + public Optional getView(ConnectorSession session, SchemaTableName viewName) { - return ImmutableMap.of(); + return getIcebergView(session, viewName, false).flatMap(view -> { + SQLViewRepresentation sqlView = view.sqlFor("trino"); + if (!sqlView.dialect().equalsIgnoreCase("trino")) { + throw new TrinoException(ICEBERG_UNSUPPORTED_VIEW_DIALECT, "Cannot read unsupported dialect '%s' for view '%s'".formatted(sqlView.dialect(), viewName)); + } + + Optional comment = Optional.ofNullable(view.properties().get(COMMENT)); + List viewColumns = IcebergUtil.viewColumnsFromSchema(typeManager, view.schema()); + ViewVersion currentVersion = view.currentVersion(); + Optional catalog = Optional.ofNullable(currentVersion.defaultCatalog()); + Optional schema = Optional.empty(); + if (catalog.isPresent() && !currentVersion.defaultNamespace().isEmpty()) { + schema = Optional.of(currentVersion.defaultNamespace().toString()); + } + + Optional owner = Optional.ofNullable(view.properties().get(ICEBERG_VIEW_RUN_AS_OWNER)); + return Optional.of(new ConnectorViewDefinition(sqlView.sql(), catalog, schema, viewColumns, comment, owner, owner.isEmpty())); + }); } - @Override - public Optional getView(ConnectorSession session, SchemaTableName viewName) + private Optional getIcebergView(ConnectorSession session, SchemaTableName viewName, boolean getCached) { - return Optional.empty(); + try { + return Optional.of(restSessionCatalog.loadView(convert(session), toRemoteView(session, viewName, getCached))); + } + catch (NoSuchViewException e) { + return Optional.empty(); + } + catch (RESTException e) { + throw new TrinoException(ICEBERG_CATALOG_ERROR, "Failed to load view '%s'".formatted(viewName.getTableName()), e); + } } @Override - public List listMaterializedViews(ConnectorSession session, Optional namespace) + public void createMaterializedView( + ConnectorSession session, + SchemaTableName viewName, + ConnectorMaterializedViewDefinition definition, + Map materializedViewProperties, + boolean replace, + boolean ignoreExisting) { - return ImmutableList.of(); + throw new TrinoException(NOT_SUPPORTED, "createMaterializedView is not supported for Iceberg REST catalog"); } @Override - public void createMaterializedView(ConnectorSession session, SchemaTableName viewName, ConnectorMaterializedViewDefinition definition, boolean replace, boolean ignoreExisting) + public void updateMaterializedViewColumnComment(ConnectorSession session, SchemaTableName schemaViewName, String columnName, Optional comment) { - throw new TrinoException(NOT_SUPPORTED, "createMaterializedView is not supported for Iceberg REST catalog"); + throw new TrinoException(NOT_SUPPORTED, "updateMaterializedViewColumnComment is not supported for Iceberg REST catalog"); } @Override @@ -409,10 +745,16 @@ public Optional getMaterializedView(Connect return Optional.empty(); } + @Override + public Optional getMaterializedViewStorageTable(ConnectorSession session, SchemaTableName viewName) + { + throw new TrinoException(NOT_SUPPORTED, "The Iceberg REST catalog does not support materialized views"); + } + @Override public void renameMaterializedView(ConnectorSession session, SchemaTableName source, SchemaTableName target) { - throw new TrinoException(NOT_SUPPORTED, "renameMaterializedView is not supported for Iceberg REST catalog"); + throw new TrinoException(NOT_SUPPORTED, "renameMaterialezedView is not supported for Iceberg REST catalog"); } @Override @@ -432,25 +774,39 @@ public Optional redirectTable(ConnectorSession session, @Override public void updateViewComment(ConnectorSession session, SchemaTableName schemaViewName, Optional comment) { - throw new TrinoException(NOT_SUPPORTED, "updateViewComment is not supported for Iceberg REST catalog"); + View view = getIcebergView(session, schemaViewName, true).orElseThrow(() -> new ViewNotFoundException(schemaViewName)); + UpdateViewProperties updateViewProperties = view.updateProperties(); + comment.ifPresentOrElse( + value -> updateViewProperties.set(COMMENT, value), + () -> updateViewProperties.remove(COMMENT)); + updateViewProperties.commit(); } @Override public void updateViewColumnComment(ConnectorSession session, SchemaTableName schemaViewName, String columnName, Optional comment) { - throw new TrinoException(NOT_SUPPORTED, "updateViewColumnComment is not supported for Iceberg REST catalog"); - } + View view = getIcebergView(session, schemaViewName, true) + .orElseThrow(() -> new ViewNotFoundException(schemaViewName)); + + ViewVersion current = view.currentVersion(); + Schema updatedSchema = IcebergUtil.updateColumnComment(view.schema(), columnName, comment.orElse(null)); + ReplaceViewVersion replaceViewVersion = view.replaceVersion() + .withSchema(updatedSchema) + .withDefaultCatalog(current.defaultCatalog()) + .withDefaultNamespace(current.defaultNamespace()); + for (ViewRepresentation representation : view.currentVersion().representations()) { + if (representation instanceof SQLViewRepresentation sqlViewRepresentation) { + replaceViewVersion.withQuery(sqlViewRepresentation.dialect(), sqlViewRepresentation.sql()); + } + } - @Override - public void updateMaterializedViewColumnComment(ConnectorSession session, SchemaTableName schemaViewName, String columnName, Optional comment) - { - throw new TrinoException(NOT_SUPPORTED, "updateMaterializedViewColumnComment is not supported for Iceberg REST catalog"); + replaceViewVersion.commit(); } private SessionCatalog.SessionContext convert(ConnectorSession session) { return switch (sessionType) { - case NONE -> new SessionContext(randomUUID().toString(), null, null, ImmutableMap.of(), session.getIdentity()); + case NONE -> new SessionContext(randomUUID().toString(), null, credentials, ImmutableMap.of(), session.getIdentity()); case USER -> { String sessionId = format("%s-%s", session.getUser(), session.getSource().orElse("default")); @@ -482,8 +838,152 @@ private SessionCatalog.SessionContext convert(ConnectorSession session) }; } + private void invalidateTableCache(SchemaTableName schemaTableName) + { + tableCache.invalidate(schemaTableName); + } + + private void invalidateTableMappingCache(SchemaTableName schemaTableName) + { + if (caseInsensitiveNameMatching) { + remoteTableMappingCache.invalidate(toIdentifier(schemaTableName)); + } + } + + private Namespace toNamespace(String schemaName) + { + if (!nestedNamespaceEnabled && schemaName.contains(NAMESPACE_SEPARATOR)) { + throw new TrinoException(NOT_SUPPORTED, "Nested namespace is not enabled for this catalog"); + } + return Namespace.of(Splitter.on(NAMESPACE_SEPARATOR).omitEmptyStrings().trimResults().splitToList(schemaName).toArray(new String[0])); + } + + private String toSchemaName(Namespace namespace) + { + if (!nestedNamespaceEnabled && namespace.length() != 1) { + throw new TrinoException(NOT_SUPPORTED, "Nested namespace is not enabled for this catalog"); + } + return String.join(NAMESPACE_SEPARATOR, namespace.levels()); + } + private static TableIdentifier toIdentifier(SchemaTableName schemaTableName) { return TableIdentifier.of(schemaTableName.getSchemaName(), schemaTableName.getTableName()); } + + private List listNamespaces(ConnectorSession session, Optional namespace) + { + if (namespace.isEmpty()) { + return listNamespaces(session).stream() + .map(this::toNamespace) + .collect(toImmutableList()); + } + + return ImmutableList.of(toNamespace(namespace.get())); + } + + private TableIdentifier toRemoteTable(ConnectorSession session, SchemaTableName schemaTableName, boolean getCached) + { + TableIdentifier tableIdentifier = toIdentifier(schemaTableName); + return toRemoteObject(tableIdentifier, () -> findRemoteTable(session, tableIdentifier), getCached); + } + + private TableIdentifier findRemoteTable(ConnectorSession session, TableIdentifier tableIdentifier) + { + Namespace remoteNamespace = toRemoteNamespace(session, tableIdentifier.namespace()); + List tableIdentifiers; + try { + tableIdentifiers = restSessionCatalog.listTables(convert(session), remoteNamespace); + } + catch (RESTException e) { + throw new TrinoException(ICEBERG_CATALOG_ERROR, "Failed to list tables", e); + } + TableIdentifier matchingTable = null; + for (TableIdentifier identifier : tableIdentifiers) { + if (identifier.name().equalsIgnoreCase(tableIdentifier.name())) { + if (matchingTable != null) { + throw new TrinoException(NOT_SUPPORTED, "Duplicate table names are not supported with Iceberg REST catalog: " + + Joiner.on(", ").join(matchingTable, identifier.name())); + } + matchingTable = identifier; + } + } + return matchingTable == null ? TableIdentifier.of(remoteNamespace, tableIdentifier.name()) : matchingTable; + } + + private TableIdentifier toRemoteView(ConnectorSession session, SchemaTableName schemaViewName, boolean getCached) + { + TableIdentifier tableIdentifier = toIdentifier(schemaViewName); + return toRemoteObject(tableIdentifier, () -> findRemoteView(session, tableIdentifier), getCached); + } + + private TableIdentifier findRemoteView(ConnectorSession session, TableIdentifier tableIdentifier) + { + Namespace remoteNamespace = toRemoteNamespace(session, tableIdentifier.namespace()); + List tableIdentifiers; + try { + tableIdentifiers = restSessionCatalog.listViews(convert(session), remoteNamespace); + } + catch (RESTException e) { + throw new TrinoException(ICEBERG_CATALOG_ERROR, "Failed to list views", e); + } + TableIdentifier matchingView = null; + for (TableIdentifier identifier : tableIdentifiers) { + if (identifier.name().equalsIgnoreCase(tableIdentifier.name())) { + if (matchingView != null) { + throw new TrinoException(NOT_SUPPORTED, "Duplicate view names are not supported with Iceberg REST catalog: " + + Joiner.on(", ").join(matchingView.name(), identifier.name())); + } + matchingView = identifier; + } + } + return matchingView == null ? TableIdentifier.of(remoteNamespace, tableIdentifier.name()) : matchingView; + } + + private TableIdentifier toRemoteObject(TableIdentifier tableIdentifier, Supplier remoteObjectProvider, boolean getCached) + { + if (caseInsensitiveNameMatching) { + if (getCached) { + return uncheckedCacheGet(remoteTableMappingCache, tableIdentifier, remoteObjectProvider); + } + return remoteObjectProvider.get(); + } + return tableIdentifier; + } + + private Namespace toRemoteNamespace(ConnectorSession session, Namespace trinoNamespace) + { + if (caseInsensitiveNameMatching) { + return uncheckedCacheGet(remoteNamespaceMappingCache, trinoNamespace, () -> findRemoteNamespace(session, trinoNamespace)); + } + return trinoNamespace; + } + + private Namespace findRemoteNamespace(ConnectorSession session, Namespace trinoNamespace) + { + List matchingRemoteNamespaces = listNamespaces(session, Namespace.empty()).stream() + .filter(ns -> toTrinoNamespace(ns).equals(trinoNamespace)) + .collect(toImmutableList()); + if (matchingRemoteNamespaces.size() > 1) { + throw new TrinoException(NOT_SUPPORTED, "Duplicate namespace names are not supported with Iceberg REST catalog: " + matchingRemoteNamespaces); + } + return matchingRemoteNamespaces.isEmpty() ? trinoNamespace : matchingRemoteNamespaces.get(0); + } + + private List listNamespaces(ConnectorSession session, Namespace parentNamespace) + { + List childNamespaces; + try { + childNamespaces = restSessionCatalog.listNamespaces(convert(session), parentNamespace); + } + catch (RESTException e) { + throw new TrinoException(ICEBERG_CATALOG_ERROR, "Failed to list namespaces", e); + } + return childNamespaces.stream().flatMap(childNamespace -> Stream.concat(Stream.of(childNamespace), listNamespaces(session, childNamespace).stream())).toList(); + } + + private static Namespace toTrinoNamespace(Namespace namespace) + { + return Namespace.of(Arrays.stream(namespace.levels()).map(level -> level.toLowerCase(ENGLISH)).toArray(String[]::new)); + } } diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/delete/DeleteFile.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/delete/DeleteFile.java index 17fdcad61e10..b74bd34d0cbc 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/delete/DeleteFile.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/delete/DeleteFile.java @@ -13,135 +13,70 @@ */ package io.trino.plugin.iceberg.delete; -import com.fasterxml.jackson.annotation.JsonCreator; -import com.fasterxml.jackson.annotation.JsonProperty; import com.google.common.collect.ImmutableList; -import com.google.common.collect.ImmutableMap; -import io.airlift.slice.SizeOf; import org.apache.iceberg.FileContent; import org.apache.iceberg.FileFormat; +import org.apache.iceberg.types.Conversions; -import java.nio.ByteBuffer; import java.util.List; -import java.util.Map; import java.util.Optional; -import static com.google.common.base.MoreObjects.firstNonNull; import static com.google.common.base.MoreObjects.toStringHelper; -import static com.google.common.collect.ImmutableMap.toImmutableMap; import static io.airlift.slice.SizeOf.SIZE_OF_INT; import static io.airlift.slice.SizeOf.estimatedSizeOf; import static io.airlift.slice.SizeOf.instanceSize; -import static io.trino.plugin.base.io.ByteBuffers.getWrappedBytes; import static java.util.Objects.requireNonNull; - -public final class DeleteFile +import static org.apache.iceberg.MetadataColumns.DELETE_FILE_POS; + +public record DeleteFile( + FileContent content, + String path, + FileFormat format, + long recordCount, + long fileSizeInBytes, + List equalityFieldIds, + Optional rowPositionLowerBound, + Optional rowPositionUpperBound, + long dataSequenceNumber) { private static final long INSTANCE_SIZE = instanceSize(DeleteFile.class); - private final FileContent content; - private final String path; - private final FileFormat format; - private final long recordCount; - private final long fileSizeInBytes; - private final List equalityFieldIds; - private final Map lowerBounds; - private final Map upperBounds; - public static DeleteFile fromIceberg(org.apache.iceberg.DeleteFile deleteFile) { - Map lowerBounds = firstNonNull(deleteFile.lowerBounds(), ImmutableMap.of()) - .entrySet().stream().collect(toImmutableMap(Map.Entry::getKey, entry -> getWrappedBytes(entry.getValue()).clone())); - Map upperBounds = firstNonNull(deleteFile.upperBounds(), ImmutableMap.of()) - .entrySet().stream().collect(toImmutableMap(Map.Entry::getKey, entry -> getWrappedBytes(entry.getValue()).clone())); + Optional rowPositionLowerBound = Optional.ofNullable(deleteFile.lowerBounds()) + .map(bounds -> bounds.get(DELETE_FILE_POS.fieldId())) + .map(bytes -> Conversions.fromByteBuffer(DELETE_FILE_POS.type(), bytes)); + Optional rowPositionUpperBound = Optional.ofNullable(deleteFile.upperBounds()) + .map(bounds -> bounds.get(DELETE_FILE_POS.fieldId())) + .map(bytes -> Conversions.fromByteBuffer(DELETE_FILE_POS.type(), bytes)); return new DeleteFile( deleteFile.content(), - deleteFile.path().toString(), + deleteFile.location(), deleteFile.format(), deleteFile.recordCount(), deleteFile.fileSizeInBytes(), Optional.ofNullable(deleteFile.equalityFieldIds()).orElseGet(ImmutableList::of), - lowerBounds, - upperBounds); - } - - @JsonCreator - public DeleteFile( - FileContent content, - String path, - FileFormat format, - long recordCount, - long fileSizeInBytes, - List equalityFieldIds, - Map lowerBounds, - Map upperBounds) - { - this.content = requireNonNull(content, "content is null"); - this.path = requireNonNull(path, "path is null"); - this.format = requireNonNull(format, "format is null"); - this.recordCount = recordCount; - this.fileSizeInBytes = fileSizeInBytes; - this.equalityFieldIds = ImmutableList.copyOf(requireNonNull(equalityFieldIds, "equalityFieldIds is null")); - this.lowerBounds = ImmutableMap.copyOf(requireNonNull(lowerBounds, "lowerBounds is null")); - this.upperBounds = ImmutableMap.copyOf(requireNonNull(upperBounds, "upperBounds is null")); - } - - @JsonProperty - public FileContent content() - { - return content; - } - - @JsonProperty - public String path() - { - return path; - } - - @JsonProperty - public FileFormat format() - { - return format; - } - - @JsonProperty - public long recordCount() - { - return recordCount; - } - - @JsonProperty - public long fileSizeInBytes() - { - return fileSizeInBytes; - } - - @JsonProperty - public List equalityFieldIds() - { - return equalityFieldIds; - } - - @JsonProperty - public Map getLowerBounds() - { - return lowerBounds; + rowPositionLowerBound, + rowPositionUpperBound, + deleteFile.dataSequenceNumber()); } - @JsonProperty - public Map getUpperBounds() + public DeleteFile { - return upperBounds; + requireNonNull(content, "content is null"); + requireNonNull(path, "path is null"); + requireNonNull(format, "format is null"); + equalityFieldIds = ImmutableList.copyOf(requireNonNull(equalityFieldIds, "equalityFieldIds is null")); + requireNonNull(rowPositionLowerBound, "rowPositionLowerBound is null"); + requireNonNull(rowPositionUpperBound, "rowPositionUpperBound is null"); } - public long getRetainedSizeInBytes() + public long retainedSizeInBytes() { return INSTANCE_SIZE + estimatedSizeOf(path) - + estimatedSizeOf(equalityFieldIds, ignored -> SIZE_OF_INT) - + estimatedSizeOf(lowerBounds, entry -> SIZE_OF_INT, SizeOf::sizeOf) - + estimatedSizeOf(upperBounds, entry -> SIZE_OF_INT, SizeOf::sizeOf); + + estimatedSizeOf(equalityFieldIds, ignore -> SIZE_OF_INT); } @Override diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/delete/DeleteFilter.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/delete/DeleteFilter.java index b061ec276efc..68ba48f49848 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/delete/DeleteFilter.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/delete/DeleteFilter.java @@ -19,5 +19,5 @@ public interface DeleteFilter { - RowPredicate createPredicate(List columns); + RowPredicate createPredicate(List columns, long dataSequenceNumber); } diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/delete/DeleteManager.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/delete/DeleteManager.java new file mode 100644 index 000000000000..5dc29b8482eb --- /dev/null +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/delete/DeleteManager.java @@ -0,0 +1,214 @@ +/* + * 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.iceberg.delete; + +import com.google.common.base.VerifyException; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.common.util.concurrent.Futures; +import com.google.common.util.concurrent.ListenableFuture; +import io.airlift.slice.Slice; +import io.trino.plugin.iceberg.IcebergColumnHandle; +import io.trino.plugin.iceberg.IcebergPageSourceProvider.ReaderPageSourceWithRowPositions; +import io.trino.plugin.iceberg.delete.EqualityDeleteFilter.EqualityDeleteFilterBuilder; +import io.trino.spi.TrinoException; +import io.trino.spi.connector.ConnectorPageSource; +import io.trino.spi.predicate.Domain; +import io.trino.spi.predicate.NullableValue; +import io.trino.spi.predicate.Range; +import io.trino.spi.predicate.TupleDomain; +import io.trino.spi.predicate.ValueSet; +import io.trino.spi.type.TypeManager; +import org.apache.iceberg.Schema; +import org.roaringbitmap.longlong.LongBitmapDataProvider; +import org.roaringbitmap.longlong.Roaring64Bitmap; + +import java.io.IOException; +import java.io.UncheckedIOException; +import java.util.ArrayList; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ExecutionException; + +import static com.google.common.base.Verify.verify; +import static com.google.common.collect.ImmutableList.toImmutableList; +import static io.airlift.slice.Slices.utf8Slice; +import static io.trino.plugin.iceberg.IcebergErrorCode.ICEBERG_BAD_DATA; +import static io.trino.plugin.iceberg.IcebergUtil.getColumnHandle; +import static io.trino.plugin.iceberg.IcebergUtil.schemaFromHandles; +import static io.trino.plugin.iceberg.delete.PositionDeleteFilter.readPositionDeletes; +import static io.trino.spi.type.VarcharType.VARCHAR; +import static java.util.Objects.requireNonNull; +import static org.apache.iceberg.MetadataColumns.DELETE_FILE_PATH; +import static org.apache.iceberg.MetadataColumns.DELETE_FILE_POS; + +public class DeleteManager +{ + private final TypeManager typeManager; + private final Map, EqualityDeleteFilterBuilder> equalityDeleteFiltersBySchema = new ConcurrentHashMap<>(); + + public DeleteManager(TypeManager typeManager) + { + this.typeManager = requireNonNull(typeManager, "typeManager is null"); + } + + public Optional getDeletePredicate( + String dataFilePath, + long dataSequenceNumber, + List deleteFiles, + List readColumns, + Schema tableSchema, + ReaderPageSourceWithRowPositions readerPageSourceWithRowPositions, + DeletePageSourceProvider deletePageSourceProvider) + { + if (deleteFiles.isEmpty()) { + return Optional.empty(); + } + + List positionDeleteFiles = new ArrayList<>(); + List equalityDeleteFiles = new ArrayList<>(); + for (DeleteFile deleteFile : deleteFiles) { + switch (deleteFile.content()) { + case POSITION_DELETES -> positionDeleteFiles.add(deleteFile); + case EQUALITY_DELETES -> equalityDeleteFiles.add(deleteFile); + case DATA -> throw new VerifyException("DATA is not delete file type"); + } + } + + Optional positionDeletes = createPositionDeleteFilter(dataFilePath, positionDeleteFiles, readerPageSourceWithRowPositions, deletePageSourceProvider) + .map(filter -> filter.createPredicate(readColumns, dataSequenceNumber)); + Optional equalityDeletes = createEqualityDeleteFilter(equalityDeleteFiles, tableSchema, deletePageSourceProvider).stream() + .map(filter -> filter.createPredicate(readColumns, dataSequenceNumber)) + .reduce(RowPredicate::and); + + if (positionDeletes.isEmpty()) { + return equalityDeletes; + } + return equalityDeletes + .map(rowPredicate -> positionDeletes.get().and(rowPredicate)) + .or(() -> positionDeletes); + } + + public interface DeletePageSourceProvider + { + ConnectorPageSource openDeletes( + DeleteFile delete, + List deleteColumns, + TupleDomain tupleDomain); + } + + private Optional createPositionDeleteFilter( + String dataFilePath, + List positionDeleteFiles, + ReaderPageSourceWithRowPositions readerPageSourceWithRowPositions, + DeletePageSourceProvider deletePageSourceProvider) + { + if (positionDeleteFiles.isEmpty()) { + return Optional.empty(); + } + + Slice targetPath = utf8Slice(dataFilePath); + + Optional startRowPosition = readerPageSourceWithRowPositions.startRowPosition(); + Optional endRowPosition = readerPageSourceWithRowPositions.endRowPosition(); + verify(startRowPosition.isPresent() == endRowPosition.isPresent(), "startRowPosition and endRowPosition must be specified together"); + IcebergColumnHandle deleteFilePath = getColumnHandle(DELETE_FILE_PATH, typeManager); + IcebergColumnHandle deleteFilePos = getColumnHandle(DELETE_FILE_POS, typeManager); + List deleteColumns = ImmutableList.of(deleteFilePath, deleteFilePos); + TupleDomain deleteDomain = TupleDomain.fromFixedValues(ImmutableMap.of(deleteFilePath, NullableValue.of(VARCHAR, targetPath))); + if (startRowPosition.isPresent()) { + Range positionRange = Range.range(deleteFilePos.getType(), startRowPosition.get(), true, endRowPosition.get(), true); + TupleDomain positionDomain = TupleDomain.withColumnDomains(ImmutableMap.of(deleteFilePos, Domain.create(ValueSet.ofRanges(positionRange), false))); + deleteDomain = deleteDomain.intersect(positionDomain); + } + + LongBitmapDataProvider deletedRows = new Roaring64Bitmap(); + for (DeleteFile deleteFile : positionDeleteFiles) { + if (shouldLoadPositionDeleteFile(deleteFile, startRowPosition, endRowPosition)) { + try (ConnectorPageSource pageSource = deletePageSourceProvider.openDeletes(deleteFile, deleteColumns, deleteDomain)) { + readPositionDeletes(pageSource, targetPath, deletedRows); + } + catch (IOException e) { + throw new UncheckedIOException(e); + } + } + } + + if (deletedRows.isEmpty()) { + return Optional.empty(); + } + return Optional.of(new PositionDeleteFilter(deletedRows)); + } + + private static boolean shouldLoadPositionDeleteFile(DeleteFile deleteFile, Optional startRowPosition, Optional endRowPosition) + { + if (startRowPosition.isEmpty()) { + return true; + } + + Optional positionLowerBound = deleteFile.rowPositionLowerBound(); + Optional positionUpperBound = deleteFile.rowPositionUpperBound(); + return (positionLowerBound.isEmpty() || positionLowerBound.get() <= endRowPosition.orElseThrow()) && + (positionUpperBound.isEmpty() || positionUpperBound.get() >= startRowPosition.get()); + } + + private List createEqualityDeleteFilter(List equalityDeleteFiles, Schema schema, DeletePageSourceProvider deletePageSourceProvider) + { + if (equalityDeleteFiles.isEmpty()) { + return List.of(); + } + + // The equality delete files can be loaded in parallel. There may be multiple split threads attempting to load the + // same files. The current thread will only load a file if it is not already being loaded by another thread. + List> pendingLoads = new ArrayList<>(); + Set deleteFilters = new HashSet<>(); + for (DeleteFile deleteFile : equalityDeleteFiles) { + List fieldIds = deleteFile.equalityFieldIds(); + verify(!fieldIds.isEmpty(), "equality field IDs are missing"); + List deleteColumns = fieldIds.stream() + .map(id -> getColumnHandle(schema.findField(id), typeManager)) + .collect(toImmutableList()); + + // each file can have a different set of columns for the equality delete, so we need to create a new builder for each set of columns + EqualityDeleteFilterBuilder builder = equalityDeleteFiltersBySchema.computeIfAbsent(fieldIds, ignore -> EqualityDeleteFilter.builder(schemaFromHandles(deleteColumns))); + deleteFilters.add(builder); + + ListenableFuture loadFuture = builder.readEqualityDeletes(deleteFile, deleteColumns, deletePageSourceProvider); + if (!loadFuture.isDone()) { + pendingLoads.add(loadFuture); + } + } + + // Wait loads happening in other threads + try { + Futures.allAsList(pendingLoads).get(); + } + catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException(e); + } + catch (ExecutionException e) { + // Since execution can happen on another thread, it is not safe to unwrap the exception + throw new TrinoException(ICEBERG_BAD_DATA, "Failed to load equality deletes", e); + } + + return deleteFilters.stream() + .map(EqualityDeleteFilterBuilder::build) + .toList(); + } +} diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/delete/EqualityDeleteFilter.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/delete/EqualityDeleteFilter.java index a41fbc679479..c027fbf43fc9 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/delete/EqualityDeleteFilter.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/delete/EqualityDeleteFilter.java @@ -13,68 +13,139 @@ */ package io.trino.plugin.iceberg.delete; +import com.google.common.util.concurrent.Futures; +import com.google.common.util.concurrent.ListenableFuture; +import com.google.common.util.concurrent.ListenableFutureTask; +import com.google.errorprone.annotations.ThreadSafe; import io.trino.plugin.iceberg.IcebergColumnHandle; +import io.trino.plugin.iceberg.delete.DeleteManager.DeletePageSourceProvider; import io.trino.spi.Page; +import io.trino.spi.TrinoException; import io.trino.spi.connector.ConnectorPageSource; +import io.trino.spi.predicate.TupleDomain; import io.trino.spi.type.Type; import org.apache.iceberg.Schema; -import org.apache.iceberg.StructLike; -import org.apache.iceberg.util.StructLikeSet; +import org.apache.iceberg.types.Types.StructType; +import org.apache.iceberg.util.StructLikeWrapper; import org.apache.iceberg.util.StructProjection; +import java.io.IOException; +import java.io.UncheckedIOException; import java.util.List; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; -import static io.trino.plugin.iceberg.IcebergUtil.schemaFromHandles; +import static com.google.common.base.Verify.verify; +import static io.trino.plugin.iceberg.IcebergErrorCode.ICEBERG_CANNOT_OPEN_SPLIT; +import static io.trino.plugin.iceberg.IcebergUtil.structTypeFromHandles; import static java.util.Objects.requireNonNull; public final class EqualityDeleteFilter implements DeleteFilter { - private final Schema schema; - private final StructLikeSet deleteSet; + private final Schema deleteSchema; + private final Map deletedRows; - private EqualityDeleteFilter(Schema schema, StructLikeSet deleteSet) + private EqualityDeleteFilter(Schema deleteSchema, Map deletedRows) { - this.schema = requireNonNull(schema, "schema is null"); - this.deleteSet = requireNonNull(deleteSet, "deleteSet is null"); + this.deleteSchema = requireNonNull(deleteSchema, "deleteSchema is null"); + this.deletedRows = requireNonNull(deletedRows, "deletedRows is null"); } @Override - public RowPredicate createPredicate(List columns) + public RowPredicate createPredicate(List columns, long splitDataSequenceNumber) { + StructType fileStructType = structTypeFromHandles(columns); + StructType deleteStructType = deleteSchema.asStruct(); + if (deleteSchema.columns().stream().anyMatch(column -> fileStructType.field(column.fieldId()) == null)) { + throw new TrinoException(ICEBERG_CANNOT_OPEN_SPLIT, "columns list doesn't contain all equality delete columns"); + } + + StructLikeWrapper structLikeWrapper = StructLikeWrapper.forType(deleteStructType); + StructProjection projection = StructProjection.create(fileStructType, deleteStructType); Type[] types = columns.stream() .map(IcebergColumnHandle::getType) .toArray(Type[]::new); - Schema fileSchema = schemaFromHandles(columns); - StructProjection projection = StructProjection.create(fileSchema, schema); - return (page, position) -> { - StructLike row = new LazyTrinoRow(types, page, position); - return !deleteSet.contains(projection.wrap(row)); + StructProjection row = projection.wrap(new LazyTrinoRow(types, page, position)); + DataSequenceNumber maxDeleteVersion = deletedRows.get(structLikeWrapper.set(row)); + // clear reference to avoid memory leak + structLikeWrapper.set(null); + return maxDeleteVersion == null || maxDeleteVersion.dataSequenceNumber() <= splitDataSequenceNumber; }; } - public static DeleteFilter readEqualityDeletes(ConnectorPageSource pageSource, List columns, Schema tableSchema) + public static EqualityDeleteFilterBuilder builder(Schema deleteSchema) { - Type[] types = columns.stream() - .map(IcebergColumnHandle::getType) - .toArray(Type[]::new); + return new EqualityDeleteFilterBuilder(deleteSchema); + } - Schema deleteSchema = schemaFromHandles(columns); - StructLikeSet deleteSet = StructLikeSet.create(deleteSchema.asStruct()); + @ThreadSafe + public static class EqualityDeleteFilterBuilder + { + private final Schema deleteSchema; + private final Map deletedRows; + private final Map> loadingFiles = new ConcurrentHashMap<>(); - while (!pageSource.isFinished()) { - Page page = pageSource.getNextPage(); - if (page == null) { - continue; - } + private EqualityDeleteFilterBuilder(Schema deleteSchema) + { + this.deleteSchema = requireNonNull(deleteSchema, "deleteSchema is null"); + this.deletedRows = new ConcurrentHashMap<>(); + } + + public ListenableFuture readEqualityDeletes(DeleteFile deleteFile, List deleteColumns, DeletePageSourceProvider deletePageSourceProvider) + { + verify(deleteColumns.size() == deleteSchema.columns().size(), "delete columns size doesn't match delete schema size"); + + // ensure only one thread loads the file + ListenableFutureTask futureTask = loadingFiles.computeIfAbsent( + deleteFile.path(), + key -> ListenableFutureTask.create(() -> readEqualityDeletesInternal(deleteFile, deleteColumns, deletePageSourceProvider), null)); + futureTask.run(); + return Futures.nonCancellationPropagating(futureTask); + } + + private void readEqualityDeletesInternal(DeleteFile deleteFile, List deleteColumns, DeletePageSourceProvider deletePageSourceProvider) + { + DataSequenceNumber sequenceNumber = new DataSequenceNumber(deleteFile.dataSequenceNumber()); + try (ConnectorPageSource pageSource = deletePageSourceProvider.openDeletes(deleteFile, deleteColumns, TupleDomain.all())) { + Type[] types = deleteColumns.stream() + .map(IcebergColumnHandle::getType) + .toArray(Type[]::new); - for (int position = 0; position < page.getPositionCount(); position++) { - deleteSet.add(new TrinoRow(types, page, position)); + StructLikeWrapper wrapper = StructLikeWrapper.forType(deleteSchema.asStruct()); + while (!pageSource.isFinished()) { + Page page = pageSource.getNextPage(); + if (page == null) { + continue; + } + + for (int position = 0; position < page.getPositionCount(); position++) { + TrinoRow row = new TrinoRow(types, page, position); + deletedRows.merge(wrapper.copyFor(row), sequenceNumber, (existing, newValue) -> { + if (existing.dataSequenceNumber() > newValue.dataSequenceNumber()) { + return existing; + } + return newValue; + }); + } + } + } + catch (IOException e) { + throw new UncheckedIOException(e); } } - return new EqualityDeleteFilter(deleteSchema, deleteSet); + /** + * Builds the EqualityDeleteFilter. + * After building the EqualityDeleteFilter, additional rows can be added to this builder, and the filter can be rebuilt. + */ + public EqualityDeleteFilter build() + { + return new EqualityDeleteFilter(deleteSchema, deletedRows); + } } + + private record DataSequenceNumber(long dataSequenceNumber) {} } diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/delete/PositionDeleteFilter.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/delete/PositionDeleteFilter.java index 952f215fb16b..d45e807e5aa8 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/delete/PositionDeleteFilter.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/delete/PositionDeleteFilter.java @@ -39,7 +39,7 @@ public PositionDeleteFilter(ImmutableLongBitmapDataProvider deletedRows) } @Override - public RowPredicate createPredicate(List columns) + public RowPredicate createPredicate(List columns, long dataSequenceNumber) { int filePosChannel = rowPositionChannel(columns); return (page, position) -> { diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/delete/IcebergPositionDeletePageSink.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/delete/PositionDeleteWriter.java similarity index 57% rename from plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/delete/IcebergPositionDeletePageSink.java rename to plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/delete/PositionDeleteWriter.java index 0eb7b80318f5..1f6bc8f71543 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/delete/IcebergPositionDeletePageSink.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/delete/PositionDeleteWriter.java @@ -24,34 +24,33 @@ import io.trino.plugin.iceberg.MetricsWrapper; import io.trino.plugin.iceberg.PartitionData; import io.trino.spi.Page; +import io.trino.spi.PageBuilder; import io.trino.spi.block.Block; import io.trino.spi.block.RunLengthEncodedBlock; -import io.trino.spi.connector.ConnectorPageSink; import io.trino.spi.connector.ConnectorSession; import org.apache.iceberg.FileContent; import org.apache.iceberg.PartitionSpec; import org.apache.iceberg.PartitionSpecParser; import org.apache.iceberg.io.LocationProvider; +import org.roaringbitmap.longlong.ImmutableLongBitmapDataProvider; -import java.util.ArrayList; import java.util.Collection; +import java.util.List; import java.util.Map; import java.util.Optional; -import java.util.concurrent.CompletableFuture; -import static com.google.common.base.Preconditions.checkArgument; import static io.airlift.slice.Slices.utf8Slice; import static io.airlift.slice.Slices.wrappedBuffer; import static io.trino.spi.predicate.Utils.nativeValueToBlock; +import static io.trino.spi.type.BigintType.BIGINT; import static io.trino.spi.type.VarcharType.VARCHAR; import static java.util.Objects.requireNonNull; import static java.util.UUID.randomUUID; -import static java.util.concurrent.CompletableFuture.completedFuture; -public class IcebergPositionDeletePageSink - implements ConnectorPageSink +public class PositionDeleteWriter { private final String dataFilePath; + private final Block dataFilePathBlock; private final PartitionSpec partitionSpec; private final Optional partition; private final String outputPath; @@ -59,10 +58,7 @@ public class IcebergPositionDeletePageSink private final IcebergFileWriter writer; private final IcebergFileFormat fileFormat; - private long validationCpuNanos; - private boolean writtenData; - - public IcebergPositionDeletePageSink( + public PositionDeleteWriter( String dataFilePath, PartitionSpec partitionSpec, Optional partition, @@ -75,12 +71,14 @@ public IcebergPositionDeletePageSink( Map storageProperties) { this.dataFilePath = requireNonNull(dataFilePath, "dataFilePath is null"); + this.dataFilePathBlock = nativeValueToBlock(VARCHAR, utf8Slice(dataFilePath)); this.jsonCodec = requireNonNull(jsonCodec, "jsonCodec is null"); this.partitionSpec = requireNonNull(partitionSpec, "partitionSpec is null"); this.partition = requireNonNull(partition, "partition is null"); this.fileFormat = requireNonNull(fileFormat, "fileFormat is null"); - // prepend query id to a file name so we can determine which files were written by which query. This is needed for opportunistic cleanup of extra files - // which may be present for successfully completing query in presence of failure recovery mechanisms. + // Prepend query ID to the file name, allowing us to determine the files written by a query. + // This is necessary for opportunistic cleanup of extra files, which may be present for + // successfully completed queries in the presence of failure recovery mechanisms. String fileName = fileFormat.toIceberg().addExtension(session.getQueryId() + "-" + randomUUID()); this.outputPath = partition .map(partitionData -> locationProvider.newDataLocation(partitionSpec, partitionData, fileName)) @@ -88,69 +86,52 @@ public IcebergPositionDeletePageSink( this.writer = fileWriterFactory.createPositionDeleteWriter(fileSystem, Location.of(outputPath), session, fileFormat, storageProperties); } - @Override - public long getCompletedBytes() + public Collection write(ImmutableLongBitmapDataProvider rowsToDelete) { - return writer.getWrittenBytes(); - } + writeDeletes(rowsToDelete); + writer.commit(); - @Override - public long getMemoryUsage() - { - return writer.getMemoryUsage(); - } + CommitTaskData task = new CommitTaskData( + outputPath, + fileFormat, + writer.getWrittenBytes(), + new MetricsWrapper(writer.getFileMetrics().metrics()), + PartitionSpecParser.toJson(partitionSpec), + partition.map(PartitionData::toJson), + FileContent.POSITION_DELETES, + Optional.of(dataFilePath), + writer.getFileMetrics().splitOffsets()); - @Override - public long getValidationCpuNanos() - { - return validationCpuNanos; + return List.of(wrappedBuffer(jsonCodec.toJsonBytes(task))); } - @Override - public CompletableFuture appendPage(Page page) + public void abort() { - checkArgument(page.getChannelCount() == 1, "IcebergPositionDeletePageSink expected a Page with only one channel, but got " + page.getChannelCount()); - - Block[] blocks = new Block[2]; - blocks[0] = RunLengthEncodedBlock.create(nativeValueToBlock(VARCHAR, utf8Slice(dataFilePath)), page.getPositionCount()); - blocks[1] = page.getBlock(0); - writer.appendRows(new Page(blocks)); - - writtenData = true; - return NOT_BLOCKED; + writer.rollback(); } - @Override - public CompletableFuture> finish() + private void writeDeletes(ImmutableLongBitmapDataProvider rowsToDelete) { - Collection commitTasks = new ArrayList<>(); - if (writtenData) { - writer.commit(); - CommitTaskData task = new CommitTaskData( - outputPath, - fileFormat, - writer.getWrittenBytes(), - new MetricsWrapper(writer.getMetrics()), - PartitionSpecParser.toJson(partitionSpec), - partition.map(PartitionData::toJson), - FileContent.POSITION_DELETES, - Optional.of(dataFilePath)); - Long recordCount = task.getMetrics().recordCount(); - if (recordCount != null && recordCount > 0) { - commitTasks.add(wrappedBuffer(jsonCodec.toJsonBytes(task))); + PageBuilder pageBuilder = new PageBuilder(List.of(BIGINT)); + + rowsToDelete.forEach(rowPosition -> { + pageBuilder.declarePosition(); + BIGINT.writeLong(pageBuilder.getBlockBuilder(0), rowPosition); + if (pageBuilder.isFull()) { + writePage(pageBuilder.build()); + pageBuilder.reset(); } - validationCpuNanos = writer.getValidationCpuNanos(); - } - else { - // clean up the empty delete file - writer.rollback(); + }); + + if (!pageBuilder.isEmpty()) { + writePage(pageBuilder.build()); } - return completedFuture(commitTasks); } - @Override - public void abort() + private void writePage(Page page) { - writer.rollback(); + writer.appendRows(new Page( + RunLengthEncodedBlock.create(dataFilePathBlock, page.getPositionCount()), + page.getBlock(0))); } } diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/fileio/ForwardingFileIo.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/fileio/ForwardingFileIo.java index 03dcb9d109d9..cfb511c3fccb 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/fileio/ForwardingFileIo.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/fileio/ForwardingFileIo.java @@ -13,6 +13,7 @@ */ package io.trino.plugin.iceberg.fileio; +import com.google.common.collect.ImmutableMap; import com.google.common.collect.Iterables; import io.trino.filesystem.Location; import io.trino.filesystem.TrinoFileSystem; @@ -24,6 +25,7 @@ import java.io.IOException; import java.io.UncheckedIOException; import java.util.List; +import java.util.Map; import java.util.stream.Stream; import static java.util.Objects.requireNonNull; @@ -36,10 +38,17 @@ public class ForwardingFileIo private static final int BATCH_DELETE_PATHS_MESSAGE_LIMIT = 5; private final TrinoFileSystem fileSystem; + private final Map properties; public ForwardingFileIo(TrinoFileSystem fileSystem) + { + this(fileSystem, ImmutableMap.of()); + } + + public ForwardingFileIo(TrinoFileSystem fileSystem, Map properties) { this.fileSystem = requireNonNull(fileSystem, "fileSystem is null"); + this.properties = ImmutableMap.copyOf(requireNonNull(properties, "properties is null")); } @Override @@ -57,7 +66,7 @@ public InputFile newInputFile(String path, long length) @Override public OutputFile newOutputFile(String path) { - return new ForwardingOutputFile(fileSystem, path); + return new ForwardingOutputFile(fileSystem, Location.of(path)); } @Override @@ -95,4 +104,16 @@ private void deleteBatch(List filesToDelete) e); } } + + @Override + public Map properties() + { + return properties; + } + + @Override + public void initialize(Map properties) + { + throw new UnsupportedOperationException("ForwardingFileIO does not support initialization by properties"); + } } diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/fileio/ForwardingOutputFile.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/fileio/ForwardingOutputFile.java index 40a65d5b7a36..8d0f956688b7 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/fileio/ForwardingOutputFile.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/fileio/ForwardingOutputFile.java @@ -33,10 +33,10 @@ public class ForwardingOutputFile private final TrinoFileSystem fileSystem; private final TrinoOutputFile outputFile; - public ForwardingOutputFile(TrinoFileSystem fileSystem, String path) + public ForwardingOutputFile(TrinoFileSystem fileSystem, Location location) { this.fileSystem = requireNonNull(fileSystem, "fileSystem is null"); - this.outputFile = fileSystem.newOutputFile(Location.of(path)); + this.outputFile = fileSystem.newOutputFile(location); } @Override diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/functions/IcebergFunctionProvider.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/functions/IcebergFunctionProvider.java new file mode 100644 index 000000000000..57bbea9479df --- /dev/null +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/functions/IcebergFunctionProvider.java @@ -0,0 +1,45 @@ +/* + * 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.iceberg.functions; + +import com.google.inject.Inject; +import io.trino.plugin.iceberg.functions.tablechanges.TableChangesFunctionHandle; +import io.trino.plugin.iceberg.functions.tablechanges.TableChangesFunctionProcessorProvider; +import io.trino.spi.function.FunctionProvider; +import io.trino.spi.function.table.ConnectorTableFunctionHandle; +import io.trino.spi.function.table.TableFunctionProcessorProvider; + +import static java.util.Objects.requireNonNull; + +public class IcebergFunctionProvider + implements FunctionProvider +{ + private final TableChangesFunctionProcessorProvider tableChangesFunctionProcessorProvider; + + @Inject + public IcebergFunctionProvider(TableChangesFunctionProcessorProvider tableChangesFunctionProcessorProvider) + { + this.tableChangesFunctionProcessorProvider = requireNonNull(tableChangesFunctionProcessorProvider, "tableChangesFunctionProcessorProvider is null"); + } + + @Override + public TableFunctionProcessorProvider getTableFunctionProcessorProvider(ConnectorTableFunctionHandle functionHandle) + { + if (functionHandle instanceof TableChangesFunctionHandle) { + return tableChangesFunctionProcessorProvider; + } + + throw new UnsupportedOperationException("Unsupported function: " + functionHandle); + } +} diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/functions/tablechanges/TableChangesFunction.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/functions/tablechanges/TableChangesFunction.java new file mode 100644 index 000000000000..989c7ada0a66 --- /dev/null +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/functions/tablechanges/TableChangesFunction.java @@ -0,0 +1,192 @@ +/* + * 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.iceberg.functions.tablechanges; + +import com.google.common.collect.ImmutableList; +import com.google.inject.Inject; +import io.airlift.slice.Slice; +import io.trino.plugin.iceberg.ColumnIdentity; +import io.trino.plugin.iceberg.IcebergColumnHandle; +import io.trino.plugin.iceberg.IcebergUtil; +import io.trino.plugin.iceberg.catalog.TrinoCatalogFactory; +import io.trino.spi.TrinoException; +import io.trino.spi.connector.ConnectorAccessControl; +import io.trino.spi.connector.ConnectorSession; +import io.trino.spi.connector.ConnectorTransactionHandle; +import io.trino.spi.connector.SchemaTableName; +import io.trino.spi.function.table.AbstractConnectorTableFunction; +import io.trino.spi.function.table.Argument; +import io.trino.spi.function.table.Descriptor; +import io.trino.spi.function.table.ScalarArgument; +import io.trino.spi.function.table.ScalarArgumentSpecification; +import io.trino.spi.function.table.TableFunctionAnalysis; +import io.trino.spi.type.TypeManager; +import io.trino.spi.type.VarcharType; +import org.apache.iceberg.Schema; +import org.apache.iceberg.SchemaParser; +import org.apache.iceberg.Table; +import org.apache.iceberg.TableProperties; + +import java.util.List; +import java.util.Map; +import java.util.Optional; + +import static com.google.common.collect.ImmutableSet.toImmutableSet; +import static io.trino.plugin.iceberg.ColumnIdentity.TypeCategory.PRIMITIVE; +import static io.trino.plugin.iceberg.IcebergColumnHandle.DATA_CHANGE_ORDINAL_ID; +import static io.trino.plugin.iceberg.IcebergColumnHandle.DATA_CHANGE_ORDINAL_NAME; +import static io.trino.plugin.iceberg.IcebergColumnHandle.DATA_CHANGE_TIMESTAMP_ID; +import static io.trino.plugin.iceberg.IcebergColumnHandle.DATA_CHANGE_TIMESTAMP_NAME; +import static io.trino.plugin.iceberg.IcebergColumnHandle.DATA_CHANGE_TYPE_ID; +import static io.trino.plugin.iceberg.IcebergColumnHandle.DATA_CHANGE_TYPE_NAME; +import static io.trino.plugin.iceberg.IcebergColumnHandle.DATA_CHANGE_VERSION_ID; +import static io.trino.plugin.iceberg.IcebergColumnHandle.DATA_CHANGE_VERSION_NAME; +import static io.trino.plugin.iceberg.TypeConverter.toTrinoType; +import static io.trino.spi.StandardErrorCode.INVALID_FUNCTION_ARGUMENT; +import static io.trino.spi.function.table.ReturnTypeSpecification.GenericTable.GENERIC_TABLE; +import static io.trino.spi.type.BigintType.BIGINT; +import static io.trino.spi.type.IntegerType.INTEGER; +import static io.trino.spi.type.TimestampWithTimeZoneType.TIMESTAMP_TZ_MILLIS; +import static java.util.Objects.requireNonNull; + +public class TableChangesFunction + extends AbstractConnectorTableFunction +{ + private static final String FUNCTION_NAME = "table_changes"; + private static final String SCHEMA_VAR_NAME = "SCHEMA"; + private static final String TABLE_VAR_NAME = "TABLE"; + private static final String START_SNAPSHOT_VAR_NAME = "START_SNAPSHOT_ID"; + private static final String END_SNAPSHOT_VAR_NAME = "END_SNAPSHOT_ID"; + + private final TrinoCatalogFactory trinoCatalogFactory; + private final TypeManager typeManager; + + @Inject + public TableChangesFunction(TrinoCatalogFactory trinoCatalogFactory, TypeManager typeManager) + { + super( + "system", + FUNCTION_NAME, + ImmutableList.of( + ScalarArgumentSpecification.builder() + .name(SCHEMA_VAR_NAME) + .type(VarcharType.createUnboundedVarcharType()) + .build(), + ScalarArgumentSpecification.builder() + .name(TABLE_VAR_NAME) + .type(VarcharType.createUnboundedVarcharType()) + .build(), + ScalarArgumentSpecification.builder() + .name(START_SNAPSHOT_VAR_NAME) + .type(BIGINT) + .build(), + ScalarArgumentSpecification.builder() + .name(END_SNAPSHOT_VAR_NAME) + .type(BIGINT) + .build()), + GENERIC_TABLE); + + this.trinoCatalogFactory = requireNonNull(trinoCatalogFactory, "trinoCatalogFactory is null"); + this.typeManager = requireNonNull(typeManager, "typeManager is null"); + } + + @Override + public TableFunctionAnalysis analyze(ConnectorSession session, ConnectorTransactionHandle transaction, Map arguments, ConnectorAccessControl accessControl) + { + String schema = ((Slice) checkNonNull(((ScalarArgument) arguments.get(SCHEMA_VAR_NAME)).getValue())).toStringUtf8(); + String table = ((Slice) checkNonNull(((ScalarArgument) arguments.get(TABLE_VAR_NAME)).getValue())).toStringUtf8(); + long startSnapshotId = (long) checkNonNull(((ScalarArgument) arguments.get(START_SNAPSHOT_VAR_NAME)).getValue()); + long endSnapshotId = (long) checkNonNull(((ScalarArgument) arguments.get(END_SNAPSHOT_VAR_NAME)).getValue()); + + SchemaTableName schemaTableName = new SchemaTableName(schema, table); + Table icebergTable = trinoCatalogFactory.create(session.getIdentity()) + .loadTable(session, schemaTableName); + + checkSnapshotExists(icebergTable, startSnapshotId); + checkSnapshotExists(icebergTable, endSnapshotId); + + ImmutableList.Builder columns = ImmutableList.builder(); + Schema tableSchema = icebergTable.schemas().get(icebergTable.snapshot(endSnapshotId).schemaId()); + tableSchema.columns().stream() + .map(column -> new Descriptor.Field(column.name(), Optional.of(toTrinoType(column.type(), typeManager)))) + .forEach(columns::add); + + columns.add(new Descriptor.Field(DATA_CHANGE_TYPE_NAME, Optional.of(VarcharType.createUnboundedVarcharType()))); + columns.add(new Descriptor.Field(DATA_CHANGE_VERSION_NAME, Optional.of(BIGINT))); + columns.add(new Descriptor.Field(DATA_CHANGE_TIMESTAMP_NAME, Optional.of(TIMESTAMP_TZ_MILLIS))); + columns.add(new Descriptor.Field(DATA_CHANGE_ORDINAL_NAME, Optional.of(INTEGER))); + + ImmutableList.Builder columnHandlesBuilder = ImmutableList.builder(); + IcebergUtil.getColumns(tableSchema, typeManager).forEach(columnHandlesBuilder::add); + columnHandlesBuilder.add(new IcebergColumnHandle( + new ColumnIdentity(DATA_CHANGE_TYPE_ID, DATA_CHANGE_TYPE_NAME, PRIMITIVE, ImmutableList.of()), + VarcharType.createUnboundedVarcharType(), + ImmutableList.of(), + VarcharType.createUnboundedVarcharType(), + false, + Optional.empty())); + columnHandlesBuilder.add(new IcebergColumnHandle( + new ColumnIdentity(DATA_CHANGE_VERSION_ID, DATA_CHANGE_VERSION_NAME, PRIMITIVE, ImmutableList.of()), + BIGINT, + ImmutableList.of(), + BIGINT, + false, + Optional.empty())); + columnHandlesBuilder.add(new IcebergColumnHandle( + new ColumnIdentity(DATA_CHANGE_TIMESTAMP_ID, DATA_CHANGE_TIMESTAMP_NAME, PRIMITIVE, ImmutableList.of()), + TIMESTAMP_TZ_MILLIS, + ImmutableList.of(), + TIMESTAMP_TZ_MILLIS, + false, + Optional.empty())); + columnHandlesBuilder.add(new IcebergColumnHandle( + new ColumnIdentity(DATA_CHANGE_ORDINAL_ID, DATA_CHANGE_ORDINAL_NAME, PRIMITIVE, ImmutableList.of()), + INTEGER, + ImmutableList.of(), + INTEGER, + false, + Optional.empty())); + List columnHandles = columnHandlesBuilder.build(); + + accessControl.checkCanSelectFromColumns(null, schemaTableName, columnHandles.stream() + .map(IcebergColumnHandle::getName) + .collect(toImmutableSet())); + + return TableFunctionAnalysis.builder() + .returnedType(new Descriptor(columns.build())) + .handle(new TableChangesFunctionHandle( + schemaTableName, + SchemaParser.toJson(tableSchema), + columnHandles, + Optional.ofNullable(icebergTable.properties().get(TableProperties.DEFAULT_NAME_MAPPING)), + startSnapshotId, + endSnapshotId)) + .build(); + } + + private static Object checkNonNull(Object argumentValue) + { + if (argumentValue == null) { + throw new TrinoException(INVALID_FUNCTION_ARGUMENT, FUNCTION_NAME + " arguments may not be null"); + } + return argumentValue; + } + + private static void checkSnapshotExists(Table icebergTable, long snapshotId) + { + if (icebergTable.snapshot(snapshotId) == null) { + throw new TrinoException(INVALID_FUNCTION_ARGUMENT, "Snapshot not found in Iceberg table history: " + snapshotId); + } + } +} diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/functions/tablechanges/TableChangesFunctionHandle.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/functions/tablechanges/TableChangesFunctionHandle.java new file mode 100644 index 000000000000..97354093476c --- /dev/null +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/functions/tablechanges/TableChangesFunctionHandle.java @@ -0,0 +1,41 @@ +/* + * 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.iceberg.functions.tablechanges; + +import com.google.common.collect.ImmutableList; +import io.trino.plugin.iceberg.IcebergColumnHandle; +import io.trino.spi.connector.SchemaTableName; +import io.trino.spi.function.table.ConnectorTableFunctionHandle; + +import java.util.List; +import java.util.Optional; + +import static java.util.Objects.requireNonNull; + +public record TableChangesFunctionHandle( + SchemaTableName schemaTableName, + String tableSchemaJson, + List columns, + Optional nameMappingJson, + long startSnapshotId, + long endSnapshotId) implements ConnectorTableFunctionHandle +{ + public TableChangesFunctionHandle + { + requireNonNull(schemaTableName, "schemaTableName is null"); + requireNonNull(tableSchemaJson, "tableSchemaJson is null"); + columns = ImmutableList.copyOf(requireNonNull(columns, "columns is null")); + requireNonNull(nameMappingJson, "nameMappingJson is null"); + } +} diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/functions/tablechanges/TableChangesFunctionProcessor.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/functions/tablechanges/TableChangesFunctionProcessor.java new file mode 100644 index 000000000000..d93f661bf231 --- /dev/null +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/functions/tablechanges/TableChangesFunctionProcessor.java @@ -0,0 +1,180 @@ +/* + * 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.iceberg.functions.tablechanges; + +import com.google.common.collect.ImmutableList; +import io.trino.plugin.iceberg.IcebergColumnHandle; +import io.trino.plugin.iceberg.IcebergPageSourceProvider; +import io.trino.plugin.iceberg.PartitionData; +import io.trino.spi.Page; +import io.trino.spi.block.Block; +import io.trino.spi.block.RunLengthEncodedBlock; +import io.trino.spi.connector.ConnectorPageSource; +import io.trino.spi.connector.ConnectorSession; +import io.trino.spi.connector.DynamicFilter; +import io.trino.spi.function.table.TableFunctionProcessorState; +import io.trino.spi.function.table.TableFunctionSplitProcessor; +import io.trino.spi.predicate.TupleDomain; +import org.apache.iceberg.PartitionSpec; +import org.apache.iceberg.PartitionSpecParser; +import org.apache.iceberg.Schema; +import org.apache.iceberg.SchemaParser; +import org.apache.iceberg.mapping.NameMappingParser; + +import java.util.Optional; + +import static io.airlift.slice.Slices.utf8Slice; +import static io.trino.plugin.iceberg.IcebergColumnHandle.DATA_CHANGE_ORDINAL_ID; +import static io.trino.plugin.iceberg.IcebergColumnHandle.DATA_CHANGE_TIMESTAMP_ID; +import static io.trino.plugin.iceberg.IcebergColumnHandle.DATA_CHANGE_TYPE_ID; +import static io.trino.plugin.iceberg.IcebergColumnHandle.DATA_CHANGE_VERSION_ID; +import static io.trino.spi.function.table.TableFunctionProcessorState.Finished.FINISHED; +import static io.trino.spi.function.table.TableFunctionProcessorState.Processed.produced; +import static io.trino.spi.predicate.Utils.nativeValueToBlock; +import static io.trino.spi.type.BigintType.BIGINT; +import static io.trino.spi.type.IntegerType.INTEGER; +import static io.trino.spi.type.TimestampWithTimeZoneType.TIMESTAMP_TZ_MILLIS; +import static io.trino.spi.type.VarcharType.createUnboundedVarcharType; +import static java.util.Objects.requireNonNull; + +public class TableChangesFunctionProcessor + implements TableFunctionSplitProcessor +{ + private static final Page EMPTY_PAGE = new Page(0); + + private final ConnectorPageSource pageSource; + private final int[] delegateColumnMap; + private final Optional changeTypeIndex; + private final Block changeTypeValue; + private final Optional changeVersionIndex; + private final Block changeVersionValue; + private final Optional changeTimestampIndex; + private final Block changeTimestampValue; + private final Optional changeOrdinalIndex; + private final Block changeOrdinalValue; + + public TableChangesFunctionProcessor( + ConnectorSession session, + TableChangesFunctionHandle functionHandle, + TableChangesSplit split, + IcebergPageSourceProvider icebergPageSourceProvider) + { + requireNonNull(session, "session is null"); + requireNonNull(functionHandle, "functionHandle is null"); + requireNonNull(split, "split is null"); + requireNonNull(icebergPageSourceProvider, "icebergPageSourceProvider is null"); + + Schema tableSchema = SchemaParser.fromJson(functionHandle.tableSchemaJson()); + PartitionSpec partitionSpec = PartitionSpecParser.fromJson(tableSchema, split.partitionSpecJson()); + org.apache.iceberg.types.Type[] partitionColumnTypes = partitionSpec.fields().stream() + .map(field -> field.transform().getResultType(tableSchema.findType(field.sourceId()))) + .toArray(org.apache.iceberg.types.Type[]::new); + + int delegateColumnIndex = 0; + int[] delegateColumnMap = new int[functionHandle.columns().size()]; + Optional changeTypeIndex = Optional.empty(); + Optional changeVersionIndex = Optional.empty(); + Optional changeTimestampIndex = Optional.empty(); + Optional changeOrdinalIndex = Optional.empty(); + for (int columnIndex = 0; columnIndex < functionHandle.columns().size(); columnIndex++) { + IcebergColumnHandle column = functionHandle.columns().get(columnIndex); + if (column.getId() == DATA_CHANGE_TYPE_ID) { + changeTypeIndex = Optional.of(columnIndex); + delegateColumnMap[columnIndex] = -1; + } + else if (column.getId() == DATA_CHANGE_VERSION_ID) { + changeVersionIndex = Optional.of(columnIndex); + delegateColumnMap[columnIndex] = -1; + } + else if (column.getId() == DATA_CHANGE_TIMESTAMP_ID) { + changeTimestampIndex = Optional.of(columnIndex); + delegateColumnMap[columnIndex] = -1; + } + else if (column.getId() == DATA_CHANGE_ORDINAL_ID) { + changeOrdinalIndex = Optional.of(columnIndex); + delegateColumnMap[columnIndex] = -1; + } + else { + delegateColumnMap[columnIndex] = delegateColumnIndex; + delegateColumnIndex++; + } + } + + this.pageSource = icebergPageSourceProvider.createPageSource( + session, + functionHandle.columns(), + tableSchema, + partitionSpec, + PartitionData.fromJson(split.partitionDataJson(), partitionColumnTypes), + ImmutableList.of(), + DynamicFilter.EMPTY, + TupleDomain.all(), + TupleDomain.all(), + split.path(), + split.start(), + split.length(), + split.fileSize(), + split.fileRecordCount(), + split.partitionDataJson(), + split.fileFormat(), + split.fileIoProperties(), + 0, + functionHandle.nameMappingJson().map(NameMappingParser::fromJson)); + this.delegateColumnMap = delegateColumnMap; + + this.changeTypeIndex = changeTypeIndex; + this.changeTypeValue = nativeValueToBlock(createUnboundedVarcharType(), utf8Slice(split.changeType().getTableValue())); + + this.changeVersionIndex = changeVersionIndex; + this.changeVersionValue = nativeValueToBlock(BIGINT, split.snapshotId()); + + this.changeTimestampIndex = changeTimestampIndex; + this.changeTimestampValue = nativeValueToBlock(TIMESTAMP_TZ_MILLIS, split.snapshotTimestamp()); + + this.changeOrdinalIndex = changeOrdinalIndex; + this.changeOrdinalValue = nativeValueToBlock(INTEGER, (long) split.changeOrdinal()); + } + + @Override + public TableFunctionProcessorState process() + { + if (pageSource.isFinished()) { + return FINISHED; + } + + Page dataPage = pageSource.getNextPage(); + if (dataPage == null) { + return TableFunctionProcessorState.Processed.produced(EMPTY_PAGE); + } + + Block[] blocks = new Block[delegateColumnMap.length]; + for (int targetChannel = 0; targetChannel < delegateColumnMap.length; targetChannel++) { + int delegateIndex = delegateColumnMap[targetChannel]; + if (delegateIndex != -1) { + blocks[targetChannel] = dataPage.getBlock(delegateIndex); + } + } + + changeTypeIndex.ifPresent(columnChannel -> + blocks[columnChannel] = RunLengthEncodedBlock.create(changeTypeValue, dataPage.getPositionCount())); + changeVersionIndex.ifPresent(columnChannel -> + blocks[columnChannel] = RunLengthEncodedBlock.create(changeVersionValue, dataPage.getPositionCount())); + changeTimestampIndex.ifPresent(columnChannel -> + blocks[columnChannel] = RunLengthEncodedBlock.create(changeTimestampValue, dataPage.getPositionCount())); + changeOrdinalIndex.ifPresent(columnChannel -> + blocks[columnChannel] = RunLengthEncodedBlock.create(changeOrdinalValue, dataPage.getPositionCount())); + + return produced(new Page(dataPage.getPositionCount(), blocks)); + } +} diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/functions/tablechanges/TableChangesFunctionProcessorProvider.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/functions/tablechanges/TableChangesFunctionProcessorProvider.java new file mode 100644 index 000000000000..80ef0fd49152 --- /dev/null +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/functions/tablechanges/TableChangesFunctionProcessorProvider.java @@ -0,0 +1,50 @@ +/* + * 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.iceberg.functions.tablechanges; + +import com.google.inject.Inject; +import io.trino.plugin.iceberg.IcebergPageSourceProvider; +import io.trino.plugin.iceberg.IcebergPageSourceProviderFactory; +import io.trino.spi.connector.ConnectorSession; +import io.trino.spi.connector.ConnectorSplit; +import io.trino.spi.function.table.ConnectorTableFunctionHandle; +import io.trino.spi.function.table.TableFunctionProcessorProvider; +import io.trino.spi.function.table.TableFunctionSplitProcessor; + +import static java.util.Objects.requireNonNull; + +public class TableChangesFunctionProcessorProvider + implements TableFunctionProcessorProvider +{ + private final IcebergPageSourceProviderFactory icebergPageSourceProviderFactory; + + @Inject + public TableChangesFunctionProcessorProvider(IcebergPageSourceProviderFactory icebergPageSourceProviderFactory) + { + this.icebergPageSourceProviderFactory = requireNonNull(icebergPageSourceProviderFactory, "icebergPageSourceProviderFactory is null"); + } + + @Override + public TableFunctionSplitProcessor getSplitProcessor( + ConnectorSession session, + ConnectorTableFunctionHandle handle, + ConnectorSplit split) + { + return new TableChangesFunctionProcessor( + session, + (TableChangesFunctionHandle) handle, + (TableChangesSplit) split, + (IcebergPageSourceProvider) icebergPageSourceProviderFactory.createPageSourceProvider()); + } +} diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/functions/tablechanges/TableChangesFunctionProvider.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/functions/tablechanges/TableChangesFunctionProvider.java new file mode 100644 index 000000000000..18d0d4a748f3 --- /dev/null +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/functions/tablechanges/TableChangesFunctionProvider.java @@ -0,0 +1,45 @@ +/* + * 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.iceberg.functions.tablechanges; + +import com.google.inject.Inject; +import com.google.inject.Provider; +import io.trino.plugin.base.classloader.ClassLoaderSafeConnectorTableFunction; +import io.trino.plugin.iceberg.catalog.TrinoCatalogFactory; +import io.trino.spi.function.table.ConnectorTableFunction; +import io.trino.spi.type.TypeManager; + +import static java.util.Objects.requireNonNull; + +public class TableChangesFunctionProvider + implements Provider +{ + private final TrinoCatalogFactory trinoCatalogFactory; + private final TypeManager typeManager; + + @Inject + public TableChangesFunctionProvider(TrinoCatalogFactory trinoCatalogFactory, TypeManager typeManager) + { + this.trinoCatalogFactory = requireNonNull(trinoCatalogFactory, "trinoCatalogFactory is null"); + this.typeManager = requireNonNull(typeManager, "typeManager is null"); + } + + @Override + public ConnectorTableFunction get() + { + return new ClassLoaderSafeConnectorTableFunction( + new TableChangesFunction(trinoCatalogFactory, typeManager), + getClass().getClassLoader()); + } +} diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/functions/tablechanges/TableChangesSplit.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/functions/tablechanges/TableChangesSplit.java new file mode 100644 index 000000000000..5209e97d8115 --- /dev/null +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/functions/tablechanges/TableChangesSplit.java @@ -0,0 +1,126 @@ +/* + * 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.iceberg.functions.tablechanges; + +import com.google.common.collect.ImmutableMap; +import io.airlift.slice.SizeOf; +import io.trino.plugin.iceberg.IcebergFileFormat; +import io.trino.spi.HostAddress; +import io.trino.spi.SplitWeight; +import io.trino.spi.connector.ConnectorSplit; + +import java.util.List; +import java.util.Map; + +import static com.google.common.base.MoreObjects.toStringHelper; +import static io.airlift.slice.SizeOf.estimatedSizeOf; +import static java.util.Objects.requireNonNull; + +public record TableChangesSplit( + ChangeType changeType, + long snapshotId, + long snapshotTimestamp, + int changeOrdinal, + String path, + long start, + long length, + long fileSize, + long fileRecordCount, + IcebergFileFormat fileFormat, + String partitionSpecJson, + String partitionDataJson, + SplitWeight splitWeight, + Map fileIoProperties) implements ConnectorSplit +{ + private static final int INSTANCE_SIZE = SizeOf.instanceSize(TableChangesSplit.class); + + public TableChangesSplit + { + requireNonNull(changeType, "changeType is null"); + requireNonNull(path, "path is null"); + requireNonNull(fileFormat, "fileFormat is null"); + requireNonNull(partitionSpecJson, "partitionSpecJson is null"); + requireNonNull(partitionDataJson, "partitionDataJson is null"); + requireNonNull(splitWeight, "splitWeight is null"); + fileIoProperties = ImmutableMap.copyOf(requireNonNull(fileIoProperties, "fileIoProperties is null")); + } + + @Override + public SplitWeight getSplitWeight() + { + return splitWeight; + } + + @Override + public boolean isRemotelyAccessible() + { + return true; + } + + @Override + public List getAddresses() + { + return List.of(); + } + + @Override + public Map getInfo() + { + return ImmutableMap.builder() + .put("path", path) + .put("start", String.valueOf(start)) + .put("length", String.valueOf(length)) + .buildOrThrow(); + } + + @Override + public long getRetainedSizeInBytes() + { + return INSTANCE_SIZE + + estimatedSizeOf(path) + + estimatedSizeOf(partitionSpecJson) + + estimatedSizeOf(partitionDataJson) + + splitWeight.getRetainedSizeInBytes() + + estimatedSizeOf(fileIoProperties, SizeOf::estimatedSizeOf, SizeOf::estimatedSizeOf); + } + + @Override + public String toString() + { + return toStringHelper(this) + .addValue(path) + .add("start", start) + .add("length", length) + .add("records", fileRecordCount) + .toString(); + } + + public enum ChangeType { + ADDED_FILE("insert"), + DELETED_FILE("delete"), + POSITIONAL_DELETE("delete"); + + private final String tableValue; + + ChangeType(String tableValue) + { + this.tableValue = tableValue; + } + + public String getTableValue() + { + return tableValue; + } + } +} diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/functions/tablechanges/TableChangesSplitSource.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/functions/tablechanges/TableChangesSplitSource.java new file mode 100644 index 000000000000..097c0d04919d --- /dev/null +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/functions/tablechanges/TableChangesSplitSource.java @@ -0,0 +1,178 @@ +/* + * 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.iceberg.functions.tablechanges; + +import com.google.common.io.Closer; +import io.trino.plugin.iceberg.IcebergFileFormat; +import io.trino.plugin.iceberg.PartitionData; +import io.trino.spi.SplitWeight; +import io.trino.spi.TrinoException; +import io.trino.spi.connector.ConnectorSplit; +import io.trino.spi.connector.ConnectorSplitSource; +import io.trino.spi.type.DateTimeEncoding; +import org.apache.iceberg.AddedRowsScanTask; +import org.apache.iceberg.ChangelogScanTask; +import org.apache.iceberg.DeletedDataFileScanTask; +import org.apache.iceberg.IncrementalChangelogScan; +import org.apache.iceberg.PartitionSpecParser; +import org.apache.iceberg.SplittableScanTask; +import org.apache.iceberg.Table; +import org.apache.iceberg.io.CloseableIterable; +import org.apache.iceberg.io.CloseableIterator; + +import java.io.IOException; +import java.io.UncheckedIOException; +import java.util.ArrayList; +import java.util.Iterator; +import java.util.List; +import java.util.concurrent.CompletableFuture; + +import static com.google.common.collect.Iterators.singletonIterator; +import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED; +import static io.trino.spi.type.TimeZoneKey.UTC_KEY; +import static java.util.Collections.emptyIterator; +import static java.util.Objects.requireNonNull; +import static java.util.concurrent.CompletableFuture.completedFuture; + +public class TableChangesSplitSource + implements ConnectorSplitSource +{ + private final Table icebergTable; + private final IncrementalChangelogScan tableScan; + private final long targetSplitSize; + private final Closer closer = Closer.create(); + + private CloseableIterable changelogScanIterable; + private CloseableIterator changelogScanIterator; + private Iterator fileTasksIterator = emptyIterator(); + + public TableChangesSplitSource( + Table icebergTable, + IncrementalChangelogScan tableScan) + { + this.icebergTable = requireNonNull(icebergTable, "table is null"); + this.tableScan = requireNonNull(tableScan, "tableScan is null"); + this.targetSplitSize = tableScan.targetSplitSize(); + } + + @Override + public CompletableFuture getNextBatch(int maxSize) + { + if (changelogScanIterable == null) { + try { + this.changelogScanIterable = closer.register(tableScan.planFiles()); + this.changelogScanIterator = closer.register(changelogScanIterable.iterator()); + } + catch (UnsupportedOperationException e) { + throw new TrinoException(NOT_SUPPORTED, "Table uses features which are not yet supported by the table_changes function", e); + } + } + + List splits = new ArrayList<>(maxSize); + while (splits.size() < maxSize && (fileTasksIterator.hasNext() || changelogScanIterator.hasNext())) { + if (!fileTasksIterator.hasNext()) { + ChangelogScanTask wholeFileTask = changelogScanIterator.next(); + fileTasksIterator = splitIfPossible(wholeFileTask, targetSplitSize); + continue; + } + + ChangelogScanTask next = fileTasksIterator.next(); + splits.add(toIcebergSplit(next)); + } + return completedFuture(new ConnectorSplitBatch(splits, isFinished())); + } + + @Override + public boolean isFinished() + { + return changelogScanIterator != null && !changelogScanIterator.hasNext(); + } + + @Override + public void close() + { + try { + closer.close(); + } + catch (IOException e) { + throw new UncheckedIOException(e); + } + } + + @SuppressWarnings("unchecked") + private static Iterator splitIfPossible(ChangelogScanTask wholeFileScan, long targetSplitSize) + { + if (wholeFileScan instanceof AddedRowsScanTask) { + return ((SplittableScanTask) wholeFileScan).split(targetSplitSize).iterator(); + } + + if (wholeFileScan instanceof DeletedDataFileScanTask) { + return ((SplittableScanTask) wholeFileScan).split(targetSplitSize).iterator(); + } + + return singletonIterator(wholeFileScan); + } + + private ConnectorSplit toIcebergSplit(ChangelogScanTask task) + { + // TODO: Support DeletedRowsScanTask (requires https://github.com/apache/iceberg/pull/6182) + if (task instanceof AddedRowsScanTask) { + return toSplit((AddedRowsScanTask) task); + } + else if (task instanceof DeletedDataFileScanTask) { + return toSplit((DeletedDataFileScanTask) task); + } + else { + throw new TrinoException(NOT_SUPPORTED, "ChangelogScanTask type is not supported:" + task); + } + } + + private TableChangesSplit toSplit(AddedRowsScanTask task) + { + return new TableChangesSplit( + TableChangesSplit.ChangeType.ADDED_FILE, + task.commitSnapshotId(), + DateTimeEncoding.packDateTimeWithZone(icebergTable.snapshot(task.commitSnapshotId()).timestampMillis(), UTC_KEY), + task.changeOrdinal(), + task.file().location(), + task.start(), + task.length(), + task.file().fileSizeInBytes(), + task.file().recordCount(), + IcebergFileFormat.fromIceberg(task.file().format()), + PartitionSpecParser.toJson(task.spec()), + PartitionData.toJson(task.file().partition()), + SplitWeight.standard(), + icebergTable.io().properties()); + } + + private TableChangesSplit toSplit(DeletedDataFileScanTask task) + { + return new TableChangesSplit( + TableChangesSplit.ChangeType.DELETED_FILE, + task.commitSnapshotId(), + DateTimeEncoding.packDateTimeWithZone(icebergTable.snapshot(task.commitSnapshotId()).timestampMillis(), UTC_KEY), + task.changeOrdinal(), + task.file().location(), + task.start(), + task.length(), + task.file().fileSizeInBytes(), + task.file().recordCount(), + IcebergFileFormat.fromIceberg(task.file().format()), + PartitionSpecParser.toJson(task.spec()), + PartitionData.toJson(task.file().partition()), + SplitWeight.standard(), + icebergTable.io().properties()); + } +} diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/procedure/IcebergAddFilesFromTableHandle.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/procedure/IcebergAddFilesFromTableHandle.java new file mode 100644 index 000000000000..67f455a28c38 --- /dev/null +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/procedure/IcebergAddFilesFromTableHandle.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.iceberg.procedure; + +import io.trino.plugin.iceberg.procedure.MigrationUtils.RecursiveDirectory; +import jakarta.annotation.Nullable; + +import java.util.Map; + +import static java.util.Objects.requireNonNull; + +public record IcebergAddFilesFromTableHandle( + io.trino.plugin.hive.metastore.Table table, + @Nullable Map partitionFilter, + RecursiveDirectory recursiveDirectory) + implements IcebergProcedureHandle +{ + public IcebergAddFilesFromTableHandle + { + requireNonNull(table, "table is null"); + requireNonNull(recursiveDirectory, "recursiveDirectory is null"); + } +} diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/procedure/IcebergAddFilesHandle.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/procedure/IcebergAddFilesHandle.java new file mode 100644 index 000000000000..6fd365a54089 --- /dev/null +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/procedure/IcebergAddFilesHandle.java @@ -0,0 +1,30 @@ +/* + * 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.iceberg.procedure; + +import io.trino.plugin.hive.HiveStorageFormat; +import io.trino.plugin.iceberg.procedure.MigrationUtils.RecursiveDirectory; + +import static java.util.Objects.requireNonNull; + +public record IcebergAddFilesHandle(String location, HiveStorageFormat format, RecursiveDirectory recursiveDirectory) + implements IcebergProcedureHandle +{ + public IcebergAddFilesHandle + { + requireNonNull(location, "location is null"); + requireNonNull(format, "format is null"); + requireNonNull(recursiveDirectory, "recursiveDirectory is null"); + } +} diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/procedure/IcebergDropExtendedStatsHandle.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/procedure/IcebergDropExtendedStatsHandle.java index 55bf7e092a7d..a42fc0bd42d6 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/procedure/IcebergDropExtendedStatsHandle.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/procedure/IcebergDropExtendedStatsHandle.java @@ -16,7 +16,7 @@ import static com.google.common.base.MoreObjects.toStringHelper; public class IcebergDropExtendedStatsHandle - extends IcebergProcedureHandle + implements IcebergProcedureHandle { @Override public String toString() diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/procedure/IcebergExpireSnapshotsHandle.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/procedure/IcebergExpireSnapshotsHandle.java index 02687389b233..fa403ee1ed5c 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/procedure/IcebergExpireSnapshotsHandle.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/procedure/IcebergExpireSnapshotsHandle.java @@ -21,7 +21,7 @@ import static java.util.Objects.requireNonNull; public class IcebergExpireSnapshotsHandle - extends IcebergProcedureHandle + implements IcebergProcedureHandle { private final Duration retentionThreshold; diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/procedure/IcebergOptimizeHandle.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/procedure/IcebergOptimizeHandle.java index b64238376536..5e23ce5aaa02 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/procedure/IcebergOptimizeHandle.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/procedure/IcebergOptimizeHandle.java @@ -30,7 +30,7 @@ import static java.util.Objects.requireNonNull; public class IcebergOptimizeHandle - extends IcebergProcedureHandle + implements IcebergProcedureHandle { private final Optional snapshotId; private final String schemaAsJson; diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/procedure/IcebergOptimizeManifestsHandle.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/procedure/IcebergOptimizeManifestsHandle.java new file mode 100644 index 000000000000..cc1a44174ec5 --- /dev/null +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/procedure/IcebergOptimizeManifestsHandle.java @@ -0,0 +1,17 @@ +/* + * 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.iceberg.procedure; + +public record IcebergOptimizeManifestsHandle() + implements IcebergProcedureHandle {} diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/procedure/IcebergProcedureHandle.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/procedure/IcebergProcedureHandle.java index e9ce4199eeac..bd2c64e0778c 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/procedure/IcebergProcedureHandle.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/procedure/IcebergProcedureHandle.java @@ -20,9 +20,13 @@ use = JsonTypeInfo.Id.NAME, property = "@type") @JsonSubTypes({ - @JsonSubTypes.Type(value = IcebergOptimizeHandle.class, name = "optimize"), @JsonSubTypes.Type(value = IcebergDropExtendedStatsHandle.class, name = "drop_extended_stats"), + @JsonSubTypes.Type(value = IcebergRollbackToSnapshotHandle.class, name = "rollback_to_snapshot"), @JsonSubTypes.Type(value = IcebergExpireSnapshotsHandle.class, name = "expire_snapshots"), + @JsonSubTypes.Type(value = IcebergOptimizeHandle.class, name = "optimize"), + @JsonSubTypes.Type(value = IcebergOptimizeManifestsHandle.class, name = "optimize_manifests"), @JsonSubTypes.Type(value = IcebergRemoveOrphanFilesHandle.class, name = "remove_orphan_files"), + @JsonSubTypes.Type(value = IcebergAddFilesHandle.class, name = "add_files"), + @JsonSubTypes.Type(value = IcebergAddFilesFromTableHandle.class, name = "add_files_from_table"), }) -public abstract class IcebergProcedureHandle {} +public interface IcebergProcedureHandle {} diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/procedure/IcebergRemoveOrphanFilesHandle.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/procedure/IcebergRemoveOrphanFilesHandle.java index e4ac5d9f5890..13c64b3f3c2c 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/procedure/IcebergRemoveOrphanFilesHandle.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/procedure/IcebergRemoveOrphanFilesHandle.java @@ -21,7 +21,7 @@ import static java.util.Objects.requireNonNull; public class IcebergRemoveOrphanFilesHandle - extends IcebergProcedureHandle + implements IcebergProcedureHandle { private final Duration retentionThreshold; diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/procedure/IcebergRollbackToSnapshotHandle.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/procedure/IcebergRollbackToSnapshotHandle.java new file mode 100644 index 000000000000..aad1ba3218b1 --- /dev/null +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/procedure/IcebergRollbackToSnapshotHandle.java @@ -0,0 +1,17 @@ +/* + * 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.iceberg.procedure; + +public record IcebergRollbackToSnapshotHandle(long snapshotId) + implements IcebergProcedureHandle {} diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/procedure/IcebergTableExecuteHandle.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/procedure/IcebergTableExecuteHandle.java index 1ffbf7e14d3f..9c9bc399d801 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/procedure/IcebergTableExecuteHandle.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/procedure/IcebergTableExecuteHandle.java @@ -18,6 +18,8 @@ import io.trino.spi.connector.ConnectorTableExecuteHandle; import io.trino.spi.connector.SchemaTableName; +import java.util.Map; + import static java.util.Objects.requireNonNull; public class IcebergTableExecuteHandle @@ -27,18 +29,21 @@ public class IcebergTableExecuteHandle private final IcebergTableProcedureId procedureId; private final IcebergProcedureHandle procedureHandle; private final String tableLocation; + private final Map fileIoProperties; @JsonCreator public IcebergTableExecuteHandle( SchemaTableName schemaTableName, IcebergTableProcedureId procedureId, IcebergProcedureHandle procedureHandle, - String tableLocation) + String tableLocation, + Map fileIoProperties) { this.schemaTableName = requireNonNull(schemaTableName, "schemaTableName is null"); this.procedureId = requireNonNull(procedureId, "procedureId is null"); this.procedureHandle = requireNonNull(procedureHandle, "procedureHandle is null"); this.tableLocation = requireNonNull(tableLocation, "tableLocation is null"); + this.fileIoProperties = requireNonNull(fileIoProperties, "fileIoProperties is null"); } @JsonProperty @@ -65,13 +70,20 @@ public String getTableLocation() return tableLocation; } + @JsonProperty + public Map getFileIoProperties() + { + return fileIoProperties; + } + public IcebergTableExecuteHandle withProcedureHandle(IcebergProcedureHandle procedureHandle) { return new IcebergTableExecuteHandle( schemaTableName, procedureId, procedureHandle, - tableLocation); + tableLocation, + fileIoProperties); } @Override diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/procedure/IcebergTableProcedureId.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/procedure/IcebergTableProcedureId.java index 8b1c68fb23ed..6230f8b779b6 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/procedure/IcebergTableProcedureId.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/procedure/IcebergTableProcedureId.java @@ -16,7 +16,11 @@ public enum IcebergTableProcedureId { OPTIMIZE, + OPTIMIZE_MANIFESTS, DROP_EXTENDED_STATS, + ROLLBACK_TO_SNAPSHOT, EXPIRE_SNAPSHOTS, REMOVE_ORPHAN_FILES, + ADD_FILES, + ADD_FILES_FROM_TABLE, } diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/procedure/MigrateProcedure.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/procedure/MigrateProcedure.java index 6c57853c9e4b..a9601675121a 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/procedure/MigrateProcedure.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/procedure/MigrateProcedure.java @@ -246,7 +246,7 @@ public void doMigrate(ConnectorSession session, String schemaName, String tableN schema, parsePartitionFields(schema, toPartitionFields(hiveTable)), unsorted(), - location, + Optional.of(location), properties); List dataFiles = dataFilesBuilder.build(); @@ -262,7 +262,7 @@ public void doMigrate(ConnectorSession session, String schemaName, String tableN .setParameter(METADATA_LOCATION_PROP, location) .setParameter(TABLE_TYPE_PROP, ICEBERG_TABLE_TYPE_VALUE.toUpperCase(ENGLISH)) .build(); - metastore.replaceTable(schemaName, tableName, newTable, principalPrivileges); + metastore.replaceTable(schemaName, tableName, newTable, principalPrivileges, ImmutableMap.of()); transaction.commitTransaction(); log.debug("Successfully migrated %s table to Iceberg format", sourceTableName); diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/procedure/MigrationUtils.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/procedure/MigrationUtils.java new file mode 100644 index 000000000000..c80e7cd081da --- /dev/null +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/procedure/MigrationUtils.java @@ -0,0 +1,328 @@ +/* + * 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.iceberg.procedure; + +import com.google.common.base.Joiner; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableSet; +import io.airlift.log.Logger; +import io.trino.filesystem.FileEntry; +import io.trino.filesystem.FileIterator; +import io.trino.filesystem.Location; +import io.trino.filesystem.TrinoFileSystem; +import io.trino.filesystem.TrinoInputFile; +import io.trino.parquet.ParquetDataSource; +import io.trino.parquet.ParquetReaderOptions; +import io.trino.parquet.metadata.ParquetMetadata; +import io.trino.parquet.reader.MetadataReader; +import io.trino.plugin.hive.FileFormatDataSourceStats; +import io.trino.plugin.hive.HiveStorageFormat; +import io.trino.plugin.hive.metastore.HiveMetastore; +import io.trino.plugin.hive.metastore.HiveMetastoreFactory; +import io.trino.plugin.hive.metastore.Partition; +import io.trino.plugin.hive.metastore.Storage; +import io.trino.plugin.hive.parquet.TrinoParquetDataSource; +import io.trino.plugin.iceberg.catalog.TrinoCatalog; +import io.trino.plugin.iceberg.fileio.ForwardingInputFile; +import io.trino.plugin.iceberg.util.OrcMetrics; +import io.trino.plugin.iceberg.util.ParquetUtil; +import io.trino.spi.TrinoException; +import io.trino.spi.connector.ConnectorSession; +import io.trino.spi.connector.SchemaTableName; +import org.apache.iceberg.AppendFiles; +import org.apache.iceberg.DataFile; +import org.apache.iceberg.DataFiles; +import org.apache.iceberg.FileScanTask; +import org.apache.iceberg.Metrics; +import org.apache.iceberg.MetricsConfig; +import org.apache.iceberg.PartitionSpec; +import org.apache.iceberg.Schema; +import org.apache.iceberg.StructLike; +import org.apache.iceberg.Table; +import org.apache.iceberg.Transaction; +import org.apache.iceberg.avro.Avro; +import org.apache.iceberg.io.CloseableIterable; +import org.apache.iceberg.mapping.MappingUtil; +import org.apache.iceberg.mapping.NameMapping; +import org.apache.iceberg.types.Types; + +import java.io.IOException; +import java.io.UncheckedIOException; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.stream.Stream; + +import static com.google.common.base.MoreObjects.firstNonNull; +import static com.google.common.base.Verify.verify; +import static com.google.common.collect.ImmutableSet.toImmutableSet; +import static io.trino.plugin.base.util.Procedures.checkProcedureArgument; +import static io.trino.plugin.hive.HiveMetadata.extractHiveStorageFormat; +import static io.trino.plugin.iceberg.IcebergErrorCode.ICEBERG_COMMIT_ERROR; +import static io.trino.plugin.iceberg.IcebergSessionProperties.isMergeManifestsOnWrite; +import static io.trino.spi.StandardErrorCode.ALREADY_EXISTS; +import static io.trino.spi.StandardErrorCode.CONSTRAINT_VIOLATION; +import static io.trino.spi.StandardErrorCode.NOT_FOUND; +import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED; +import static org.apache.iceberg.TableProperties.DEFAULT_NAME_MAPPING; +import static org.apache.iceberg.mapping.NameMappingParser.toJson; + +public final class MigrationUtils +{ + private static final Logger log = Logger.get(MigrationUtils.class); + private static final Joiner.MapJoiner PARTITION_JOINER = Joiner.on("/").withKeyValueSeparator("="); + + private static final MetricsConfig METRICS_CONFIG = MetricsConfig.getDefault(); + + public enum RecursiveDirectory + { + TRUE, + FALSE, + FAIL, + /**/ + } + + private MigrationUtils() {} + + public static List buildDataFiles( + TrinoFileSystem fileSystem, + RecursiveDirectory recursive, + HiveStorageFormat format, + String location, + PartitionSpec partitionSpec, + Optional partition, + Schema schema) + throws IOException + { + // TODO: Introduce parallelism + FileIterator files = fileSystem.listFiles(Location.of(location)); + ImmutableList.Builder dataFilesBuilder = ImmutableList.builder(); + while (files.hasNext()) { + FileEntry file = files.next(); + String fileLocation = file.location().toString(); + String relativePath = fileLocation.substring(location.length()); + if (relativePath.contains("/_") || relativePath.contains("/.")) { + continue; + } + if (recursive == RecursiveDirectory.FALSE && isRecursive(location, fileLocation)) { + continue; + } + if (recursive == RecursiveDirectory.FAIL && isRecursive(location, fileLocation)) { + throw new TrinoException(NOT_SUPPORTED, "Recursive directory must not exist when recursive_directory argument is 'fail': " + file.location()); + } + + Metrics metrics = loadMetrics(fileSystem.newInputFile(file.location(), file.length()), format, schema); + DataFile dataFile = buildDataFile(fileLocation, file.length(), partition, partitionSpec, format.name(), metrics); + dataFilesBuilder.add(dataFile); + } + List dataFiles = dataFilesBuilder.build(); + log.debug("Found %d files in '%s'", dataFiles.size(), location); + return dataFiles; + } + + private static boolean isRecursive(String baseLocation, String location) + { + verify(location.startsWith(baseLocation), "%s should start with %s", location, baseLocation); + String suffix = location.substring(baseLocation.length() + 1).replaceFirst("^/+", ""); + return suffix.contains("/"); + } + + public static Metrics loadMetrics(TrinoInputFile file, HiveStorageFormat storageFormat, Schema schema) + { + return switch (storageFormat) { + case ORC -> OrcMetrics.fileMetrics(file, METRICS_CONFIG, schema); + case PARQUET -> parquetMetrics(file, METRICS_CONFIG, MappingUtil.create(schema)); + case AVRO -> new Metrics(Avro.rowCount(new ForwardingInputFile(file)), null, null, null, null); + default -> throw new TrinoException(NOT_SUPPORTED, "Unsupported storage format: " + storageFormat); + }; + } + + private static Metrics parquetMetrics(TrinoInputFile file, MetricsConfig metricsConfig, NameMapping nameMapping) + { + try (ParquetDataSource dataSource = new TrinoParquetDataSource(file, new ParquetReaderOptions(), new FileFormatDataSourceStats())) { + ParquetMetadata metadata = MetadataReader.readFooter(dataSource, Optional.empty()); + return ParquetUtil.footerMetrics(metadata, Stream.empty(), metricsConfig, nameMapping); + } + catch (IOException e) { + throw new UncheckedIOException("Failed to read file footer: " + file.location(), e); + } + } + + public static void addFiles( + ConnectorSession session, + TrinoFileSystem fileSystem, + TrinoCatalog catalog, + SchemaTableName targetName, + String location, + HiveStorageFormat format, + RecursiveDirectory recursiveDirectory) + { + Table table = catalog.loadTable(session, targetName); + PartitionSpec partitionSpec = table.spec(); + + checkProcedureArgument(partitionSpec.isUnpartitioned(), "The procedure does not support partitioned tables"); + + try { + List dataFiles = buildDataFilesFromLocation(fileSystem, recursiveDirectory, format, location, partitionSpec, Optional.empty(), table.schema()); + addFiles(session, table, dataFiles); + } + catch (Exception e) { + throw new TrinoException(ICEBERG_COMMIT_ERROR, "Failed to add files: " + firstNonNull(e.getMessage(), e), e); + } + } + + private static List buildDataFilesFromLocation( + TrinoFileSystem fileSystem, + RecursiveDirectory recursive, + HiveStorageFormat format, + String location, + PartitionSpec partitionSpec, + Optional partition, + Schema schema) + throws IOException + { + if (fileSystem.directoryExists(Location.of(location)).orElse(false)) { + return MigrationUtils.buildDataFiles(fileSystem, recursive, format, location, partitionSpec, partition, schema); + } + + TrinoInputFile file = fileSystem.newInputFile(Location.of(location)); + if (file.exists()) { + Metrics metrics = loadMetrics(file, format, schema); + return ImmutableList.of(buildDataFile(file.location().toString(), file.length(), partition, partitionSpec, format.name(), metrics)); + } + + throw new TrinoException(NOT_FOUND, "Location not found: " + location); + } + + public static void addFilesFromTable( + ConnectorSession session, + TrinoFileSystem fileSystem, + HiveMetastoreFactory metastoreFactory, + Table targetTable, + io.trino.plugin.hive.metastore.Table sourceTable, + Map partitionFilter, + RecursiveDirectory recursiveDirectory) + { + HiveMetastore metastore = metastoreFactory.createMetastore(Optional.of(session.getIdentity())); + + PartitionSpec partitionSpec = targetTable.spec(); + Schema schema = targetTable.schema(); + NameMapping nameMapping = MappingUtil.create(schema); + + HiveStorageFormat storageFormat = extractHiveStorageFormat(sourceTable.getStorage().getStorageFormat()); + String location = sourceTable.getStorage().getLocation(); + + try { + ImmutableList.Builder dataFilesBuilder = ImmutableList.builder(); + if (partitionSpec.isUnpartitioned()) { + log.debug("Building data files from %s", location); + dataFilesBuilder.addAll(buildDataFiles(fileSystem, recursiveDirectory, storageFormat, location, partitionSpec, Optional.empty(), schema)); + } + else { + List partitionNames = partitionFilter == null ? ImmutableList.of() : ImmutableList.of(PARTITION_JOINER.join(partitionFilter)); + Map> partitions = metastore.getPartitionsByNames(sourceTable, partitionNames); + for (Map.Entry> partition : partitions.entrySet()) { + Storage storage = partition.getValue().orElseThrow(() -> new IllegalArgumentException("Invalid partition: " + partition.getKey())).getStorage(); + log.debug("Building data files from partition: %s", partition); + HiveStorageFormat partitionStorageFormat = extractHiveStorageFormat(storage.getStorageFormat()); + StructLike partitionData = DataFiles.data(partitionSpec, partition.getKey()); + dataFilesBuilder.addAll(buildDataFiles(fileSystem, recursiveDirectory, partitionStorageFormat, storage.getLocation(), partitionSpec, Optional.of(partitionData), schema)); + } + } + + log.debug("Start new transaction"); + Transaction transaction = targetTable.newTransaction(); + if (!targetTable.properties().containsKey(DEFAULT_NAME_MAPPING)) { + log.debug("Update default name mapping property"); + transaction.updateProperties() + .set(DEFAULT_NAME_MAPPING, toJson(nameMapping)) + .commit(); + } + addFiles(session, targetTable, dataFilesBuilder.build()); + } + catch (Exception e) { + throw new TrinoException(ICEBERG_COMMIT_ERROR, "Failed to add files: " + firstNonNull(e.getMessage(), e), e); + } + } + + public static DataFile buildDataFile(String path, long length, Optional partition, PartitionSpec spec, String format, Metrics metrics) + { + DataFiles.Builder dataFile = DataFiles.builder(spec) + .withPath(path) + .withFormat(format) + .withFileSizeInBytes(length) + .withMetrics(metrics); + partition.ifPresent(dataFile::withPartition); + return dataFile.build(); + } + + public static void addFiles(ConnectorSession session, Table table, List dataFiles) + { + Schema schema = table.schema(); + Set requiredFields = schema.columns().stream() + .filter(Types.NestedField::isRequired) + .map(Types.NestedField::fieldId) + .collect(toImmutableSet()); + + ImmutableSet.Builder existingFilesBuilder = ImmutableSet.builder(); + try (CloseableIterable iterator = table.newScan().planFiles()) { + for (FileScanTask fileScanTask : iterator) { + DataFile dataFile = fileScanTask.file(); + existingFilesBuilder.add(dataFile.location()); + } + } + catch (IOException e) { + throw new UncheckedIOException(e); + } + Set existingFiles = existingFilesBuilder.build(); + + if (!requiredFields.isEmpty()) { + for (DataFile dataFile : dataFiles) { + Map nullValueCounts = firstNonNull(dataFile.nullValueCounts(), Map.of()); + for (Integer field : requiredFields) { + Long nullCount = nullValueCounts.get(field); + if (nullCount == null || nullCount > 0) { + throw new TrinoException(CONSTRAINT_VIOLATION, "NULL value not allowed for NOT NULL column: " + schema.findField(field).name()); + } + } + } + } + + try { + log.debug("Start new transaction"); + Transaction transaction = table.newTransaction(); + if (!table.properties().containsKey(DEFAULT_NAME_MAPPING)) { + log.debug("Update default name mapping property"); + transaction.updateProperties() + .set(DEFAULT_NAME_MAPPING, toJson(MappingUtil.create(schema))) + .commit(); + } + log.debug("Append data %d data files", dataFiles.size()); + AppendFiles appendFiles = isMergeManifestsOnWrite(session) ? transaction.newAppend() : transaction.newFastAppend(); + for (DataFile dataFile : dataFiles) { + if (existingFiles.contains(dataFile.location())) { + throw new TrinoException(ALREADY_EXISTS, "File already exists: " + dataFile.location()); + } + appendFiles.appendFile(dataFile); + } + appendFiles.commit(); + transaction.commitTransaction(); + log.debug("Successfully added files to %s table", table.name()); + } + catch (Exception e) { + throw new TrinoException(ICEBERG_COMMIT_ERROR, "Failed to add files: " + firstNonNull(e.getMessage(), e), e); + } + } +} diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/procedure/OptimizeManifestsTableProcedure.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/procedure/OptimizeManifestsTableProcedure.java new file mode 100644 index 000000000000..62604a27fa80 --- /dev/null +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/procedure/OptimizeManifestsTableProcedure.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.iceberg.procedure; + +import com.google.common.collect.ImmutableList; +import com.google.inject.Provider; +import io.trino.spi.connector.TableProcedureMetadata; + +import static io.trino.plugin.iceberg.procedure.IcebergTableProcedureId.OPTIMIZE_MANIFESTS; +import static io.trino.spi.connector.TableProcedureExecutionMode.coordinatorOnly; + +public class OptimizeManifestsTableProcedure + implements Provider +{ + @Override + public TableProcedureMetadata get() + { + return new TableProcedureMetadata( + OPTIMIZE_MANIFESTS.name(), + coordinatorOnly(), + ImmutableList.of()); + } +} diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/RollbackToSnapshotProcedure.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/procedure/RollbackToSnapshotProcedure.java similarity index 96% rename from plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/RollbackToSnapshotProcedure.java rename to plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/procedure/RollbackToSnapshotProcedure.java index 7a64bd520ffb..5f0e5b7ae5a5 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/RollbackToSnapshotProcedure.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/procedure/RollbackToSnapshotProcedure.java @@ -11,7 +11,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package io.trino.plugin.iceberg; +package io.trino.plugin.iceberg.procedure; import com.google.common.collect.ImmutableList; import com.google.inject.Inject; @@ -34,7 +34,7 @@ public class RollbackToSnapshotProcedure implements Provider { - private static final MethodHandle ROLLBACK_TO_SNAPSHOT; + public static final MethodHandle ROLLBACK_TO_SNAPSHOT; static { try { diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/util/DefaultLocationProvider.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/util/DefaultLocationProvider.java new file mode 100644 index 000000000000..a63b3ca54066 --- /dev/null +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/util/DefaultLocationProvider.java @@ -0,0 +1,61 @@ +/* + * 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.iceberg.util; + +import org.apache.iceberg.PartitionSpec; +import org.apache.iceberg.StructLike; +import org.apache.iceberg.TableProperties; +import org.apache.iceberg.io.LocationProvider; + +import java.util.Map; + +import static java.lang.String.format; +import static org.apache.iceberg.util.LocationUtil.stripTrailingSlash; + +// based on org.apache.iceberg.LocationProviders.DefaultLocationProvider +public class DefaultLocationProvider + implements LocationProvider +{ + private final String dataLocation; + + public DefaultLocationProvider(String tableLocation, Map properties) + { + this.dataLocation = stripTrailingSlash(dataLocation(properties, tableLocation)); + } + + @SuppressWarnings("deprecation") + private static String dataLocation(Map properties, String tableLocation) + { + String dataLocation = properties.get(TableProperties.WRITE_DATA_LOCATION); + if (dataLocation == null) { + dataLocation = properties.get(TableProperties.WRITE_FOLDER_STORAGE_LOCATION); + if (dataLocation == null) { + dataLocation = format("%s/data", stripTrailingSlash(tableLocation)); + } + } + return dataLocation; + } + + @Override + public String newDataLocation(PartitionSpec spec, StructLike partitionData, String filename) + { + return "%s/%s/%s".formatted(dataLocation, spec.partitionToPath(partitionData), filename); + } + + @Override + public String newDataLocation(String filename) + { + return "%s/%s".formatted(dataLocation, filename); + } +} diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/util/HiveSchemaUtil.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/util/HiveSchemaUtil.java index 7ff436f7df13..3f8bacde2a2b 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/util/HiveSchemaUtil.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/util/HiveSchemaUtil.java @@ -14,10 +14,12 @@ package io.trino.plugin.iceberg.util; import io.trino.plugin.hive.type.TypeInfo; +import io.trino.spi.TrinoException; import org.apache.iceberg.types.Type; import org.apache.iceberg.types.Types.DecimalType; import static io.trino.plugin.hive.type.TypeInfoUtils.getTypeInfoFromTypeString; +import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED; import static java.util.stream.Collectors.joining; // based on org.apache.iceberg.hive.HiveSchemaUtil @@ -41,8 +43,12 @@ private static String convertToTypeString(Type type) case DATE -> "date"; case TIME, STRING, UUID -> "string"; case TIMESTAMP -> "timestamp"; + case TIMESTAMP_NANO -> throw new TrinoException(NOT_SUPPORTED, "Unsupported Iceberg type: TIMESTAMP_NANO"); case FIXED, BINARY -> "binary"; case DECIMAL -> "decimal(%s,%s)".formatted(((DecimalType) type).precision(), ((DecimalType) type).scale()); + case UNKNOWN -> throw new TrinoException(NOT_SUPPORTED, "Unsupported Iceberg type: UNKNOWN"); + // TODO https://github.com/trinodb/trino/issues/24538 Support variant type + case VARIANT -> throw new TrinoException(NOT_SUPPORTED, "Unsupported Iceberg type: VARIANT"); case LIST -> "array<%s>".formatted(convert(type.asListType().elementType())); case MAP -> "map<%s,%s>".formatted(convert(type.asMapType().keyType()), convert(type.asMapType().valueType())); case STRUCT -> "struct<%s>".formatted(type.asStructType().fields().stream() diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/util/ObjectStoreLocationProvider.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/util/ObjectStoreLocationProvider.java new file mode 100644 index 000000000000..fb552c381d87 --- /dev/null +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/util/ObjectStoreLocationProvider.java @@ -0,0 +1,105 @@ +/* + * 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.iceberg.util; + +import com.google.common.hash.HashFunction; +import io.trino.filesystem.Location; +import org.apache.iceberg.PartitionSpec; +import org.apache.iceberg.StructLike; +import org.apache.iceberg.TableProperties; +import org.apache.iceberg.io.LocationProvider; + +import java.util.Base64; +import java.util.Map; + +import static com.google.common.hash.Hashing.murmur3_32_fixed; +import static java.nio.charset.StandardCharsets.UTF_8; +import static org.apache.iceberg.util.LocationUtil.stripTrailingSlash; + +// based on org.apache.iceberg.LocationProviders.ObjectStoreLocationProvider +public class ObjectStoreLocationProvider + implements LocationProvider +{ + private static final HashFunction HASH_FUNC = murmur3_32_fixed(); + private static final Base64.Encoder BASE64_ENCODER = Base64.getUrlEncoder().withoutPadding(); + private final String storageLocation; + private final String context; + + public ObjectStoreLocationProvider(String tableLocation, Map properties) + { + this.storageLocation = stripTrailingSlash(dataLocation(properties, tableLocation)); + // if the storage location is within the table prefix, don't add table and database name context + this.context = storageLocation.startsWith(stripTrailingSlash(tableLocation)) ? null : pathContext(tableLocation); + } + + @SuppressWarnings("deprecation") + private static String dataLocation(Map properties, String tableLocation) + { + String dataLocation = properties.get(TableProperties.WRITE_DATA_LOCATION); + if (dataLocation == null) { + dataLocation = properties.get(TableProperties.OBJECT_STORE_PATH); + if (dataLocation == null) { + dataLocation = properties.get(TableProperties.WRITE_FOLDER_STORAGE_LOCATION); + if (dataLocation == null) { + dataLocation = "%s/data".formatted(stripTrailingSlash(tableLocation)); + } + } + } + return dataLocation; + } + + @Override + public String newDataLocation(PartitionSpec spec, StructLike partitionData, String filename) + { + return newDataLocation("%s/%s".formatted(spec.partitionToPath(partitionData), filename)); + } + + @Override + public String newDataLocation(String filename) + { + String hash = computeHash(filename); + if (context != null) { + return "%s/%s/%s/%s".formatted(storageLocation, hash, context, filename); + } + return "%s/%s/%s".formatted(storageLocation, hash, filename); + } + + private static String pathContext(String tableLocation) + { + Location location; + String name; + try { + location = Location.of(stripTrailingSlash(tableLocation)); + name = location.fileName(); + } + catch (IllegalArgumentException | IllegalStateException e) { + return null; + } + + try { + String parent = stripTrailingSlash(location.parentDirectory().path()); + parent = parent.substring(parent.lastIndexOf('/') + 1); + return "%s/%s".formatted(parent, name); + } + catch (IllegalArgumentException | IllegalStateException e) { + return name; + } + } + + private static String computeHash(String fileName) + { + byte[] bytes = HASH_FUNC.hashString(fileName, UTF_8).asBytes(); + return BASE64_ENCODER.encodeToString(bytes); + } +} diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/util/OrcIcebergIds.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/util/OrcIcebergIds.java new file mode 100644 index 000000000000..3d805e33494d --- /dev/null +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/util/OrcIcebergIds.java @@ -0,0 +1,107 @@ +/* + * 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.iceberg.util; + +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.common.graph.Traverser; +import io.trino.orc.OrcColumn; +import io.trino.orc.OrcReader; +import io.trino.orc.metadata.OrcType.OrcTypeKind; +import org.apache.iceberg.mapping.MappedField; +import org.apache.iceberg.mapping.NameMapping; + +import java.util.List; +import java.util.Map; +import java.util.Optional; + +import static com.google.common.collect.ImmutableList.toImmutableList; +import static io.trino.plugin.iceberg.TypeConverter.ORC_ICEBERG_ID_KEY; + +public final class OrcIcebergIds +{ + private OrcIcebergIds() {} + + public static Map fileColumnsByIcebergId(OrcReader reader, Optional nameMapping) + { + List fileColumns = reader.getRootColumn().getNestedColumns(); + + if (nameMapping.isPresent() && !hasIds(reader.getRootColumn())) { + fileColumns = fileColumns.stream() + .map(orcColumn -> setMissingFieldIds(orcColumn, nameMapping.get(), ImmutableList.of(orcColumn.getColumnName()))) + .collect(toImmutableList()); + } + + return mapIdsToOrcFileColumns(fileColumns); + } + + private static boolean hasIds(OrcColumn column) + { + if (column.getAttributes().containsKey(ORC_ICEBERG_ID_KEY)) { + return true; + } + + return column.getNestedColumns().stream().anyMatch(OrcIcebergIds::hasIds); + } + + private static OrcColumn setMissingFieldIds(OrcColumn column, NameMapping nameMapping, List qualifiedPath) + { + MappedField mappedField = nameMapping.find(qualifiedPath); + + ImmutableMap.Builder attributes = ImmutableMap.builder(); + attributes.putAll(column.getAttributes()); + if ((mappedField != null) && (mappedField.id() != null)) { + attributes.put(ORC_ICEBERG_ID_KEY, String.valueOf(mappedField.id())); + } + + List orcColumns = column.getNestedColumns().stream() + .map(nestedColumn -> setMissingFieldIds(nestedColumn, nameMapping, ImmutableList.builder() + .addAll(qualifiedPath) + .add(pathName(column, nestedColumn)) + .build())) + .collect(toImmutableList()); + + return new OrcColumn( + column.getPath(), + column.getColumnId(), + column.getColumnName(), + column.getColumnType(), + column.getOrcDataSourceId(), + orcColumns, + attributes.buildOrThrow()); + } + + private static String pathName(OrcColumn column, OrcColumn nestedColumn) + { + // Trino ORC reader uses "item" for list element names, but NameMapper expects "element" + if (column.getColumnType() == OrcTypeKind.LIST) { + return "element"; + } + return nestedColumn.getColumnName(); + } + + private static Map mapIdsToOrcFileColumns(List columns) + { + ImmutableMap.Builder columnsById = ImmutableMap.builder(); + Traverser.forTree(OrcColumn::getNestedColumns) + .depthFirstPreOrder(columns) + .forEach(column -> { + String fieldId = column.getAttributes().get(ORC_ICEBERG_ID_KEY); + if (fieldId != null) { + columnsById.put(Integer.parseInt(fieldId), column); + } + }); + return columnsById.buildOrThrow(); + } +} diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/util/OrcMetrics.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/util/OrcMetrics.java new file mode 100644 index 000000000000..ba6fcea929a0 --- /dev/null +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/util/OrcMetrics.java @@ -0,0 +1,351 @@ +/* + * 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.iceberg.util; + +import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableSet; +import io.airlift.slice.Slice; +import io.trino.filesystem.TrinoInputFile; +import io.trino.orc.OrcColumn; +import io.trino.orc.OrcDataSource; +import io.trino.orc.OrcReader; +import io.trino.orc.OrcReaderOptions; +import io.trino.orc.metadata.ColumnMetadata; +import io.trino.orc.metadata.Footer; +import io.trino.orc.metadata.OrcColumnId; +import io.trino.orc.metadata.OrcType; +import io.trino.orc.metadata.statistics.BooleanStatistics; +import io.trino.orc.metadata.statistics.ColumnStatistics; +import io.trino.orc.metadata.statistics.DateStatistics; +import io.trino.orc.metadata.statistics.DecimalStatistics; +import io.trino.orc.metadata.statistics.DoubleStatistics; +import io.trino.orc.metadata.statistics.IntegerStatistics; +import io.trino.orc.metadata.statistics.StringStatistics; +import io.trino.orc.metadata.statistics.TimestampStatistics; +import io.trino.plugin.hive.FileFormatDataSourceStats; +import io.trino.plugin.iceberg.TrinoOrcDataSource; +import org.apache.iceberg.Metrics; +import org.apache.iceberg.MetricsConfig; +import org.apache.iceberg.MetricsModes; +import org.apache.iceberg.MetricsUtil; +import org.apache.iceberg.Schema; +import org.apache.iceberg.expressions.Literal; +import org.apache.iceberg.mapping.MappingUtil; +import org.apache.iceberg.mapping.NameMapping; +import org.apache.iceberg.types.Conversions; +import org.apache.iceberg.types.Type; +import org.apache.iceberg.types.Type.TypeID; +import org.apache.iceberg.types.Types; +import org.apache.iceberg.util.BinaryUtil; +import org.apache.iceberg.util.UnicodeUtil; + +import java.io.IOException; +import java.io.UncheckedIOException; +import java.math.BigDecimal; +import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; + +import static com.google.common.base.Verify.verify; +import static com.google.common.collect.ImmutableList.toImmutableList; +import static com.google.common.collect.ImmutableMap.toImmutableMap; +import static io.trino.orc.OrcReader.createOrcReader; +import static io.trino.orc.metadata.OrcColumnId.ROOT_COLUMN; +import static io.trino.plugin.iceberg.TypeConverter.ORC_ICEBERG_ID_KEY; +import static io.trino.plugin.iceberg.util.OrcIcebergIds.fileColumnsByIcebergId; +import static io.trino.spi.type.Timestamps.MICROSECONDS_PER_MILLISECOND; +import static java.lang.Math.toIntExact; +import static java.math.RoundingMode.UNNECESSARY; +import static java.util.function.Function.identity; + +public final class OrcMetrics +{ + private OrcMetrics() {} + + public static Metrics fileMetrics(TrinoInputFile file, MetricsConfig metricsConfig, Schema schema) + { + OrcReaderOptions options = new OrcReaderOptions(); + try (OrcDataSource dataSource = new TrinoOrcDataSource(file, options, new FileFormatDataSourceStats())) { + Optional reader = createOrcReader(dataSource, options); + if (reader.isEmpty()) { + return new Metrics(0L, null, null, null, null); + } + Footer footer = reader.get().getFooter(); + + // use name mapping to compute missing Iceberg field IDs + Optional nameMapping = Optional.of(MappingUtil.create(schema)); + Map mappedColumns = fileColumnsByIcebergId(reader.get(), nameMapping) + .values().stream() + .collect(toImmutableMap(OrcColumn::getColumnId, identity())); + + // rebuild type list with mapped columns + List mappedTypes = new ArrayList<>(); + ColumnMetadata types = footer.getTypes(); + for (int i = 0; i < types.size(); i++) { + OrcColumnId id = new OrcColumnId(i); + mappedTypes.add(Optional.ofNullable(mappedColumns.get(id)) + .map(OrcMetrics::toBasicOrcType) + .orElseGet(() -> types.get(id))); + } + + return computeMetrics(metricsConfig, schema, new ColumnMetadata<>(mappedTypes), footer.getNumberOfRows(), footer.getFileStats()); + } + catch (IOException e) { + throw new UncheckedIOException("Failed to read file footer: " + file.location(), e); + } + } + + private static OrcType toBasicOrcType(OrcColumn column) + { + return new OrcType( + column.getColumnType(), + column.getNestedColumns().stream() + .map(OrcColumn::getColumnId) + .collect(toImmutableList()), + null, + Optional.empty(), + Optional.empty(), + Optional.empty(), + column.getAttributes()); + } + + public static Metrics computeMetrics( + MetricsConfig metricsConfig, + Schema icebergSchema, + ColumnMetadata orcColumns, + long fileRowCount, + Optional> columnStatistics) + { + if (columnStatistics.isEmpty()) { + return new Metrics(fileRowCount, null, null, null, null, null, null); + } + // Columns that are descendants of LIST or MAP types are excluded because: + // 1. Their stats are not used by Apache Iceberg to filter out data files + // 2. Their record count can be larger than table-level row count. There's no good way to calculate nullCounts for them. + // See https://github.com/apache/iceberg/pull/199#discussion_r429443627 + Set excludedColumns = getExcludedColumns(orcColumns); + + ImmutableMap.Builder valueCountsBuilder = ImmutableMap.builder(); + ImmutableMap.Builder nullCountsBuilder = ImmutableMap.builder(); + ImmutableMap.Builder nanCountsBuilder = ImmutableMap.builder(); + ImmutableMap.Builder lowerBoundsBuilder = ImmutableMap.builder(); + ImmutableMap.Builder upperBoundsBuilder = ImmutableMap.builder(); + + // OrcColumnId(0) is the root column that represents file-level schema + for (int i = 1; i < orcColumns.size(); i++) { + OrcColumnId orcColumnId = new OrcColumnId(i); + if (excludedColumns.contains(orcColumnId)) { + continue; + } + OrcType orcColumn = orcColumns.get(orcColumnId); + ColumnStatistics orcColumnStats = columnStatistics.get().get(orcColumnId); + int icebergId = getIcebergId(orcColumn); + Types.NestedField icebergField = icebergSchema.findField(icebergId); + MetricsModes.MetricsMode metricsMode = MetricsUtil.metricsMode(icebergSchema, metricsConfig, icebergId); + if (metricsMode.equals(MetricsModes.None.get())) { + continue; + } + verify(icebergField != null, "Cannot find Iceberg column with ID %s in schema %s", icebergId, icebergSchema); + valueCountsBuilder.put(icebergId, fileRowCount); + if (orcColumnStats.hasNumberOfValues()) { + nullCountsBuilder.put(icebergId, fileRowCount - orcColumnStats.getNumberOfValues()); + } + if (orcColumnStats.getNumberOfNanValues() > 0) { + nanCountsBuilder.put(icebergId, orcColumnStats.getNumberOfNanValues()); + } + + if (!metricsMode.equals(MetricsModes.Counts.get())) { + toIcebergMinMax(orcColumnStats, icebergField.type(), metricsMode).ifPresent(minMax -> { + lowerBoundsBuilder.put(icebergId, minMax.getMin()); + upperBoundsBuilder.put(icebergId, minMax.getMax()); + }); + } + } + Map valueCounts = valueCountsBuilder.buildOrThrow(); + Map nullCounts = nullCountsBuilder.buildOrThrow(); + Map nanCounts = nanCountsBuilder.buildOrThrow(); + Map lowerBounds = lowerBoundsBuilder.buildOrThrow(); + Map upperBounds = upperBoundsBuilder.buildOrThrow(); + return new Metrics( + fileRowCount, + null, // TODO: Add column size accounting to ORC column writers + valueCounts.isEmpty() ? null : valueCounts, + nullCounts.isEmpty() ? null : nullCounts, + nanCounts.isEmpty() ? null : nanCounts, + lowerBounds.isEmpty() ? null : lowerBounds, + upperBounds.isEmpty() ? null : upperBounds); + } + + private static Set getExcludedColumns(ColumnMetadata orcColumns) + { + ImmutableSet.Builder excludedColumns = ImmutableSet.builder(); + populateExcludedColumns(orcColumns, ROOT_COLUMN, false, excludedColumns); + return excludedColumns.build(); + } + + private static void populateExcludedColumns(ColumnMetadata orcColumns, OrcColumnId orcColumnId, boolean exclude, ImmutableSet.Builder excludedColumns) + { + if (exclude) { + excludedColumns.add(orcColumnId); + } + OrcType orcColumn = orcColumns.get(orcColumnId); + switch (orcColumn.getOrcTypeKind()) { + case LIST: + case MAP: + for (OrcColumnId child : orcColumn.getFieldTypeIndexes()) { + populateExcludedColumns(orcColumns, child, true, excludedColumns); + } + return; + case STRUCT: + for (OrcColumnId child : orcColumn.getFieldTypeIndexes()) { + populateExcludedColumns(orcColumns, child, exclude, excludedColumns); + } + return; + default: + // unexpected, TODO throw + } + } + + private static int getIcebergId(OrcType orcColumn) + { + String icebergId = orcColumn.getAttributes().get(ORC_ICEBERG_ID_KEY); + verify(icebergId != null, "ORC column %s doesn't have an associated Iceberg ID", orcColumn); + return Integer.parseInt(icebergId); + } + + private static Optional toIcebergMinMax(ColumnStatistics orcColumnStats, Type icebergType, MetricsModes.MetricsMode metricsModes) + { + BooleanStatistics booleanStatistics = orcColumnStats.getBooleanStatistics(); + if (booleanStatistics != null) { + boolean hasTrueValues = booleanStatistics.getTrueValueCount() != 0; + boolean hasFalseValues = orcColumnStats.getNumberOfValues() != booleanStatistics.getTrueValueCount(); + return Optional.of(new IcebergMinMax(icebergType, !hasFalseValues, hasTrueValues, metricsModes)); + } + + IntegerStatistics integerStatistics = orcColumnStats.getIntegerStatistics(); + if (integerStatistics != null) { + Object min = integerStatistics.getMin(); + Object max = integerStatistics.getMax(); + if (min == null || max == null) { + return Optional.empty(); + } + if (icebergType.typeId() == TypeID.INTEGER) { + min = toIntExact((Long) min); + max = toIntExact((Long) max); + } + return Optional.of(new IcebergMinMax(icebergType, min, max, metricsModes)); + } + DoubleStatistics doubleStatistics = orcColumnStats.getDoubleStatistics(); + if (doubleStatistics != null) { + Object min = doubleStatistics.getMin(); + Object max = doubleStatistics.getMax(); + if (min == null || max == null) { + return Optional.empty(); + } + if (icebergType.typeId() == TypeID.FLOAT) { + min = ((Double) min).floatValue(); + max = ((Double) max).floatValue(); + } + return Optional.of(new IcebergMinMax(icebergType, min, max, metricsModes)); + } + StringStatistics stringStatistics = orcColumnStats.getStringStatistics(); + if (stringStatistics != null) { + Slice min = stringStatistics.getMin(); + Slice max = stringStatistics.getMax(); + if (min == null || max == null) { + return Optional.empty(); + } + return Optional.of(new IcebergMinMax(icebergType, min.toStringUtf8(), max.toStringUtf8(), metricsModes)); + } + DateStatistics dateStatistics = orcColumnStats.getDateStatistics(); + if (dateStatistics != null) { + Integer min = dateStatistics.getMin(); + Integer max = dateStatistics.getMax(); + if (min == null || max == null) { + return Optional.empty(); + } + return Optional.of(new IcebergMinMax(icebergType, min, max, metricsModes)); + } + DecimalStatistics decimalStatistics = orcColumnStats.getDecimalStatistics(); + if (decimalStatistics != null) { + BigDecimal min = decimalStatistics.getMin(); + BigDecimal max = decimalStatistics.getMax(); + if (min == null || max == null) { + return Optional.empty(); + } + min = min.setScale(((Types.DecimalType) icebergType).scale(), UNNECESSARY); + max = max.setScale(((Types.DecimalType) icebergType).scale(), UNNECESSARY); + return Optional.of(new IcebergMinMax(icebergType, min, max, metricsModes)); + } + TimestampStatistics timestampStatistics = orcColumnStats.getTimestampStatistics(); + if (timestampStatistics != null) { + Long min = timestampStatistics.getMin(); + Long max = timestampStatistics.getMax(); + if (min == null || max == null) { + return Optional.empty(); + } + // Since ORC timestamp statistics are truncated to millisecond precision, this can cause some column values to fall outside the stats range. + // We are appending 999 microseconds to account for the fact that Trino ORC writer truncates timestamps. + return Optional.of(new IcebergMinMax(icebergType, min * MICROSECONDS_PER_MILLISECOND, (max * MICROSECONDS_PER_MILLISECOND) + (MICROSECONDS_PER_MILLISECOND - 1), metricsModes)); + } + return Optional.empty(); + } + + private static class IcebergMinMax + { + private final ByteBuffer min; + private final ByteBuffer max; + + private IcebergMinMax(Type type, Object min, Object max, MetricsModes.MetricsMode metricsMode) + { + if (metricsMode instanceof MetricsModes.Full) { + this.min = Conversions.toByteBuffer(type, min); + this.max = Conversions.toByteBuffer(type, max); + } + else if (metricsMode instanceof MetricsModes.Truncate truncateMode) { + int truncateLength = truncateMode.length(); + switch (type.typeId()) { + case STRING: + this.min = UnicodeUtil.truncateStringMin(Literal.of((CharSequence) min), truncateLength).toByteBuffer(); + this.max = UnicodeUtil.truncateStringMax(Literal.of((CharSequence) max), truncateLength).toByteBuffer(); + break; + case FIXED: + case BINARY: + this.min = BinaryUtil.truncateBinaryMin(Literal.of((ByteBuffer) min), truncateLength).toByteBuffer(); + this.max = BinaryUtil.truncateBinaryMax(Literal.of((ByteBuffer) max), truncateLength).toByteBuffer(); + break; + default: + this.min = Conversions.toByteBuffer(type, min); + this.max = Conversions.toByteBuffer(type, max); + } + } + else { + throw new UnsupportedOperationException("Unsupported metrics mode for Iceberg Max/Min Bound: " + metricsMode); + } + } + + public ByteBuffer getMin() + { + return min; + } + + public ByteBuffer getMax() + { + return max; + } + } +} diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/util/OrcTypeConverter.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/util/OrcTypeConverter.java new file mode 100644 index 000000000000..d1f1250372d5 --- /dev/null +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/util/OrcTypeConverter.java @@ -0,0 +1,179 @@ +/* + * 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.iceberg.util; + +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import io.trino.orc.metadata.ColumnMetadata; +import io.trino.orc.metadata.OrcColumnId; +import io.trino.orc.metadata.OrcType; +import io.trino.orc.metadata.OrcType.OrcTypeKind; +import io.trino.spi.TrinoException; +import org.apache.iceberg.Schema; +import org.apache.iceberg.types.Type; +import org.apache.iceberg.types.Types.DecimalType; +import org.apache.iceberg.types.Types.ListType; +import org.apache.iceberg.types.Types.MapType; +import org.apache.iceberg.types.Types.NestedField; +import org.apache.iceberg.types.Types.StructType; +import org.apache.iceberg.types.Types.TimestampType; + +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.Optional; + +import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED; + +public final class OrcTypeConverter +{ + public static final String ORC_ICEBERG_ID_KEY = "iceberg.id"; + public static final String ORC_ICEBERG_REQUIRED_KEY = "iceberg.required"; + public static final String ICEBERG_LONG_TYPE = "iceberg.long-type"; + public static final String ICEBERG_BINARY_TYPE = "iceberg.binary-type"; + + private OrcTypeConverter() {} + + public static ColumnMetadata toOrcType(Schema schema) + { + return new ColumnMetadata<>(toOrcStructType(0, schema.asStruct(), ImmutableMap.of())); + } + + private static List toOrcType(int nextFieldTypeIndex, Type type, Map attributes) + { + return switch (type.typeId()) { + case BOOLEAN -> ImmutableList.of(new OrcType(OrcTypeKind.BOOLEAN, ImmutableList.of(), ImmutableList.of(), Optional.empty(), Optional.empty(), Optional.empty(), attributes)); + case INTEGER -> ImmutableList.of(new OrcType(OrcTypeKind.INT, ImmutableList.of(), ImmutableList.of(), Optional.empty(), Optional.empty(), Optional.empty(), attributes)); + case LONG -> ImmutableList.of(new OrcType(OrcTypeKind.LONG, ImmutableList.of(), ImmutableList.of(), Optional.empty(), Optional.empty(), Optional.empty(), attributes)); + case FLOAT -> ImmutableList.of(new OrcType(OrcTypeKind.FLOAT, ImmutableList.of(), ImmutableList.of(), Optional.empty(), Optional.empty(), Optional.empty(), attributes)); + case DOUBLE -> ImmutableList.of(new OrcType(OrcTypeKind.DOUBLE, ImmutableList.of(), ImmutableList.of(), Optional.empty(), Optional.empty(), Optional.empty(), attributes)); + case DATE -> ImmutableList.of(new OrcType(OrcTypeKind.DATE, ImmutableList.of(), ImmutableList.of(), Optional.empty(), Optional.empty(), Optional.empty(), attributes)); + case TIME -> { + attributes = ImmutableMap.builder() + .putAll(attributes) + .put(ICEBERG_LONG_TYPE, "TIME") + .buildOrThrow(); + yield ImmutableList.of(new OrcType(OrcTypeKind.LONG, ImmutableList.of(), ImmutableList.of(), Optional.empty(), Optional.empty(), Optional.empty(), attributes)); + } + case TIMESTAMP -> { + OrcTypeKind timestampKind = ((TimestampType) type).shouldAdjustToUTC() ? OrcTypeKind.TIMESTAMP_INSTANT : OrcTypeKind.TIMESTAMP; + yield ImmutableList.of(new OrcType(timestampKind, ImmutableList.of(), ImmutableList.of(), Optional.empty(), Optional.empty(), Optional.empty(), attributes)); + } + // TODO https://github.com/trinodb/trino/issues/19753 Support Iceberg timestamp types with nanosecond precision + case TIMESTAMP_NANO -> throw new TrinoException(NOT_SUPPORTED, "Unsupported Iceberg type: TIMESTAMP_NANO"); + case STRING -> ImmutableList.of(new OrcType(OrcTypeKind.STRING, ImmutableList.of(), ImmutableList.of(), Optional.empty(), Optional.empty(), Optional.empty(), attributes)); + case FIXED, BINARY -> ImmutableList.of(new OrcType(OrcTypeKind.BINARY, ImmutableList.of(), ImmutableList.of(), Optional.empty(), Optional.empty(), Optional.empty(), attributes)); + case DECIMAL -> { + DecimalType decimalType = (DecimalType) type; + yield ImmutableList.of(new OrcType(OrcTypeKind.DECIMAL, ImmutableList.of(), ImmutableList.of(), Optional.empty(), Optional.of(decimalType.precision()), Optional.of(decimalType.scale()), attributes)); + } + case UUID -> { + attributes = ImmutableMap.builder() + .putAll(attributes) + .put(ICEBERG_BINARY_TYPE, "UUID") + .buildOrThrow(); + yield ImmutableList.of(new OrcType(OrcTypeKind.BINARY, ImmutableList.of(), ImmutableList.of(), Optional.empty(), Optional.empty(), Optional.empty(), attributes)); + } + case VARIANT -> throw new TrinoException(NOT_SUPPORTED, "Unsupported Iceberg type: VARIANT"); + case UNKNOWN -> throw new TrinoException(NOT_SUPPORTED, "Unsupported Iceberg type: UNKNOWN"); + case STRUCT -> toOrcStructType(nextFieldTypeIndex, (StructType) type, attributes); + case LIST -> toOrcListType(nextFieldTypeIndex, (ListType) type, attributes); + case MAP -> toOrcMapType(nextFieldTypeIndex, (MapType) type, attributes); + }; + } + + private static List toOrcStructType(int nextFieldTypeIndex, StructType structType, Map attributes) + { + nextFieldTypeIndex++; + + List fieldTypeIndexes = new ArrayList<>(); + List fieldNames = new ArrayList<>(); + List> fieldTypesList = new ArrayList<>(); + for (NestedField field : structType.fields()) { + fieldTypeIndexes.add(new OrcColumnId(nextFieldTypeIndex)); + fieldNames.add(field.name()); + Map fieldAttributes = ImmutableMap.builder() + .put(ORC_ICEBERG_ID_KEY, Integer.toString(field.fieldId())) + .put(ORC_ICEBERG_REQUIRED_KEY, Boolean.toString(field.isRequired())) + .buildOrThrow(); + List fieldOrcTypes = toOrcType(nextFieldTypeIndex, field.type(), fieldAttributes); + fieldTypesList.add(fieldOrcTypes); + nextFieldTypeIndex += fieldOrcTypes.size(); + } + + return ImmutableList.builder() + .add(new OrcType( + OrcTypeKind.STRUCT, + fieldTypeIndexes, + fieldNames, + Optional.empty(), + Optional.empty(), + Optional.empty(), + attributes)) + .addAll(fieldTypesList.stream().flatMap(List::stream).iterator()) + .build(); + } + + private static List toOrcListType(int nextFieldTypeIndex, ListType listType, Map attributes) + { + nextFieldTypeIndex++; + + Map elementAttributes = ImmutableMap.builder() + .put(ORC_ICEBERG_ID_KEY, Integer.toString(listType.elementId())) + .put(ORC_ICEBERG_REQUIRED_KEY, Boolean.toString(listType.isElementRequired())) + .buildOrThrow(); + List itemTypes = toOrcType(nextFieldTypeIndex, listType.elementType(), elementAttributes); + + return ImmutableList.builder() + .add(new OrcType( + OrcTypeKind.LIST, + ImmutableList.of(new OrcColumnId(nextFieldTypeIndex)), + ImmutableList.of("item"), + Optional.empty(), + Optional.empty(), + Optional.empty(), + attributes)) + .addAll(itemTypes) + .build(); + } + + private static List toOrcMapType(int nextFieldTypeIndex, MapType mapType, Map attributes) + { + nextFieldTypeIndex++; + + List keyTypes = toOrcType(nextFieldTypeIndex, mapType.keyType(), ImmutableMap.builder() + .put(ORC_ICEBERG_ID_KEY, Integer.toString(mapType.keyId())) + .put(ORC_ICEBERG_REQUIRED_KEY, Boolean.toString(true)) + .buildOrThrow()); + + Map valueAttributes = ImmutableMap.builder() + .put(ORC_ICEBERG_ID_KEY, Integer.toString(mapType.valueId())) + .put(ORC_ICEBERG_REQUIRED_KEY, Boolean.toString(mapType.isValueRequired())) + .buildOrThrow(); + List valueTypes = toOrcType(nextFieldTypeIndex + keyTypes.size(), mapType.valueType(), valueAttributes); + + return ImmutableList.builder() + .add(new OrcType( + OrcTypeKind.MAP, + ImmutableList.of(new OrcColumnId(nextFieldTypeIndex), new OrcColumnId(nextFieldTypeIndex + keyTypes.size())), + ImmutableList.of("key", "value"), + Optional.empty(), + Optional.empty(), + Optional.empty(), + attributes)) + .addAll(keyTypes) + .addAll(valueTypes) + .build(); + } +} diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/util/ParquetUtil.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/util/ParquetUtil.java new file mode 100644 index 000000000000..98f50940b419 --- /dev/null +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/util/ParquetUtil.java @@ -0,0 +1,392 @@ +/* + * 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.iceberg.util; + +import com.google.common.collect.ImmutableList; +import io.trino.parquet.ParquetCorruptionException; +import io.trino.parquet.metadata.BlockMetadata; +import io.trino.parquet.metadata.ColumnChunkMetadata; +import io.trino.parquet.metadata.ParquetMetadata; +import org.apache.iceberg.FieldMetrics; +import org.apache.iceberg.Metrics; +import org.apache.iceberg.MetricsConfig; +import org.apache.iceberg.MetricsModes; +import org.apache.iceberg.MetricsModes.MetricsMode; +import org.apache.iceberg.MetricsUtil; +import org.apache.iceberg.Schema; +import org.apache.iceberg.expressions.Literal; +import org.apache.iceberg.mapping.NameMapping; +import org.apache.iceberg.parquet.ParquetSchemaUtil; +import org.apache.iceberg.types.Conversions; +import org.apache.iceberg.types.Type; +import org.apache.iceberg.types.Types; +import org.apache.iceberg.util.BinaryUtil; +import org.apache.iceberg.util.UnicodeUtil; +import org.apache.parquet.column.statistics.Statistics; +import org.apache.parquet.hadoop.metadata.ColumnPath; +import org.apache.parquet.io.api.Binary; +import org.apache.parquet.schema.LogicalTypeAnnotation.DecimalLogicalTypeAnnotation; +import org.apache.parquet.schema.MessageType; +import org.apache.parquet.schema.PrimitiveType; + +import java.math.BigDecimal; +import java.math.BigInteger; +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; +import java.util.Iterator; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.UUID; +import java.util.function.Function; +import java.util.stream.Stream; + +import static java.util.Objects.requireNonNull; +import static java.util.function.Function.identity; +import static java.util.stream.Collectors.toMap; +import static org.apache.iceberg.MetricsUtil.createNanValueCounts; +import static org.apache.iceberg.parquet.ParquetUtil.extractTimestampInt96; + +public final class ParquetUtil +{ + // based on org.apache.iceberg.parquet.ParquetUtil and on org.apache.iceberg.parquet.ParquetConversions + private ParquetUtil() {} + + public static Metrics footerMetrics(ParquetMetadata metadata, Stream> fieldMetrics, MetricsConfig metricsConfig) + throws ParquetCorruptionException + { + return footerMetrics(metadata, fieldMetrics, metricsConfig, null); + } + + public static Metrics footerMetrics( + ParquetMetadata metadata, + Stream> fieldMetrics, + MetricsConfig metricsConfig, + NameMapping nameMapping) + throws ParquetCorruptionException + { + requireNonNull(fieldMetrics, "fieldMetrics should not be null"); + + long rowCount = 0; + Map columnSizes = new HashMap<>(); + Map valueCounts = new HashMap<>(); + Map nullValueCounts = new HashMap<>(); + Map> lowerBounds = new HashMap<>(); + Map> upperBounds = new HashMap<>(); + Set missingStats = new HashSet<>(); + + // ignore metrics for fields we failed to determine reliable IDs + MessageType parquetTypeWithIds = getParquetTypeWithIds(metadata, nameMapping); + Schema fileSchema = ParquetSchemaUtil.convertAndPrune(parquetTypeWithIds); + + Map> fieldMetricsMap = fieldMetrics.collect(toMap(FieldMetrics::id, identity())); + + List blocks = metadata.getBlocks(); + for (BlockMetadata block : blocks) { + rowCount += block.rowCount(); + for (ColumnChunkMetadata column : block.columns()) { + Integer fieldId = fileSchema.aliasToId(column.getPath().toDotString()); + if (fieldId == null) { + // fileSchema may contain a subset of columns present in the file + // as we prune columns we could not assign ids + continue; + } + + increment(columnSizes, fieldId, column.getTotalSize()); + + MetricsMode metricsMode = MetricsUtil.metricsMode(fileSchema, metricsConfig, fieldId); + if (metricsMode == MetricsModes.None.get()) { + continue; + } + increment(valueCounts, fieldId, column.getValueCount()); + + Statistics stats = column.getStatistics(); + if (stats != null && !stats.isEmpty()) { + increment(nullValueCounts, fieldId, stats.getNumNulls()); + + // when there are metrics gathered by Iceberg for a column, we should use those instead + // of the ones from Parquet + if (metricsMode != MetricsModes.Counts.get() && !fieldMetricsMap.containsKey(fieldId)) { + Types.NestedField field = fileSchema.findField(fieldId); + if (field != null && stats.hasNonNullValue() && shouldStoreBounds(column, fileSchema)) { + Literal min = fromParquetPrimitive(field.type(), column.getPrimitiveType(), stats.genericGetMin()); + updateMin(lowerBounds, fieldId, field.type(), min, metricsMode); + Literal max = fromParquetPrimitive(field.type(), column.getPrimitiveType(), stats.genericGetMax()); + updateMax(upperBounds, fieldId, field.type(), max, metricsMode); + } + } + } + else { + missingStats.add(fieldId); + } + } + } + + // discard accumulated values if any stats were missing + for (Integer fieldId : missingStats) { + nullValueCounts.remove(fieldId); + lowerBounds.remove(fieldId); + upperBounds.remove(fieldId); + } + + updateFromFieldMetrics(fieldMetricsMap, metricsConfig, fileSchema, lowerBounds, upperBounds); + + return new Metrics( + rowCount, + columnSizes, + valueCounts, + nullValueCounts, + createNanValueCounts(fieldMetricsMap.values().stream(), metricsConfig, fileSchema), + toBufferMap(fileSchema, lowerBounds), + toBufferMap(fileSchema, upperBounds)); + } + + public static List getSplitOffsets(ParquetMetadata metadata) + throws ParquetCorruptionException + { + List blocks = metadata.getBlocks(); + List splitOffsets = new ArrayList<>(blocks.size()); + for (BlockMetadata blockMetaData : blocks) { + splitOffsets.add(blockMetaData.getStartingPos()); + } + Collections.sort(splitOffsets); + return ImmutableList.copyOf(splitOffsets); + } + + private static void updateFromFieldMetrics( + Map> idToFieldMetricsMap, + MetricsConfig metricsConfig, + Schema schema, + Map> lowerBounds, + Map> upperBounds) + { + idToFieldMetricsMap + .entrySet() + .forEach( + entry -> { + int fieldId = entry.getKey(); + FieldMetrics metrics = entry.getValue(); + MetricsMode metricsMode = MetricsUtil.metricsMode(schema, metricsConfig, fieldId); + + // only check for MetricsModes.None, since we don't truncate float/double values. + if (metricsMode != MetricsModes.None.get()) { + if (!metrics.hasBounds()) { + lowerBounds.remove(fieldId); + upperBounds.remove(fieldId); + } + else if (metrics.upperBound() instanceof Float) { + lowerBounds.put(fieldId, Literal.of((Float) metrics.lowerBound())); + upperBounds.put(fieldId, Literal.of((Float) metrics.upperBound())); + } + else if (metrics.upperBound() instanceof Double) { + lowerBounds.put(fieldId, Literal.of((Double) metrics.lowerBound())); + upperBounds.put(fieldId, Literal.of((Double) metrics.upperBound())); + } + else { + throw new UnsupportedOperationException("Expected only float or double column metrics"); + } + } + }); + } + + private static MessageType getParquetTypeWithIds(ParquetMetadata metadata, NameMapping nameMapping) + { + MessageType type = metadata.getFileMetaData().getSchema(); + + if (ParquetSchemaUtil.hasIds(type)) { + return type; + } + + if (nameMapping != null) { + return ParquetSchemaUtil.applyNameMapping(type, nameMapping); + } + + return ParquetSchemaUtil.addFallbackIds(type); + } + + // we allow struct nesting, but not maps or arrays + private static boolean shouldStoreBounds(ColumnChunkMetadata column, Schema schema) + { + if (column.getPrimitiveType().getPrimitiveTypeName() == PrimitiveType.PrimitiveTypeName.INT96) { + // stats for INT96 are not reliable + return false; + } + + ColumnPath columnPath = column.getPath(); + Iterator pathIterator = columnPath.iterator(); + Type currentType = schema.asStruct(); + + while (pathIterator.hasNext()) { + if (currentType == null || !currentType.isStructType()) { + return false; + } + String fieldName = pathIterator.next(); + currentType = currentType.asStructType().fieldType(fieldName); + } + + return currentType != null && currentType.isPrimitiveType(); + } + + private static void increment(Map columns, int fieldId, long amount) + { + if (columns != null) { + if (columns.containsKey(fieldId)) { + columns.put(fieldId, columns.get(fieldId) + amount); + } + else { + columns.put(fieldId, amount); + } + } + } + + @SuppressWarnings("unchecked") + private static void updateMin( + Map> lowerBounds, + int id, + Type type, + Literal min, + MetricsMode metricsMode) + { + Literal currentMin = (Literal) lowerBounds.get(id); + if (currentMin == null || min.comparator().compare(min.value(), currentMin.value()) < 0) { + if (metricsMode == MetricsModes.Full.get()) { + lowerBounds.put(id, min); + } + else { + MetricsModes.Truncate truncateMode = (MetricsModes.Truncate) metricsMode; + int truncateLength = truncateMode.length(); + switch (type.typeId()) { + case STRING: + lowerBounds.put(id, UnicodeUtil.truncateStringMin((Literal) min, truncateLength)); + break; + case FIXED: + case BINARY: + lowerBounds.put(id, BinaryUtil.truncateBinaryMin((Literal) min, truncateLength)); + break; + default: + lowerBounds.put(id, min); + } + } + } + } + + @SuppressWarnings("unchecked") + private static void updateMax( + Map> upperBounds, + int id, + Type type, + Literal max, + MetricsMode metricsMode) + { + Literal currentMax = (Literal) upperBounds.get(id); + if (currentMax == null || max.comparator().compare(max.value(), currentMax.value()) > 0) { + if (metricsMode == MetricsModes.Full.get()) { + upperBounds.put(id, max); + } + else { + MetricsModes.Truncate truncateMode = (MetricsModes.Truncate) metricsMode; + int truncateLength = truncateMode.length(); + switch (type.typeId()) { + case STRING: + Literal truncatedMaxString = UnicodeUtil.truncateStringMax((Literal) max, truncateLength); + if (truncatedMaxString != null) { + upperBounds.put(id, truncatedMaxString); + } + break; + case FIXED: + case BINARY: + Literal truncatedMaxBinary = BinaryUtil.truncateBinaryMax((Literal) max, truncateLength); + if (truncatedMaxBinary != null) { + upperBounds.put(id, truncatedMaxBinary); + } + break; + default: + upperBounds.put(id, max); + } + } + } + } + + private static Map toBufferMap(Schema schema, Map> map) + { + Map bufferMap = new HashMap<>(); + for (Map.Entry> entry : map.entrySet()) { + bufferMap.put( + entry.getKey(), + Conversions.toByteBuffer(schema.findType(entry.getKey()), entry.getValue().value())); + } + return bufferMap; + } + + @SuppressWarnings("unchecked") + public static Literal fromParquetPrimitive(Type type, PrimitiveType parquetType, Object value) + { + return switch (type.typeId()) { + case BOOLEAN -> (Literal) Literal.of((Boolean) value); + case INTEGER, DATE -> (Literal) Literal.of((Integer) value); + case LONG, TIME, TIMESTAMP -> (Literal) Literal.of((Long) value); + case FLOAT -> (Literal) Literal.of((Float) value); + case DOUBLE -> (Literal) Literal.of((Double) value); + case STRING -> { + Function stringConversion = converterFromParquet(parquetType); + yield (Literal) Literal.of((CharSequence) stringConversion.apply(value)); + } + case UUID -> { + Function uuidConversion = converterFromParquet(parquetType); + yield (Literal) Literal.of((UUID) uuidConversion.apply(value)); + } + case FIXED, BINARY -> { + Function binaryConversion = converterFromParquet(parquetType); + yield (Literal) Literal.of((ByteBuffer) binaryConversion.apply(value)); + } + case DECIMAL -> { + Function decimalConversion = converterFromParquet(parquetType); + yield (Literal) Literal.of((BigDecimal) decimalConversion.apply(value)); + } + default -> throw new IllegalArgumentException("Unsupported primitive type: " + type); + }; + } + + static Function converterFromParquet(PrimitiveType type) + { + if (type.getOriginalType() != null) { + switch (type.getOriginalType()) { + case UTF8: + // decode to CharSequence to avoid copying into a new String + return binary -> StandardCharsets.UTF_8.decode(((Binary) binary).toByteBuffer()); + case DECIMAL: + DecimalLogicalTypeAnnotation decimal = (DecimalLogicalTypeAnnotation) type.getLogicalTypeAnnotation(); + int scale = decimal.getScale(); + return switch (type.getPrimitiveTypeName()) { + case INT32, INT64 -> number -> BigDecimal.valueOf(((Number) number).longValue(), scale); + case FIXED_LEN_BYTE_ARRAY, BINARY -> binary -> new BigDecimal(new BigInteger(((Binary) binary).getBytes()), scale); + default -> throw new IllegalArgumentException("Unsupported primitive type for decimal: " + type.getPrimitiveTypeName()); + }; + default: + } + } + + return switch (type.getPrimitiveTypeName()) { + case FIXED_LEN_BYTE_ARRAY, BINARY -> binary -> ByteBuffer.wrap(((Binary) binary).getBytes()); + case INT96 -> binary -> extractTimestampInt96(ByteBuffer.wrap(((Binary) binary).getBytes()).order(ByteOrder.LITTLE_ENDIAN)); + default -> obj -> obj; + }; + } +} diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/BaseIcebergConnectorSmokeTest.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/BaseIcebergConnectorSmokeTest.java index e2d3fec64e1b..1a967a5959f4 100644 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/BaseIcebergConnectorSmokeTest.java +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/BaseIcebergConnectorSmokeTest.java @@ -14,6 +14,7 @@ package io.trino.plugin.iceberg; import com.google.common.collect.ImmutableList; +import com.google.common.collect.Iterables; import com.google.common.collect.Streams; import io.trino.Session; import io.trino.filesystem.FileIterator; @@ -30,6 +31,7 @@ import org.testng.annotations.BeforeClass; import org.testng.annotations.Test; +import java.time.ZonedDateTime; import java.util.List; import java.util.Optional; import java.util.concurrent.CyclicBarrier; @@ -47,6 +49,7 @@ import static io.trino.testing.TestingConnectorSession.SESSION; import static io.trino.testing.TestingNames.randomNameSuffix; import static java.lang.String.format; +import static java.time.format.DateTimeFormatter.ISO_INSTANT; import static java.util.Objects.requireNonNull; import static java.util.concurrent.Executors.newFixedThreadPool; import static java.util.concurrent.TimeUnit.SECONDS; @@ -130,7 +133,9 @@ public void testDeleteRowsConcurrently() ExecutorService executor = newFixedThreadPool(threads); List rows = ImmutableList.of("(1, 0, 0, 0)", "(0, 1, 0, 0)", "(0, 0, 1, 0)", "(0, 0, 0, 1)"); - String[] expectedErrors = new String[]{"Failed to commit Iceberg update to table:", "Failed to replace table due to concurrent updates:"}; + String[] expectedErrors = new String[] {"Failed to commit the transaction during write:", + "Failed to replace table due to concurrent updates:", + "Failed to commit during write:"}; try (TestTable table = new TestTable( getQueryRunner()::execute, "test_concurrent_delete", @@ -154,7 +159,7 @@ public void testDeleteRowsConcurrently() .collect(toImmutableList()); Stream> expectedRows = Streams.mapWithIndex(futures.stream(), (future, index) -> { - boolean deleteSuccessful = tryGetFutureValue(future, 10, SECONDS).orElseThrow(); + boolean deleteSuccessful = tryGetFutureValue(future, 20, SECONDS).orElseThrow(); return deleteSuccessful ? Optional.empty() : Optional.of(rows.get((int) index)); }); String expectedValues = expectedRows.filter(Optional::isPresent).map(Optional::get).collect(joining(", ")); @@ -640,6 +645,76 @@ public void testMetadataTables() protected abstract boolean isFileSorted(Location path, String sortColumnName); + @Test + public void testTableChangesFunction() + { + try (TestTable table = new TestTable( + getQueryRunner()::execute, + "test_table_changes_function_", + "AS SELECT nationkey, name FROM tpch.tiny.nation WITH NO DATA")) { + long initialSnapshot = getMostRecentSnapshotId(table.getName()); + assertUpdate("INSERT INTO " + table.getName() + " SELECT nationkey, name FROM nation", 25); + long snapshotAfterInsert = getMostRecentSnapshotId(table.getName()); + String snapshotAfterInsertTime = getSnapshotTime(table.getName(), snapshotAfterInsert).format(ISO_INSTANT); + + assertQuery( + "SELECT nationkey, name, _change_type, _change_version_id, to_iso8601(_change_timestamp), _change_ordinal " + + "FROM TABLE(system.table_changes('tpch', '%s', %s, %s))".formatted(table.getName(), initialSnapshot, snapshotAfterInsert), + "SELECT nationkey, name, 'insert', %s, '%s', 0 FROM nation".formatted(snapshotAfterInsert, snapshotAfterInsertTime)); + + assertUpdate("DELETE FROM " + table.getName(), 25); + long snapshotAfterDelete = getMostRecentSnapshotId(table.getName()); + String snapshotAfterDeleteTime = getSnapshotTime(table.getName(), snapshotAfterDelete).format(ISO_INSTANT); + + assertQuery( + "SELECT nationkey, name, _change_type, _change_version_id, to_iso8601(_change_timestamp), _change_ordinal " + + "FROM TABLE(system.table_changes('tpch', '%s', %s, %s))".formatted(table.getName(), snapshotAfterInsert, snapshotAfterDelete), + "SELECT nationkey, name, 'delete', %s, '%s', 0 FROM nation".formatted(snapshotAfterDelete, snapshotAfterDeleteTime)); + + assertQuery( + "SELECT nationkey, name, _change_type, _change_version_id, to_iso8601(_change_timestamp), _change_ordinal " + + "FROM TABLE(system.table_changes('tpch', '%s', %s, %s))".formatted(table.getName(), initialSnapshot, snapshotAfterDelete), + "SELECT nationkey, name, 'insert', %s, '%s', 0 FROM nation UNION SELECT nationkey, name, 'delete', %s, '%s', 1 FROM nation".formatted( + snapshotAfterInsert, snapshotAfterInsertTime, snapshotAfterDelete, snapshotAfterDeleteTime)); + } + } + + @Test + public void testRowLevelDeletesWithTableChangesFunction() + { + try (TestTable table = new TestTable( + getQueryRunner()::execute, + "test_row_level_deletes_with_table_changes_function_", + "AS SELECT nationkey, regionkey, name FROM tpch.tiny.nation WITH NO DATA")) { + assertUpdate("INSERT INTO " + table.getName() + " SELECT nationkey, regionkey, name FROM nation", 25); + long snapshotAfterInsert = getMostRecentSnapshotId(table.getName()); + + assertUpdate("DELETE FROM " + table.getName() + " WHERE regionkey = 2", 5); + long snapshotAfterDelete = getMostRecentSnapshotId(table.getName()); + + assertQueryFails( + "SELECT * FROM TABLE(system.table_changes(CURRENT_SCHEMA, '%s', %s, %s))".formatted(table.getName(), snapshotAfterInsert, snapshotAfterDelete), + "Table uses features which are not yet supported by the table_changes function"); + } + } + + protected AutoCloseable createAdditionalTables(String schema) + { + return () -> {}; + } + + private long getMostRecentSnapshotId(String tableName) + { + return (long) Iterables.getOnlyElement(getQueryRunner().execute(format("SELECT snapshot_id FROM \"%s$snapshots\" ORDER BY committed_at DESC LIMIT 1", tableName)) + .getOnlyColumnAsSet()); + } + + private ZonedDateTime getSnapshotTime(String tableName, long snapshotId) + { + return (ZonedDateTime) Iterables.getOnlyElement(getQueryRunner().execute(format("SELECT committed_at FROM \"%s$snapshots\" WHERE snapshot_id = %s", tableName, snapshotId)) + .getOnlyColumnAsSet()); + } + private String getTableLocation(String tableName) { return (String) computeScalar("SELECT DISTINCT regexp_replace(\"$path\", '/[^/]*/[^/]*$', '') FROM " + tableName); diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/BaseIcebergConnectorTest.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/BaseIcebergConnectorTest.java index 9fec165a10cb..6763cb326bab 100644 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/BaseIcebergConnectorTest.java +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/BaseIcebergConnectorTest.java @@ -19,6 +19,7 @@ import com.fasterxml.jackson.databind.node.ObjectNode; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableSet; import io.airlift.units.DataSize; import io.airlift.units.Duration; import io.trino.Session; @@ -109,7 +110,6 @@ import static io.trino.SystemSessionProperties.TASK_WRITER_COUNT; import static io.trino.SystemSessionProperties.USE_PREFERRED_WRITE_PARTITIONING; import static io.trino.plugin.hive.HiveTestUtils.HDFS_ENVIRONMENT; -import static io.trino.plugin.hive.metastore.file.TestingFileHiveMetastore.createTestingFileHiveMetastore; import static io.trino.plugin.iceberg.IcebergFileFormat.AVRO; import static io.trino.plugin.iceberg.IcebergFileFormat.ORC; import static io.trino.plugin.iceberg.IcebergFileFormat.PARQUET; @@ -160,6 +160,7 @@ import static org.testng.Assert.assertNull; import static org.testng.Assert.assertTrue; +@Test(singleThreaded = true) public abstract class BaseIcebergConnectorTest extends BaseConnectorTest { @@ -188,12 +189,24 @@ protected IcebergQueryRunner.Builder createQueryRunnerBuilder() return IcebergQueryRunner.builder() .setIcebergProperties(ImmutableMap.builder() .put("iceberg.file-format", format.name()) + // Only allow some extra properties. Add "sorted_by" so that we can test that the property is disallowed by the connector explicitly. + .put("iceberg.allowed-extra-properties", "extra.property.one,extra.property.two,extra.property.three,sorted_by") // Allows testing the sorting writer flushing to the file system with smaller tables .put("iceberg.writer-sort-buffer-size", "1MB") .buildOrThrow()) .setInitialTables(REQUIRED_TPCH_TABLES); } + protected TestTable newTrinoTable(String namePrefix, @Language("SQL") String tableDefinition) + { + return newTrinoTable(namePrefix, tableDefinition, ImmutableList.of()); + } + + protected TestTable newTrinoTable(String namePrefix, @Language("SQL") String tableDefinition, List rowsToInsert) + { + return new TestTable(getQueryRunner()::execute, namePrefix, tableDefinition, rowsToInsert); + } + @BeforeClass public void initFileSystem() { @@ -226,9 +239,6 @@ protected boolean hasBehavior(TestingConnectorBehavior connectorBehavior) case SUPPORTS_TOPN_PUSHDOWN: return false; - case SUPPORTS_DROP_SCHEMA_CASCADE: - return false; - case SUPPORTS_RENAME_MATERIALIZED_VIEW_ACROSS_SCHEMAS: return false; @@ -259,17 +269,6 @@ public void testAddRowFieldCaseInsensitivity() } } - @Override - public void testAddAndDropColumnName(String columnName) - { - if (columnName.equals("a.dot")) { - assertThatThrownBy(() -> super.testAddAndDropColumnName(columnName)) - .hasMessage("Failed to add column: Cannot add column with ambiguous name: a.dot, use addColumn(parent, name, type)"); - return; - } - super.testAddAndDropColumnName(columnName); - } - @Override protected void verifyVersionedQueryFailurePermissible(Exception e) { @@ -278,13 +277,24 @@ protected void verifyVersionedQueryFailurePermissible(Exception e) "Unsupported type for temporal table version: .*|" + "Unsupported type for table version: .*|" + "No version history table tpch.nation at or before .*|" + - "Iceberg snapshot ID does not exists: .*"); + "Iceberg snapshot ID does not exists: .*|" + + "Cannot find snapshot with reference name: .*"); + } + + @Override + protected Session withoutSmallFileThreshold(Session session) + { + return Session.builder(session) + .setCatalogSessionProperty(getSession().getCatalog().orElseThrow(), "parquet_small_file_threshold", "0B") + .setCatalogSessionProperty(getSession().getCatalog().orElseThrow(), "orc_tiny_stripe_threshold", "0B") + .build(); } @Override protected void verifyConcurrentUpdateFailurePermissible(Exception e) { - assertThat(e).hasMessageContaining("Failed to commit Iceberg update to table"); + assertThat(e).hasMessageMatching("Failed to commit the transaction during write.*|" + + "Failed to commit during write.*"); } @Override @@ -358,7 +368,8 @@ public void testShowCreateTable() "WITH (\n" + " format = '" + format.name() + "',\n" + " format_version = 2,\n" + - " location = '\\E.*/iceberg_data/tpch/orders-.*\\Q'\n" + + " location = '\\E.*/iceberg_data/tpch/orders-.*\\Q',\n" + + " max_commit_retry = 4\n" + ")\\E"); } @@ -1043,20 +1054,6 @@ public void testCreatePartitionedTableWithNestedTypes() dropTable("test_partitioned_table_nested_type"); } - @Test - public void testCreatePartitionedTableWithNestedField() - { - assertQueryFails( - "CREATE TABLE test_partitioned_table_nested_field(parent ROW(child VARCHAR)) WITH (partitioning = ARRAY['\"parent.child\"'])", - "\\QPartitioning by nested field is unsupported: parent.child"); - assertQueryFails( - "CREATE TABLE test_partitioned_table_nested_field(grandparent ROW(parent ROW(child VARCHAR))) WITH (partitioning = ARRAY['\"grandparent.parent.child\"'])", - "\\QPartitioning by nested field is unsupported: grandparent.parent.child"); - assertQueryFails( - "CREATE TABLE test_partitioned_table_nested_field(grandparent ROW(parent ROW(child VARCHAR))) WITH (partitioning = ARRAY['\"grandparent.parent\"'])", - "\\QUnable to parse partitioning value: Cannot partition by non-primitive source field: struct<3: child: optional string>"); - } - @Test public void testCreatePartitionedTableAs() { @@ -1086,6 +1083,7 @@ public void testCreatePartitionedTableAs() " format = '%s',\n" + " format_version = 2,\n" + " location = '%s',\n" + + " max_commit_retry = 4,\n" + " partitioning = ARRAY['order_status','ship_priority','bucket(\"order key\", 9)']\n" + ")", getSession().getCatalog().orElseThrow(), @@ -1379,18 +1377,16 @@ public void testUpdateWithSortOrder() try (TestTable table = new TestTable( getQueryRunner()::execute, "test_sorted_update", - "WITH (sorted_by = ARRAY['comment']) AS TABLE tpch.tiny.lineitem WITH NO DATA")) { + "WITH (sorted_by = ARRAY['comment']) AS TABLE tpch.tiny.customer WITH NO DATA")) { assertUpdate( withSmallRowGroups, - "INSERT INTO " + table.getName() + " TABLE tpch.tiny.lineitem", - "VALUES 60175"); - assertUpdate(withSmallRowGroups, "UPDATE " + table.getName() + " SET comment = substring(comment, 2)", 60175); + "INSERT INTO " + table.getName() + " TABLE tpch.tiny.customer", + "VALUES 1500"); + assertUpdate(withSmallRowGroups, "UPDATE " + table.getName() + " SET comment = substring(comment, 2)", 1500); assertQuery( - "SELECT orderkey, partkey, suppkey, linenumber, quantity, extendedprice, discount, tax, returnflag, linestatus, shipdate, " + - "commitdate, receiptdate, shipinstruct, shipmode, comment FROM " + table.getName(), - "SELECT orderkey, partkey, suppkey, linenumber, quantity, extendedprice, discount, tax, returnflag, linestatus, shipdate, " + - "commitdate, receiptdate, shipinstruct, shipmode, substring(comment, 2) FROM lineitem"); - for (Object filePath : computeActual("SELECT file_path from \"" + table.getName() + "$files\"").getOnlyColumnAsSet()) { + "SELECT custkey, name, address, nationkey, phone, acctbal, mktsegment, comment FROM " + table.getName(), + "SELECT custkey, name, address, nationkey, phone, acctbal, mktsegment, substring(comment, 2) FROM customer"); + for (Object filePath : computeActual("SELECT file_path from \"" + table.getName() + "$files\" WHERE content != 1").getOnlyColumnAsSet()) { assertTrue(isFileSorted((String) filePath, "comment")); } } @@ -1440,7 +1436,8 @@ public void testTableComments() "WITH (\n" + format(" format = '%s',\n", format) + " format_version = 2,\n" + - format(" location = '%s'\n", tempDirPath) + + format(" location = '%s',\n", tempDirPath) + + " max_commit_retry = 4\n" + ")"; String createTableWithoutComment = "" + "CREATE TABLE iceberg.tpch.test_table_comments (\n" + @@ -1449,7 +1446,8 @@ public void testTableComments() "WITH (\n" + " format = '" + format + "',\n" + " format_version = 2,\n" + - " location = '" + tempDirPath + "'\n" + + " location = '" + tempDirPath + "',\n" + + " max_commit_retry = 4\n" + ")"; String createTableSql = format(createTableTemplate, "test table comment", format); assertUpdate(createTableSql); @@ -1539,7 +1537,14 @@ public void testDropRowFieldWhenDuplicates() { // Override because Iceberg doesn't allow duplicated field names in a row type assertThatThrownBy(super::testDropRowFieldWhenDuplicates) - .hasMessage("Invalid schema: multiple fields for name col.a: 2 and 3"); + .hasMessage("Field name 'a' specified more than once"); + } + + @Override // Override because ambiguous field name is disallowed in the connector + public void testDropAmbiguousRowFieldCaseSensitivity() + { + assertThatThrownBy(super::testDropAmbiguousRowFieldCaseSensitivity) + .hasMessage("Field name 'some_field' specified more than once"); } @Test @@ -1694,6 +1699,7 @@ private void testCreateTableLikeForFormat(IcebergFileFormat otherFormat) format = '%s', format_version = 2, location = '%s', + max_commit_retry = 4, partitioning = ARRAY['adate'] )""", format, @@ -1711,7 +1717,8 @@ private void testCreateTableLikeForFormat(IcebergFileFormat otherFormat) WITH ( format = '%s', format_version = 2, - location = '%s' + location = '%s', + max_commit_retry = 4 )""", format, getTableLocation("test_create_table_like_copy1"))); @@ -1724,7 +1731,8 @@ private void testCreateTableLikeForFormat(IcebergFileFormat otherFormat) WITH ( format = '%s', format_version = 2, - location = '%s' + location = '%s', + max_commit_retry = 4 )""", format, getTableLocation("test_create_table_like_copy2"))); @@ -4710,23 +4718,6 @@ public void testSplitPruningForFilterOnNonPartitionColumn(DataMappingTestSetup t } } - @Test - public void testGetIcebergTableProperties() - { - assertUpdate("CREATE TABLE test_iceberg_get_table_props (x BIGINT)"); - verifyIcebergTableProperties(computeActual("SELECT * FROM \"test_iceberg_get_table_props$properties\"")); - dropTable("test_iceberg_get_table_props"); - } - - protected void verifyIcebergTableProperties(MaterializedResult actual) - { - assertThat(actual).isNotNull(); - MaterializedResult expected = resultBuilder(getSession()) - .row("write.format.default", format.name()) - .build(); - assertEqualsIgnoreOrder(actual.getMaterializedRows(), expected.getMaterializedRows()); - } - protected abstract boolean supportsIcebergFileStatistics(String typeName); @Test(dataProvider = "testDataMappingSmokeTestDataProvider") @@ -4871,7 +4862,7 @@ public void testAmbiguousColumnsWithDots() assertUpdate("CREATE TABLE ambiguous (a ROW(cow BIGINT))"); assertThatThrownBy(() -> assertUpdate("ALTER TABLE ambiguous ADD COLUMN \"a.cow\" BIGINT")) - .hasMessage("Failed to add column: Cannot add column with ambiguous name: a.cow, use addColumn(parent, name, type)"); + .hasMessage("Failed to add column: Cannot add column, name already exists: a.cow"); assertUpdate("DROP TABLE ambiguous"); } @@ -5029,7 +5020,7 @@ public void testOptimizeForPartitionedTable(int formatVersion) List initialFiles = getActiveFiles(tableName); assertThat(initialFiles).hasSize(10); - // For optimize we need to set task_writer_count to 1, otherwise it will create more than one file. + // For optimize we need to set task_min_writer_count to 1, otherwise it will create more than one file. computeActual(withSingleWriterPerTask(session), "ALTER TABLE " + tableName + " EXECUTE OPTIMIZE"); assertThat(query(session, "SELECT sum(value), listagg(key, ' ') WITHIN GROUP (ORDER BY key) FROM " + tableName)) @@ -5038,7 +5029,7 @@ public void testOptimizeForPartitionedTable(int formatVersion) List updatedFiles = getActiveFiles(tableName); // as we force repartitioning there should be only 3 partitions assertThat(updatedFiles).hasSize(3); - assertThat(getAllDataFilesFromTableDirectory(tableName)).containsExactlyInAnyOrderElementsOf(concat(initialFiles, updatedFiles)); + assertThat(getAllDataFilesFromTableDirectory(tableName)).containsExactlyInAnyOrderElementsOf(ImmutableSet.copyOf(concat(initialFiles, updatedFiles))); assertUpdate("DROP TABLE " + tableName); } @@ -5185,24 +5176,26 @@ public void testOptimizeCleansUpDeleteFiles() List allDataFilesAfterDelete = getAllDataFilesFromTableDirectory(tableName); assertThat(allDataFilesAfterDelete).hasSize(6); - // For optimize we need to set task_writer_count to 1, otherwise it will create more than one file. - computeActual(withSingleWriterPerTask(getSession()), "ALTER TABLE " + tableName + " EXECUTE OPTIMIZE WHERE regionkey = 4"); + // For optimize we need to set task_min_writer_count to 1, otherwise it will create more than one file. + computeActual(withSingleWriterPerTask(getSession()), "ALTER TABLE " + tableName + " EXECUTE OPTIMIZE WHERE regionkey = 3"); computeActual(sessionWithShortRetentionUnlocked, "ALTER TABLE " + tableName + " EXECUTE EXPIRE_SNAPSHOTS (retention_threshold => '0s')"); computeActual(sessionWithShortRetentionUnlocked, "ALTER TABLE " + tableName + " EXECUTE REMOVE_ORPHAN_FILES (retention_threshold => '0s')"); assertQuery( "SELECT summary['total-delete-files'] FROM \"" + tableName + "$snapshots\" WHERE snapshot_id = " + getCurrentSnapshotId(tableName), - "VALUES '1'"); + "VALUES '0'"); List allDataFilesAfterOptimizeWithWhere = getAllDataFilesFromTableDirectory(tableName); assertThat(allDataFilesAfterOptimizeWithWhere) - .hasSize(6) - .doesNotContain(allDataFilesInitially.stream().filter(file -> file.contains("regionkey=4")) + .hasSize(5) + .doesNotContain(allDataFilesInitially.stream().filter(file -> file.contains("regionkey=3")) + .toArray(String[]::new)) + .contains(allDataFilesInitially.stream().filter(file -> !file.contains("regionkey=3")) .toArray(String[]::new)); assertThat(query("SELECT * FROM " + tableName)) .matches("SELECT * FROM nation WHERE nationkey != 7"); - // For optimize we need to set task_writer_count to 1, otherwise it will create more than one file. + // For optimize we need to set task_min_writer_count to 1, otherwise it will create more than one file. computeActual(withSingleWriterPerTask(getSession()), "ALTER TABLE " + tableName + " EXECUTE OPTIMIZE"); computeActual(sessionWithShortRetentionUnlocked, "ALTER TABLE " + tableName + " EXECUTE EXPIRE_SNAPSHOTS (retention_threshold => '0s')"); computeActual(sessionWithShortRetentionUnlocked, "ALTER TABLE " + tableName + " EXECUTE REMOVE_ORPHAN_FILES (retention_threshold => '0s')"); @@ -5213,7 +5206,8 @@ public void testOptimizeCleansUpDeleteFiles() List allDataFilesAfterFullOptimize = getAllDataFilesFromTableDirectory(tableName); assertThat(allDataFilesAfterFullOptimize) .hasSize(5) - .doesNotContain(allDataFilesInitially.toArray(new String[0])); + // All files skipped from OPTIMIZE as they have no deletes and there's only one file per partition + .contains(allDataFilesAfterOptimizeWithWhere.toArray(new String[0])); assertThat(query("SELECT * FROM " + tableName)) .matches("SELECT * FROM nation WHERE nationkey != 7"); @@ -5230,7 +5224,7 @@ public void testOptimizeSnapshot() long snapshotId = getCurrentSnapshotId(tableName); assertUpdate("INSERT INTO " + tableName + " VALUES 22", 1); assertThatThrownBy(() -> query("ALTER TABLE \"%s@%d\" EXECUTE OPTIMIZE".formatted(tableName, snapshotId))) - .hasMessage(format("Invalid Iceberg table name: %s@%d", tableName, snapshotId)); + .hasMessage(format("line 1:7: Table 'iceberg.tpch.%s@%s' does not exist", tableName, snapshotId)); assertThat(query("SELECT * FROM " + tableName)) .matches("VALUES 11, 22"); @@ -5605,7 +5599,7 @@ public void testExpireSnapshotsOnSnapshot() long snapshotId = getCurrentSnapshotId(tableName); assertUpdate("INSERT INTO " + tableName + " VALUES 22", 1); assertThatThrownBy(() -> query("ALTER TABLE \"%s@%d\" EXECUTE EXPIRE_SNAPSHOTS".formatted(tableName, snapshotId))) - .hasMessage(format("Invalid Iceberg table name: %s@%d", tableName, snapshotId)); + .hasMessage(format("line 1:7: Table 'iceberg.tpch.%s@%s' does not exist", tableName, snapshotId)); assertThat(query("SELECT * FROM " + tableName)) .matches("VALUES 11, 22"); @@ -5647,7 +5641,7 @@ public void testExpireSnapshotsParameterValidation() "\\QUnable to set catalog 'iceberg' table procedure 'EXPIRE_SNAPSHOTS' property 'retention_threshold' to ['33mb']: Unknown time unit: mb"); assertQueryFails( "ALTER TABLE nation EXECUTE EXPIRE_SNAPSHOTS (retention_threshold => '33s')", - "\\QRetention specified (33.00s) is shorter than the minimum retention configured in the system (7.00d). Minimum retention can be changed with iceberg.expire_snapshots.min-retention configuration property or iceberg.expire_snapshots_min_retention session property"); + "\\QRetention specified (33.00s) is shorter than the minimum retention configured in the system (7.00d). Minimum retention can be changed with iceberg.expire-snapshots.min-retention configuration property or iceberg.expire_snapshots_min_retention session property"); } @Test @@ -5750,7 +5744,7 @@ private void testCleaningUpWithTableWithSpecifiedLocation(String suffix) List prunedMetadataFiles = getAllMetadataFilesFromTableDirectory(tableDirectory); List prunedSnapshots = getSnapshotIds(tableName); assertThat(prunedMetadataFiles).as("prunedMetadataFiles") - .hasSize(initialMetadataFiles.size() - 3); + .hasSize(initialMetadataFiles.size() - 2); assertThat(prunedSnapshots).as("prunedSnapshots") .hasSizeLessThan(initialSnapshots.size()) .hasSize(1); @@ -5785,7 +5779,7 @@ public void testRemoveOrphanFilesParameterValidation() "\\QUnable to set catalog 'iceberg' table procedure 'REMOVE_ORPHAN_FILES' property 'retention_threshold' to ['33mb']: Unknown time unit: mb"); assertQueryFails( "ALTER TABLE nation EXECUTE REMOVE_ORPHAN_FILES (retention_threshold => '33s')", - "\\QRetention specified (33.00s) is shorter than the minimum retention configured in the system (7.00d). Minimum retention can be changed with iceberg.remove_orphan_files.min-retention configuration property or iceberg.remove_orphan_files_min_retention session property"); + "\\QRetention specified (33.00s) is shorter than the minimum retention configured in the system (7.00d). Minimum retention can be changed with iceberg.remove-orphan-files.min-retention configuration property or iceberg.remove_orphan_files_min_retention session property"); } @Test @@ -5797,7 +5791,7 @@ public void testRemoveOrphanFilesOnSnapshot() long snapshotId = getCurrentSnapshotId(tableName); assertUpdate("INSERT INTO " + tableName + " VALUES 22", 1); assertThatThrownBy(() -> query("ALTER TABLE \"%s@%d\" EXECUTE REMOVE_ORPHAN_FILES".formatted(tableName, snapshotId))) - .hasMessage(format("Invalid Iceberg table name: %s@%d", tableName, snapshotId)); + .hasMessage(format("line 1:7: Table 'iceberg.tpch.%s@%s' does not exist", tableName, snapshotId)); assertThat(query("SELECT * FROM " + tableName)) .matches("VALUES 11, 22"); @@ -5940,19 +5934,19 @@ public void testModifyingOldSnapshotIsNotPossible() assertUpdate(format("INSERT INTO %s VALUES 4,5,6", tableName), 3); assertQuery(format("SELECT * FROM %s FOR VERSION AS OF %d", tableName, oldSnapshotId), "VALUES 1,2,3"); assertThatThrownBy(() -> query(format("INSERT INTO \"%s@%d\" VALUES 7,8,9", tableName, oldSnapshotId))) - .hasMessage(format("Invalid Iceberg table name: %s@%d", tableName, oldSnapshotId)); + .hasMessage(format("Table 'iceberg.tpch.%s@%s' does not exist", tableName, oldSnapshotId)); assertThatThrownBy(() -> query(format("DELETE FROM \"%s@%d\" WHERE col = 5", tableName, oldSnapshotId))) - .hasMessage(format("Invalid Iceberg table name: %s@%d", tableName, oldSnapshotId)); + .hasMessage(format("line 1:1: Table 'iceberg.tpch.%s@%s' does not exist", tableName, oldSnapshotId)); assertThatThrownBy(() -> query(format("UPDATE \"%s@%d\" SET col = 50 WHERE col = 5", tableName, oldSnapshotId))) - .hasMessage(format("Invalid Iceberg table name: %s@%d", tableName, oldSnapshotId)); + .hasMessage(format("line 1:1: Table 'iceberg.tpch.%s@%s' does not exist", tableName, oldSnapshotId)); assertThatThrownBy(() -> query(format("INSERT INTO \"%s@%d\" VALUES 7,8,9", tableName, getCurrentSnapshotId(tableName)))) - .hasMessage(format("Invalid Iceberg table name: %s@%d", tableName, getCurrentSnapshotId(tableName))); + .hasMessage(format("Table 'iceberg.tpch.%s@%s' does not exist", tableName, getCurrentSnapshotId(tableName))); assertThatThrownBy(() -> query(format("DELETE FROM \"%s@%d\" WHERE col = 9", tableName, getCurrentSnapshotId(tableName)))) - .hasMessage(format("Invalid Iceberg table name: %s@%d", tableName, getCurrentSnapshotId(tableName))); + .hasMessage(format("line 1:1: Table 'iceberg.tpch.%s@%s' does not exist", tableName, getCurrentSnapshotId(tableName))); assertThatThrownBy(() -> assertUpdate(format("UPDATE \"%s@%d\" set col = 50 WHERE col = 5", tableName, getCurrentSnapshotId(tableName)))) - .hasMessage(format("Invalid Iceberg table name: %s@%d", tableName, getCurrentSnapshotId(tableName))); + .hasMessage(format("line 1:1: Table 'iceberg.tpch.%s@%s' does not exist", tableName, getCurrentSnapshotId(tableName))); assertThatThrownBy(() -> query(format("ALTER TABLE \"%s@%d\" EXECUTE OPTIMIZE", tableName, oldSnapshotId))) - .hasMessage(format("Invalid Iceberg table name: %s@%d", tableName, oldSnapshotId)); + .hasMessage(format("line 1:1: Table 'iceberg.tpch.%s@%s' does not exist", tableName, oldSnapshotId)); assertQuery(format("SELECT * FROM %s", tableName), "VALUES 1,2,3,4,5,6"); assertUpdate("DROP TABLE " + tableName); @@ -6587,30 +6581,6 @@ public void testSnapshotSummariesHaveTrinoQueryIdFormatV2() "WHEN MATCHED THEN UPDATE SET b = t.b * 50", tableName, sourceTableName))); } - @Test - public void testMaterializedViewSnapshotSummariesHaveTrinoQueryId() - { - String matViewName = "test_materialized_view_snapshot_query_ids" + randomNameSuffix(); - String sourceTableName = "test_source_table_for_mat_view" + randomNameSuffix(); - assertUpdate(format("CREATE TABLE %s (a bigint, b bigint)", sourceTableName)); - assertUpdate(format("INSERT INTO %s VALUES (1, 1), (1, 4), (2, 2)", sourceTableName), 3); - - // create a materialized view - QueryId matViewCreateQueryId = getDistributedQueryRunner() - .executeWithQueryId(getSession(), format("CREATE MATERIALIZED VIEW %s WITH (partitioning = ARRAY['a']) AS SELECT * FROM %s", matViewName, sourceTableName)) - .getQueryId(); - - // fetch the underlying storage table name so we can inspect its snapshot summary after the REFRESH - // running queries against "materialized_view$snapshots" is not supported - String storageTable = (String) getDistributedQueryRunner() - .execute(getSession(), format("SELECT storage_table FROM system.metadata.materialized_views WHERE name = '%s'", matViewName)) - .getOnlyValue(); - - assertQueryIdStored(storageTable, matViewCreateQueryId); - - assertQueryIdStored(storageTable, executeWithQueryId(format("REFRESH MATERIALIZED VIEW %s", matViewName))); - } - @Override protected OptionalInt maxTableNameLength() { @@ -6660,12 +6630,8 @@ public void testAlterTableWithUnsupportedProperties() assertUpdate("CREATE TABLE " + tableName + " (a bigint)"); - assertQueryFails("ALTER TABLE " + tableName + " SET PROPERTIES orc_bloom_filter_columns = ARRAY['a']", - "The following properties cannot be updated: orc_bloom_filter_columns"); assertQueryFails("ALTER TABLE " + tableName + " SET PROPERTIES location = '/var/data/table/', orc_bloom_filter_fpp = 0.5", "The following properties cannot be updated: location, orc_bloom_filter_fpp"); - assertQueryFails("ALTER TABLE " + tableName + " SET PROPERTIES format = 'ORC', orc_bloom_filter_columns = ARRAY['a']", - "The following properties cannot be updated: orc_bloom_filter_columns"); assertUpdate("DROP TABLE " + tableName); } @@ -6863,7 +6829,7 @@ public void testDropCorruptedTableWithHiveRedirection() "iceberg.catalog.type", "TESTING_FILE_METASTORE", "hive.metastore.catalog.dir", dataDirectory.getPath())); - queryRunner.installPlugin(new TestingHivePlugin(createTestingFileHiveMetastore(dataDirectory))); + queryRunner.installPlugin(new TestingHivePlugin(dataDirectory.toPath())); queryRunner.createCatalog( hiveRedirectionCatalog, "hive", @@ -6918,6 +6884,56 @@ public void testNoRetryWhenMetadataFileInvalid() assertUpdate("DROP TABLE " + tableName); } + @Test + public void testTableChangesFunctionAfterSchemaChange() + { + try (TestTable table = new TestTable( + getQueryRunner()::execute, + "test_table_changes_function_", + "AS SELECT nationkey, name FROM tpch.tiny.nation WITH NO DATA")) { + long initialSnapshot = getCurrentSnapshotId(table.getName()); + assertUpdate("INSERT INTO " + table.getName() + " SELECT nationkey, name FROM nation WHERE nationkey < 5", 5); + long snapshotAfterInsert = getCurrentSnapshotId(table.getName()); + + assertUpdate("ALTER TABLE " + table.getName() + " DROP COLUMN name"); + assertUpdate("INSERT INTO " + table.getName() + " SELECT nationkey FROM nation WHERE nationkey >= 5 AND nationkey < 10", 5); + long snapshotAfterDropColumn = getCurrentSnapshotId(table.getName()); + + assertUpdate("ALTER TABLE " + table.getName() + " ADD COLUMN comment VARCHAR"); + assertUpdate("INSERT INTO " + table.getName() + " SELECT nationkey, comment FROM nation WHERE nationkey >= 10 AND nationkey < 15", 5); + long snapshotAfterAddColumn = getCurrentSnapshotId(table.getName()); + + assertUpdate("ALTER TABLE " + table.getName() + " ADD COLUMN name VARCHAR"); + assertUpdate("INSERT INTO " + table.getName() + " SELECT nationkey, comment, name FROM nation WHERE nationkey >= 15", 10); + long snapshotAfterReaddingNameColumn = getCurrentSnapshotId(table.getName()); + + assertQuery( + "SELECT nationkey, name, _change_type, _change_version_id, _change_ordinal " + + "FROM TABLE(system.table_changes('tpch', '%s', %s, %s))".formatted(table.getName(), initialSnapshot, snapshotAfterInsert), + "SELECT nationkey, name, 'insert', %s, 0 FROM nation WHERE nationkey < 5".formatted(snapshotAfterInsert)); + + assertQuery( + "SELECT nationkey, _change_type, _change_version_id, _change_ordinal " + + "FROM TABLE(system.table_changes('tpch', '%s', %s, %s))".formatted(table.getName(), initialSnapshot, snapshotAfterDropColumn), + "SELECT nationkey, 'insert', %s, 0 FROM nation WHERE nationkey < 5 UNION SELECT nationkey, 'insert', %s, 1 FROM nation WHERE nationkey >= 5 AND nationkey < 10 ".formatted(snapshotAfterInsert, snapshotAfterDropColumn)); + + assertQuery( + "SELECT nationkey, comment, _change_type, _change_version_id, _change_ordinal " + + "FROM TABLE(system.table_changes('tpch', '%s', %s, %s))".formatted(table.getName(), initialSnapshot, snapshotAfterAddColumn), + ("SELECT nationkey, NULL, 'insert', %s, 0 FROM nation WHERE nationkey < 5 " + + "UNION SELECT nationkey, NULL, 'insert', %s, 1 FROM nation WHERE nationkey >= 5 AND nationkey < 10 " + + "UNION SELECT nationkey, comment, 'insert', %s, 2 FROM nation WHERE nationkey >= 10 AND nationkey < 15").formatted(snapshotAfterInsert, snapshotAfterDropColumn, snapshotAfterAddColumn)); + + assertQuery( + "SELECT nationkey, comment, name, _change_type, _change_version_id, _change_ordinal " + + "FROM TABLE(system.table_changes('tpch', '%s', %s, %s))".formatted(table.getName(), initialSnapshot, snapshotAfterReaddingNameColumn), + ("SELECT nationkey, NULL, NULL, 'insert', %s, 0 FROM nation WHERE nationkey < 5 " + + "UNION SELECT nationkey, NULL, NULL, 'insert', %s, 1 FROM nation WHERE nationkey >= 5 AND nationkey < 10 " + + "UNION SELECT nationkey, comment, NULL, 'insert', %s, 2 FROM nation WHERE nationkey >= 10 AND nationkey < 15" + + "UNION SELECT nationkey, comment, name, 'insert', %s, 3 FROM nation WHERE nationkey >= 15").formatted(snapshotAfterInsert, snapshotAfterDropColumn, snapshotAfterAddColumn, snapshotAfterReaddingNameColumn)); + } + } + @Override protected void verifyTableNameLengthFailurePermissible(Throwable e) { diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/BaseIcebergMaterializedViewTest.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/BaseIcebergMaterializedViewTest.java index 73eaa253a4ac..771bb2bf15ae 100644 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/BaseIcebergMaterializedViewTest.java +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/BaseIcebergMaterializedViewTest.java @@ -15,32 +15,32 @@ import com.google.common.collect.ImmutableSet; import io.trino.Session; -import io.trino.metadata.MaterializedViewDefinition; -import io.trino.metadata.QualifiedObjectName; -import io.trino.spi.connector.SchemaTableName; +import io.trino.filesystem.Location; +import io.trino.filesystem.local.LocalFileSystem; +import io.trino.plugin.iceberg.fileio.ForwardingFileIo; +import io.trino.spi.QueryId; import io.trino.sql.tree.ExplainType; import io.trino.testing.AbstractTestQueryFramework; import io.trino.testing.MaterializedRow; -import io.trino.transaction.TransactionId; -import io.trino.transaction.TransactionManager; +import org.apache.iceberg.PartitionField; +import org.apache.iceberg.TableMetadata; +import org.apache.iceberg.TableMetadataParser; import org.assertj.core.api.Condition; import org.testng.annotations.BeforeClass; import org.testng.annotations.DataProvider; import org.testng.annotations.Test; +import java.nio.file.Path; import java.util.Optional; import java.util.Set; import static com.google.common.collect.ImmutableSet.toImmutableSet; +import static com.google.common.collect.Iterables.getOnlyElement; import static io.trino.SystemSessionProperties.LEGACY_MATERIALIZED_VIEW_GRACE_PERIOD; import static io.trino.testing.MaterializedResult.DEFAULT_PRECISION; -import static io.trino.testing.TestingAccessControlManager.TestingPrivilegeType.DELETE_TABLE; import static io.trino.testing.TestingAccessControlManager.TestingPrivilegeType.DROP_MATERIALIZED_VIEW; -import static io.trino.testing.TestingAccessControlManager.TestingPrivilegeType.INSERT_TABLE; import static io.trino.testing.TestingAccessControlManager.TestingPrivilegeType.REFRESH_MATERIALIZED_VIEW; import static io.trino.testing.TestingAccessControlManager.TestingPrivilegeType.RENAME_MATERIALIZED_VIEW; -import static io.trino.testing.TestingAccessControlManager.TestingPrivilegeType.SELECT_COLUMN; -import static io.trino.testing.TestingAccessControlManager.TestingPrivilegeType.UPDATE_TABLE; import static io.trino.testing.TestingAccessControlManager.privilege; import static io.trino.testing.TestingNames.randomNameSuffix; import static java.lang.String.format; @@ -52,10 +52,10 @@ public abstract class BaseIcebergMaterializedViewTest extends AbstractTestQueryFramework { - protected final String storageSchemaName = "testing_storage_schema_" + randomNameSuffix(); - protected abstract String getSchemaDirectory(); + protected abstract String getStorageMetadataLocation(String materializedViewName); + @BeforeClass public void setUp() { @@ -65,17 +65,14 @@ public void setUp() assertUpdate("CREATE TABLE base_table2 (_varchar VARCHAR, _bigint BIGINT, _date DATE) WITH (partitioning = ARRAY['_bigint', '_date'])"); assertUpdate("INSERT INTO base_table2 VALUES ('a', 0, DATE '2019-09-08'), ('a', 1, DATE '2019-09-08'), ('a', 0, DATE '2019-09-09')", 3); - - assertUpdate("CREATE SCHEMA " + storageSchemaName); } @Test public void testShowTables() { assertUpdate("CREATE MATERIALIZED VIEW materialized_view_show_tables_test AS SELECT * FROM base_table1"); - SchemaTableName storageTableName = getStorageTable("materialized_view_show_tables_test"); - Set expectedTables = ImmutableSet.of("base_table1", "base_table2", "materialized_view_show_tables_test", storageTableName.getTableName()); + Set expectedTables = ImmutableSet.of("base_table1", "base_table2", "materialized_view_show_tables_test"); Set actualTables = computeActual("SHOW TABLES").getOnlyColumnAsSet().stream() .map(String.class::cast) .collect(toImmutableSet()); @@ -107,20 +104,6 @@ public void testMaterializedViewsMetadata() computeActual("CREATE TABLE small_region AS SELECT * FROM tpch.tiny.region LIMIT 1"); computeActual(format("CREATE MATERIALIZED VIEW %s AS SELECT * FROM small_region LIMIT 1", materializedViewName)); - // test storage table name - assertQuery( - format( - "SELECT storage_catalog, storage_schema, CONCAT(storage_schema, '.', storage_table)" + - "FROM system.metadata.materialized_views WHERE schema_name = '%s' AND name = '%s'", - // TODO (https://github.com/trinodb/trino/issues/9039) remove redundant schema_name filter - schemaName, - materializedViewName), - format( - "VALUES ('%s', '%s', '%s')", - catalogName, - schemaName, - getStorageTable(catalogName, schemaName, materializedViewName))); - // test freshness update assertQuery( // TODO (https://github.com/trinodb/trino/issues/9039) remove redundant schema_name filter @@ -182,7 +165,7 @@ public void testShowCreate() "WITH (\n" + " format = 'ORC',\n" + " format_version = 2,\n" + - " location = '" + getSchemaDirectory() + "/st_\\E[0-9a-f]+-[0-9a-f]+\\Q',\n" + + " location = '" + getSchemaDirectory() + "/test_mv_show_create-\\E[0-9a-f]+\\Q',\n" + " orc_bloom_filter_columns = ARRAY['_date'],\n" + " orc_bloom_filter_fpp = 1E-1,\n" + " partitioning = ARRAY['_date'],\n" + @@ -269,22 +252,6 @@ public void testRefreshDenyPermission() assertUpdate("DROP MATERIALIZED VIEW materialized_view_refresh_deny"); } - @Test - public void testRefreshAllowedWithRestrictedStorageTable() - { - assertUpdate("CREATE MATERIALIZED VIEW materialized_view_refresh AS SELECT * FROM base_table1"); - SchemaTableName storageTable = getStorageTable("materialized_view_refresh"); - - assertAccessAllowed( - "REFRESH MATERIALIZED VIEW materialized_view_refresh", - privilege(storageTable.getTableName(), INSERT_TABLE), - privilege(storageTable.getTableName(), DELETE_TABLE), - privilege(storageTable.getTableName(), UPDATE_TABLE), - privilege(storageTable.getTableName(), SELECT_COLUMN)); - - assertUpdate("DROP MATERIALIZED VIEW materialized_view_refresh"); - } - @Test public void testCreateRefreshSelect() { @@ -524,7 +491,7 @@ public void testSqlFeatures() "WITH (\n" + " format = 'PARQUET',\n" + " format_version = 2,\n" + - " location = '" + getSchemaDirectory() + "/st_\\E[0-9a-f]+-[0-9a-f]+\\Q',\n" + + " location = '" + getSchemaDirectory() + "/materialized_view_window-\\E[0-9a-f]+\\Q',\n" + " partitioning = ARRAY['_date'],\n" + " storage_schema = '" + schema + "'\n" + ") AS\n" + @@ -632,47 +599,6 @@ public void testNestedMaterializedViews() assertUpdate("DROP MATERIALIZED VIEW materialized_view_level2"); } - @Test - public void testStorageSchemaProperty() - { - String schemaName = getSession().getSchema().orElseThrow(); - String viewName = "storage_schema_property_test"; - assertUpdate( - "CREATE MATERIALIZED VIEW " + viewName + " " + - "WITH (storage_schema = '" + storageSchemaName + "') AS " + - "SELECT * FROM base_table1"); - SchemaTableName storageTable = getStorageTable(viewName); - assertThat(storageTable.getSchemaName()).isEqualTo(storageSchemaName); - - assertUpdate("REFRESH MATERIALIZED VIEW " + viewName, 6); - assertThat(computeActual("SELECT * FROM " + viewName).getRowCount()).isEqualTo(6); - assertThat(getExplainPlan("SELECT * FROM " + viewName, ExplainType.Type.IO)) - .doesNotContain("base_table1") - .contains(storageSchemaName); - - assertThat((String) computeScalar("SHOW CREATE MATERIALIZED VIEW " + viewName)) - .contains("storage_schema = '" + storageSchemaName + "'"); - - Set storageSchemaTables = computeActual("SHOW TABLES IN " + storageSchemaName).getOnlyColumnAsSet().stream() - .map(String.class::cast) - .collect(toImmutableSet()); - assertThat(storageSchemaTables).contains(storageTable.getTableName()); - - assertUpdate("DROP MATERIALIZED VIEW " + viewName); - storageSchemaTables = computeActual("SHOW TABLES IN " + storageSchemaName).getOnlyColumnAsSet().stream() - .map(String.class::cast) - .collect(toImmutableSet()); - assertThat(storageSchemaTables).doesNotContain(storageTable.getTableName()); - - assertThatThrownBy(() -> query( - "CREATE MATERIALIZED VIEW " + viewName + " " + - "WITH (storage_schema = 'non_existent') AS " + - "SELECT * FROM base_table1")) - .hasMessageContaining("non_existent not found"); - assertThatThrownBy(() -> query("DESCRIBE " + viewName)) - .hasMessageContaining(format("'iceberg.%s.%s' does not exist", schemaName, viewName)); - } - @Test(dataProvider = "testBucketPartitioningDataProvider") public void testBucketPartitioning(String dataType, String exampleValue) { @@ -683,9 +609,11 @@ public void testBucketPartitioning(String dataType, String exampleValue) assertUpdate("CREATE MATERIALIZED VIEW test_bucket_partitioning WITH (partitioning=ARRAY['bucket(col, 4)']) AS SELECT * FROM (VALUES CAST(NULL AS %s), %s) t(col)" .formatted(dataType, exampleValue)); try { - SchemaTableName storageTable = getStorageTable("test_bucket_partitioning"); - assertThat((String) computeScalar("SHOW CREATE TABLE " + storageTable)) - .contains("partitioning = ARRAY['bucket(col, 4)']"); + TableMetadata storageMetadata = getStorageTableMetadata("test_bucket_partitioning"); + assertThat(storageMetadata.spec().fields()).hasSize(1); + PartitionField bucketPartitionField = getOnlyElement(storageMetadata.spec().fields()); + assertThat(bucketPartitionField.name()).isEqualTo("col_bucket"); + assertThat(bucketPartitionField.transform().toString()).isEqualTo("bucket[4]"); assertThat(query("SELECT * FROM test_bucket_partitioning WHERE col = " + exampleValue)) .matches("SELECT " + exampleValue); @@ -724,9 +652,11 @@ public void testTruncatePartitioning(String dataType, String exampleValue) assertUpdate("CREATE MATERIALIZED VIEW test_truncate_partitioning WITH (partitioning=ARRAY['truncate(col, 4)']) AS SELECT * FROM (VALUES CAST(NULL AS %s), %s) t(col)" .formatted(dataType, exampleValue)); try { - SchemaTableName storageTable = getStorageTable("test_truncate_partitioning"); - assertThat((String) computeScalar("SHOW CREATE TABLE " + storageTable)) - .contains("partitioning = ARRAY['truncate(col, 4)']"); + TableMetadata storageMetadata = getStorageTableMetadata("test_truncate_partitioning"); + assertThat(storageMetadata.spec().fields()).hasSize(1); + PartitionField bucketPartitionField = getOnlyElement(storageMetadata.spec().fields()); + assertThat(bucketPartitionField.name()).isEqualTo("col_trunc"); + assertThat(bucketPartitionField.transform().toString()).isEqualTo("truncate[4]"); assertThat(query("SELECT * FROM test_truncate_partitioning WHERE col = " + exampleValue)) .matches("SELECT " + exampleValue); @@ -759,9 +689,11 @@ public void testTemporalPartitioning(String partitioning, String dataType, Strin assertUpdate("CREATE MATERIALIZED VIEW test_temporal_partitioning WITH (partitioning=ARRAY['%s(col)']) AS SELECT * FROM (VALUES CAST(NULL AS %s), %s) t(col)" .formatted(partitioning, dataType, exampleValue)); try { - SchemaTableName storageTable = getStorageTable("test_temporal_partitioning"); - assertThat((String) computeScalar("SHOW CREATE TABLE " + storageTable)) - .contains("partitioning = ARRAY['%s(col)']".formatted(partitioning)); + TableMetadata storageMetadata = getStorageTableMetadata("test_temporal_partitioning"); + assertThat(storageMetadata.spec().fields()).hasSize(1); + PartitionField bucketPartitionField = getOnlyElement(storageMetadata.spec().fields()); + assertThat(bucketPartitionField.name()).isEqualTo("col_" + partitioning); + assertThat(bucketPartitionField.transform().toString()).isEqualTo(partitioning); assertThat(query("SELECT * FROM test_temporal_partitioning WHERE col = " + exampleValue)) .matches("SELECT " + exampleValue); @@ -789,25 +721,40 @@ public Object[][] testTemporalPartitioningDataProvider() }; } - protected String getColumnComment(String tableName, String columnName) + @Test + public void testMaterializedViewSnapshotSummariesHaveTrinoQueryId() { - return (String) computeScalar("SELECT comment FROM information_schema.columns WHERE table_schema = '" + getSession().getSchema().orElseThrow() + "' AND table_name = '" + tableName + "' AND column_name = '" + columnName + "'"); + String materializedViewName = "test_materialized_view_snapshot_query_ids" + randomNameSuffix(); + String sourceTableName = "test_source_table_for_mat_view" + randomNameSuffix(); + assertUpdate(format("CREATE TABLE %s (a bigint, b bigint)", sourceTableName)); + QueryId matViewCreateQueryId = getDistributedQueryRunner() + .executeWithQueryId(getSession(), format("CREATE MATERIALIZED VIEW %s WITH (partitioning = ARRAY['a']) AS SELECT * FROM %s", materializedViewName, sourceTableName)) + .getQueryId(); + + try { + assertUpdate(format("INSERT INTO %s VALUES (1, 1), (1, 4), (2, 2)", sourceTableName), 3); + + QueryId refreshQueryId = getDistributedQueryRunner() + .executeWithQueryId(getSession(), format("REFRESH MATERIALIZED VIEW %s", materializedViewName)) + .getQueryId(); + String savedQueryId = getStorageTableMetadata(materializedViewName).currentSnapshot().summary().get("trino_query_id"); + assertThat(savedQueryId).isEqualTo(refreshQueryId.getId()); + } + finally { + assertUpdate("DROP TABLE " + sourceTableName); + assertUpdate("DROP MATERIALIZED VIEW " + materializedViewName); + } } - private SchemaTableName getStorageTable(String materializedViewName) + protected String getColumnComment(String tableName, String columnName) { - return getStorageTable(getSession().getCatalog().orElseThrow(), getSession().getSchema().orElseThrow(), materializedViewName); + return (String) computeScalar("SELECT comment FROM information_schema.columns WHERE table_schema = '" + getSession().getSchema().orElseThrow() + "' AND table_name = '" + tableName + "' AND column_name = '" + columnName + "'"); } - private SchemaTableName getStorageTable(String catalogName, String schemaName, String materializedViewName) + private TableMetadata getStorageTableMetadata(String materializedViewName) { - TransactionManager transactionManager = getQueryRunner().getTransactionManager(); - TransactionId transactionId = transactionManager.beginTransaction(false); - Session session = getSession().beginTransactionId(transactionId, transactionManager, getQueryRunner().getAccessControl()); - Optional materializedView = getQueryRunner().getMetadata() - .getMaterializedView(session, new QualifiedObjectName(catalogName, schemaName, materializedViewName)); - assertThat(materializedView).isPresent(); - return materializedView.get().getStorageTable().get().getSchemaTableName(); + Location metadataLocation = Location.of(getStorageMetadataLocation(materializedViewName)); + return TableMetadataParser.read(new ForwardingFileIo(new LocalFileSystem(Path.of(metadataLocation.parentDirectory().toString()))), "local:///" + metadataLocation); } private long getLatestSnapshotId(String tableName) diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/BaseIcebergMinioConnectorSmokeTest.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/BaseIcebergMinioConnectorSmokeTest.java index d60cbc15378f..cb2af31f2ed6 100644 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/BaseIcebergMinioConnectorSmokeTest.java +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/BaseIcebergMinioConnectorSmokeTest.java @@ -224,7 +224,7 @@ protected void dropTableFromMetastore(String tableName) HiveMetastore metastore = new BridgingHiveMetastore( testingThriftHiveMetastoreBuilder() .metastoreClient(hiveMinioDataLake.getHiveHadoop().getHiveMetastoreEndpoint()) - .build()); + .build(this::closeAfterClass)); metastore.dropTable(schemaName, tableName, false); assertThat(metastore.getTable(schemaName, tableName)).isEmpty(); } @@ -235,7 +235,7 @@ protected String getMetadataLocation(String tableName) HiveMetastore metastore = new BridgingHiveMetastore( testingThriftHiveMetastoreBuilder() .metastoreClient(hiveMinioDataLake.getHiveHadoop().getHiveMetastoreEndpoint()) - .build()); + .build(this::closeAfterClass)); return metastore .getTable(schemaName, tableName).orElseThrow() .getParameters().get("metadata_location"); diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/IcebergQueryRunner.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/IcebergQueryRunner.java index b75c67b9ed10..09489bcefc08 100644 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/IcebergQueryRunner.java +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/IcebergQueryRunner.java @@ -70,6 +70,11 @@ public static Builder builder() return new Builder(); } + public static Builder builder(String schema) + { + return new Builder(schema); + } + public static class Builder extends DistributedQueryRunner.Builder { @@ -85,6 +90,14 @@ protected Builder() .build()); } + protected Builder(String schema) + { + super(testSessionBuilder() + .setCatalog(ICEBERG_CATALOG) + .setSchema(schema) + .build()); + } + public Builder setMetastoreDirectory(File metastoreDirectory) { this.metastoreDirectory = Optional.of(metastoreDirectory); @@ -100,7 +113,9 @@ public Builder setIcebergProperties(Map icebergProperties) public Builder addIcebergProperty(String key, String value) { - this.icebergProperties.put(key, value); + if (value != null) { + this.icebergProperties.put(key, value); + } return self(); } @@ -163,7 +178,7 @@ public static void main(String[] args) File warehouseLocation = Files.newTemporaryFolder(); warehouseLocation.deleteOnExit(); - Catalog backend = backendCatalog(warehouseLocation); + Catalog backend = backendCatalog(warehouseLocation.toPath()); DelegatingRestSessionCatalog delegatingCatalog = DelegatingRestSessionCatalog.builder() .delegate(backend) diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/IcebergTestUtils.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/IcebergTestUtils.java index 39199092f2c3..9a67baad269c 100644 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/IcebergTestUtils.java +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/IcebergTestUtils.java @@ -27,13 +27,26 @@ import io.trino.orc.metadata.statistics.StringStatistics; import io.trino.orc.metadata.statistics.StripeStatistics; import io.trino.parquet.ParquetReaderOptions; +import io.trino.parquet.metadata.BlockMetadata; +import io.trino.parquet.metadata.ColumnChunkMetadata; +import io.trino.parquet.metadata.ParquetMetadata; import io.trino.parquet.reader.MetadataReader; +import io.trino.plugin.base.CatalogName; import io.trino.plugin.hive.FileFormatDataSourceStats; +import io.trino.plugin.hive.TrinoViewHiveMetastore; +import io.trino.plugin.hive.metastore.HiveMetastore; +import io.trino.plugin.hive.metastore.HiveMetastoreFactory; +import io.trino.plugin.hive.metastore.cache.CachingHiveMetastore; import io.trino.plugin.hive.parquet.TrinoParquetDataSource; +import io.trino.plugin.iceberg.catalog.IcebergTableOperationsProvider; +import io.trino.plugin.iceberg.catalog.TrinoCatalog; +import io.trino.plugin.iceberg.catalog.file.FileMetastoreTableOperationsProvider; +import io.trino.plugin.iceberg.catalog.hms.TrinoHiveCatalog; +import io.trino.spi.connector.SchemaTableName; +import io.trino.spi.type.TestingTypeManager; import io.trino.testing.DistributedQueryRunner; -import org.apache.parquet.hadoop.metadata.BlockMetaData; -import org.apache.parquet.hadoop.metadata.ColumnChunkMetaData; -import org.apache.parquet.hadoop.metadata.ParquetMetadata; +import io.trino.testing.QueryRunner; +import org.apache.iceberg.BaseTable; import java.io.File; import java.io.IOException; @@ -46,7 +59,11 @@ import static com.google.common.collect.ImmutableList.toImmutableList; import static com.google.common.collect.Iterators.getOnlyElement; import static com.google.common.collect.MoreCollectors.onlyElement; +import static com.google.common.util.concurrent.MoreExecutors.directExecutor; +import static io.trino.plugin.hive.metastore.cache.CachingHiveMetastore.createPerTransactionCache; import static io.trino.plugin.iceberg.IcebergQueryRunner.ICEBERG_CATALOG; +import static io.trino.plugin.iceberg.IcebergUtil.loadIcebergTable; +import static io.trino.testing.TestingConnectorSession.SESSION; public final class IcebergTestUtils { @@ -56,9 +73,9 @@ private IcebergTestUtils() public static Session withSmallRowGroups(Session session) { return Session.builder(session) - .setCatalogSessionProperty("iceberg", "orc_writer_max_stripe_rows", "10") + .setCatalogSessionProperty("iceberg", "orc_writer_max_stripe_rows", "20") .setCatalogSessionProperty("iceberg", "parquet_writer_block_size", "1kB") - .setCatalogSessionProperty("iceberg", "parquet_writer_batch_size", "10") + .setCatalogSessionProperty("iceberg", "parquet_writer_batch_size", "20") .build(); } @@ -121,20 +138,19 @@ public static boolean checkParquetFileSorting(String path, String sortColumnName @SuppressWarnings({"unchecked", "rawtypes"}) public static boolean checkParquetFileSorting(TrinoInputFile inputFile, String sortColumnName) { - ParquetMetadata parquetMetadata; + ParquetMetadata parquetMetadata = getParquetFileMetadata(inputFile); + List blocks; try { - parquetMetadata = MetadataReader.readFooter( - new TrinoParquetDataSource(inputFile, new ParquetReaderOptions(), new FileFormatDataSourceStats()), - Optional.empty()); + blocks = parquetMetadata.getBlocks(); } catch (IOException e) { throw new UncheckedIOException(e); } Comparable previousMax = null; - verify(parquetMetadata.getBlocks().size() > 1, "Test must produce at least two row groups"); - for (BlockMetaData blockMetaData : parquetMetadata.getBlocks()) { - ColumnChunkMetaData columnMetadata = blockMetaData.getColumns().stream() + verify(blocks.size() > 1, "Test must produce at least two row groups"); + for (BlockMetadata blockMetaData : blocks) { + ColumnChunkMetadata columnMetadata = blockMetaData.columns().stream() .filter(column -> getOnlyElement(column.getPath().iterator()).equalsIgnoreCase(sortColumnName)) .collect(onlyElement()); if (previousMax != null) { @@ -152,4 +168,67 @@ public static TrinoFileSystemFactory getFileSystemFactory(DistributedQueryRunner return ((IcebergConnector) queryRunner.getCoordinator().getConnector(ICEBERG_CATALOG)) .getInjector().getInstance(TrinoFileSystemFactory.class); } + + public static HiveMetastore getHiveMetastore(QueryRunner queryRunner) + { + return ((IcebergConnector) ((DistributedQueryRunner) queryRunner).getCoordinator().getConnector(ICEBERG_CATALOG)).getInjector() + .getInstance(HiveMetastoreFactory.class) + .createMetastore(Optional.empty()); + } + + public static BaseTable loadTable(String tableName, + HiveMetastore metastore, + TrinoFileSystemFactory fileSystemFactory, + String catalogName, + String schemaName) + { + IcebergTableOperationsProvider tableOperationsProvider = new FileMetastoreTableOperationsProvider(fileSystemFactory); + CachingHiveMetastore cachingHiveMetastore = createPerTransactionCache(metastore, 1000); + TrinoCatalog catalog = new TrinoHiveCatalog( + new CatalogName(catalogName), + cachingHiveMetastore, + new TrinoViewHiveMetastore(cachingHiveMetastore, false, "trino-version", "test"), + fileSystemFactory, + new TestingTypeManager(), + tableOperationsProvider, + false, + false, + false, + new IcebergConfig().isHideMaterializedViewStorageTable(), + directExecutor()); + return loadIcebergTable(catalog, tableOperationsProvider, SESSION, new SchemaTableName(schemaName, tableName)); + } + + public static TrinoCatalog getTrinoCatalog( + HiveMetastore metastore, + TrinoFileSystemFactory fileSystemFactory, + String catalogName) + { + IcebergTableOperationsProvider tableOperationsProvider = new FileMetastoreTableOperationsProvider(fileSystemFactory); + CachingHiveMetastore cachingHiveMetastore = createPerTransactionCache(metastore, 1000); + return new TrinoHiveCatalog( + new CatalogName(catalogName), + cachingHiveMetastore, + new TrinoViewHiveMetastore(cachingHiveMetastore, false, "trino-version", "test"), + fileSystemFactory, + new TestingTypeManager(), + tableOperationsProvider, + false, + false, + false, + new IcebergConfig().isHideMaterializedViewStorageTable(), + directExecutor()); + } + + public static ParquetMetadata getParquetFileMetadata(TrinoInputFile inputFile) + { + try { + return MetadataReader.readFooter( + new TrinoParquetDataSource(inputFile, new ParquetReaderOptions(), new FileFormatDataSourceStats()), + Optional.empty()); + } + catch (IOException e) { + throw new UncheckedIOException(e); + } + } } diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestConstraintExtractor.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestConstraintExtractor.java index e781df3e2236..9c1ce289396c 100644 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestConstraintExtractor.java +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestConstraintExtractor.java @@ -381,6 +381,7 @@ private static IcebergColumnHandle newPrimitiveColumn(Type type) type, ImmutableList.of(), type, + true, Optional.empty()); } diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergAbfsConnectorSmokeTest.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergAbfsConnectorSmokeTest.java index 73498030c5bd..b5a1f6097b84 100644 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergAbfsConnectorSmokeTest.java +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergAbfsConnectorSmokeTest.java @@ -126,7 +126,7 @@ protected void dropTableFromMetastore(String tableName) HiveMetastore metastore = new BridgingHiveMetastore( testingThriftHiveMetastoreBuilder() .metastoreClient(hiveHadoop.getHiveMetastoreEndpoint()) - .build()); + .build(this::closeAfterClass)); metastore.dropTable(schemaName, tableName, false); assertThat(metastore.getTable(schemaName, tableName)).isEmpty(); } @@ -137,7 +137,7 @@ protected String getMetadataLocation(String tableName) HiveMetastore metastore = new BridgingHiveMetastore( testingThriftHiveMetastoreBuilder() .metastoreClient(hiveHadoop.getHiveMetastoreEndpoint()) - .build()); + .build(this::closeAfterClass)); return metastore .getTable(schemaName, tableName).orElseThrow() .getParameters().get("metadata_location"); diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergColumnHandle.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergColumnHandle.java index 01088be73cba..c2f2e265bc73 100644 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergColumnHandle.java +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergColumnHandle.java @@ -39,7 +39,7 @@ public class TestIcebergColumnHandle @Test public void testRoundTrip() { - testRoundTrip(new IcebergColumnHandle(primitiveColumnIdentity(12, "blah"), BIGINT, ImmutableList.of(), BIGINT, Optional.of("this is a comment"))); + testRoundTrip(new IcebergColumnHandle(primitiveColumnIdentity(12, "blah"), BIGINT, ImmutableList.of(), BIGINT, true, Optional.of("this is a comment"))); // Nested column ColumnIdentity foo1 = new ColumnIdentity(1, "foo1", PRIMITIVE, ImmutableList.of()); @@ -57,6 +57,7 @@ public void testRoundTrip() nestedColumnType, ImmutableList.of(), nestedColumnType, + true, Optional.empty()); testRoundTrip(nestedColumn); @@ -69,6 +70,7 @@ public void testRoundTrip() nestedColumnType, ImmutableList.of(2), BIGINT, + true, Optional.empty()); testRoundTrip(partialColumn); } diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergConfig.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergConfig.java index 5711641500e5..7bf3fc853c20 100644 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergConfig.java +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergConfig.java @@ -13,27 +13,30 @@ */ package io.trino.plugin.iceberg; +import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import io.airlift.units.DataSize; import io.airlift.units.Duration; -import io.trino.plugin.hive.HiveCompressionCodec; -import org.testng.annotations.Test; +import io.trino.plugin.hive.HiveCompressionOption; +import jakarta.validation.constraints.AssertFalse; +import org.junit.jupiter.api.Test; import java.util.Map; import static io.airlift.configuration.testing.ConfigAssertions.assertFullMapping; import static io.airlift.configuration.testing.ConfigAssertions.assertRecordedDefaults; import static io.airlift.configuration.testing.ConfigAssertions.recordDefaults; +import static io.airlift.testing.ValidationAssertions.assertFailsValidation; import static io.airlift.units.DataSize.Unit.GIGABYTE; import static io.airlift.units.DataSize.Unit.MEGABYTE; -import static io.trino.plugin.hive.HiveCompressionCodec.ZSTD; +import static io.trino.plugin.hive.HiveCompressionOption.ZSTD; import static io.trino.plugin.iceberg.CatalogType.GLUE; import static io.trino.plugin.iceberg.CatalogType.HIVE_METASTORE; import static io.trino.plugin.iceberg.IcebergFileFormat.ORC; import static io.trino.plugin.iceberg.IcebergFileFormat.PARQUET; import static java.util.concurrent.TimeUnit.DAYS; import static java.util.concurrent.TimeUnit.HOURS; -import static java.util.concurrent.TimeUnit.MINUTES; +import static java.util.concurrent.TimeUnit.SECONDS; public class TestIcebergConfig { @@ -43,11 +46,12 @@ public void testDefaults() assertRecordedDefaults(recordDefaults(IcebergConfig.class) .setFileFormat(PARQUET) .setCompressionCodec(ZSTD) + .setMaxCommitRetry(4) .setUseFileSizeFromMetadata(true) .setMaxPartitionsPerWriter(100) .setUniqueTableLocation(true) .setCatalogType(HIVE_METASTORE) - .setDynamicFilteringWaitTimeout(new Duration(0, MINUTES)) + .setDynamicFilteringWaitTimeout(new Duration(1, SECONDS)) .setTableStatisticsEnabled(true) .setExtendedStatisticsEnabled(true) .setCollectExtendedStatisticsOnWrite(true) @@ -58,10 +62,24 @@ public void testDefaults() .setRemoveOrphanFilesMinRetention(new Duration(7, DAYS)) .setDeleteSchemaLocationsFallback(false) .setTargetMaxFileSize(DataSize.of(1, GIGABYTE)) + .setIdleWriterMinFileSize(DataSize.of(16, MEGABYTE)) .setMinimumAssignedSplitWeight(0.05) + .setHideMaterializedViewStorageTable(true) .setMaterializedViewsStorageSchema(null) .setRegisterTableProcedureEnabled(false) - .setSortedWritingEnabled(true)); + .setAddFilesProcedureEnabled(false) + .setSortedWritingEnabled(true) + .setQueryPartitionFilterRequired(false) + .setQueryPartitionFilterRequiredSchemas(ImmutableList.of()) + .setSplitManagerThreads(Runtime.getRuntime().availableProcessors() * 2) + .setAllowedExtraProperties(ImmutableList.of()) + .setIncrementalRefreshEnabled(true) + .setMetadataCacheEnabled(true) + .setIncrementalRefreshEnabled(true) + .setObjectStoreLayoutEnabled(false) + .setMetadataParallelism(8) + .setBucketExecutionEnabled(true) + .setFileBasedConflictDetectionEnabled(true)); } @Test @@ -70,6 +88,7 @@ public void testExplicitPropertyMappings() Map properties = ImmutableMap.builder() .put("iceberg.file-format", "ORC") .put("iceberg.compression-codec", "NONE") + .put("iceberg.max-commit-retry", "100") .put("iceberg.use-file-size-from-metadata", "false") .put("iceberg.max-partitions-per-writer", "222") .put("iceberg.unique-table-location", "false") @@ -81,19 +100,33 @@ public void testExplicitPropertyMappings() .put("iceberg.projection-pushdown-enabled", "false") .put("iceberg.hive-catalog-name", "hive") .put("iceberg.format-version", "1") - .put("iceberg.expire_snapshots.min-retention", "13h") - .put("iceberg.remove_orphan_files.min-retention", "14h") + .put("iceberg.expire-snapshots.min-retention", "13h") + .put("iceberg.remove-orphan-files.min-retention", "14h") .put("iceberg.delete-schema-locations-fallback", "true") .put("iceberg.target-max-file-size", "1MB") + .put("iceberg.idle-writer-min-file-size", "1MB") .put("iceberg.minimum-assigned-split-weight", "0.01") + .put("iceberg.materialized-views.hide-storage-table", "false") .put("iceberg.materialized-views.storage-schema", "mv_storage_schema") .put("iceberg.register-table-procedure.enabled", "true") + .put("iceberg.add-files-procedure.enabled", "true") .put("iceberg.sorted-writing-enabled", "false") + .put("iceberg.query-partition-filter-required", "true") + .put("iceberg.query-partition-filter-required-schemas", "bronze,silver") + .put("iceberg.split-manager-threads", "42") + .put("iceberg.allowed-extra-properties", "propX,propY") + .put("iceberg.incremental-refresh-enabled", "false") + .put("iceberg.metadata-cache.enabled", "false") + .put("iceberg.object-store-layout.enabled", "true") + .put("iceberg.metadata.parallelism", "10") + .put("iceberg.bucket-execution", "false") + .put("iceberg.file-based-conflict-detection", "false") .buildOrThrow(); IcebergConfig expected = new IcebergConfig() .setFileFormat(ORC) - .setCompressionCodec(HiveCompressionCodec.NONE) + .setCompressionCodec(HiveCompressionOption.NONE) + .setMaxCommitRetry(100) .setUseFileSizeFromMetadata(false) .setMaxPartitionsPerWriter(222) .setUniqueTableLocation(false) @@ -109,11 +142,37 @@ public void testExplicitPropertyMappings() .setRemoveOrphanFilesMinRetention(new Duration(14, HOURS)) .setDeleteSchemaLocationsFallback(true) .setTargetMaxFileSize(DataSize.of(1, MEGABYTE)) + .setIdleWriterMinFileSize(DataSize.of(1, MEGABYTE)) .setMinimumAssignedSplitWeight(0.01) + .setHideMaterializedViewStorageTable(false) .setMaterializedViewsStorageSchema("mv_storage_schema") .setRegisterTableProcedureEnabled(true) - .setSortedWritingEnabled(false); + .setAddFilesProcedureEnabled(true) + .setSortedWritingEnabled(false) + .setQueryPartitionFilterRequired(true) + .setQueryPartitionFilterRequiredSchemas(ImmutableList.of("bronze", "silver")) + .setSplitManagerThreads(42) + .setAllowedExtraProperties(ImmutableList.of("propX", "propY")) + .setIncrementalRefreshEnabled(false) + .setMetadataCacheEnabled(false) + .setIncrementalRefreshEnabled(false) + .setObjectStoreLayoutEnabled(true) + .setMetadataParallelism(10) + .setBucketExecutionEnabled(false) + .setFileBasedConflictDetectionEnabled(false); assertFullMapping(properties, expected); } + + @Test + public void testValidation() + { + assertFailsValidation( + new IcebergConfig() + .setHideMaterializedViewStorageTable(true) + .setMaterializedViewsStorageSchema("storage_schema"), + "storageSchemaSetWhenHidingIsEnabled", + "iceberg.materialized-views.storage-schema may only be set when iceberg.materialized-views.hide-storage-table is set to false", + AssertFalse.class); + } } diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergMetadataFileOperations.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergFileOperations.java similarity index 50% rename from plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergMetadataFileOperations.java rename to plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergFileOperations.java index 14931abe0efd..27a0fbddbb32 100644 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergMetadataFileOperations.java +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergFileOperations.java @@ -14,46 +14,44 @@ package io.trino.plugin.iceberg; import com.google.common.collect.HashMultiset; +import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableMultiset; import com.google.common.collect.Multiset; -import com.google.inject.Key; import io.trino.Session; +import io.trino.SystemSessionProperties; import io.trino.filesystem.TrackingFileSystemFactory; import io.trino.filesystem.TrackingFileSystemFactory.OperationType; -import io.trino.filesystem.hdfs.HdfsFileSystemFactory; -import io.trino.plugin.hive.metastore.HiveMetastore; -import io.trino.plugin.iceberg.catalog.file.TestingIcebergFileMetastoreCatalogModule; +import io.trino.filesystem.local.LocalFileSystemFactory; import io.trino.plugin.tpch.TpchPlugin; import io.trino.testing.AbstractTestQueryFramework; import io.trino.testing.DistributedQueryRunner; +import io.trino.testing.QueryRunner; import org.apache.iceberg.util.ThreadPools; import org.intellij.lang.annotations.Language; import org.testng.annotations.DataProvider; import org.testng.annotations.Test; -import java.io.File; -import java.util.Optional; +import java.nio.file.Path; +import java.util.function.Predicate; import static com.google.common.collect.ImmutableMultiset.toImmutableMultiset; -import static com.google.inject.multibindings.OptionalBinder.newOptionalBinder; import static io.trino.SystemSessionProperties.MIN_INPUT_SIZE_PER_TASK; import static io.trino.filesystem.TrackingFileSystemFactory.OperationType.INPUT_FILE_GET_LENGTH; import static io.trino.filesystem.TrackingFileSystemFactory.OperationType.INPUT_FILE_NEW_STREAM; import static io.trino.filesystem.TrackingFileSystemFactory.OperationType.OUTPUT_FILE_CREATE; import static io.trino.filesystem.TrackingFileSystemFactory.OperationType.OUTPUT_FILE_CREATE_OR_OVERWRITE; -import static io.trino.filesystem.TrackingFileSystemFactory.OperationType.OUTPUT_FILE_LOCATION; -import static io.trino.plugin.hive.HiveTestUtils.HDFS_ENVIRONMENT; -import static io.trino.plugin.hive.HiveTestUtils.HDFS_FILE_SYSTEM_STATS; -import static io.trino.plugin.hive.metastore.file.TestingFileHiveMetastore.createTestingFileHiveMetastore; -import static io.trino.plugin.hive.util.MultisetAssertions.assertMultisetsEqual; import static io.trino.plugin.iceberg.IcebergQueryRunner.ICEBERG_CATALOG; import static io.trino.plugin.iceberg.IcebergSessionProperties.COLLECT_EXTENDED_STATISTICS_ON_WRITE; -import static io.trino.plugin.iceberg.TestIcebergMetadataFileOperations.FileType.DATA; -import static io.trino.plugin.iceberg.TestIcebergMetadataFileOperations.FileType.MANIFEST; -import static io.trino.plugin.iceberg.TestIcebergMetadataFileOperations.FileType.METADATA_JSON; -import static io.trino.plugin.iceberg.TestIcebergMetadataFileOperations.FileType.SNAPSHOT; -import static io.trino.plugin.iceberg.TestIcebergMetadataFileOperations.FileType.STATS; -import static io.trino.plugin.iceberg.TestIcebergMetadataFileOperations.FileType.fromFilePath; +import static io.trino.plugin.iceberg.TestIcebergFileOperations.FileType.DATA; +import static io.trino.plugin.iceberg.TestIcebergFileOperations.FileType.MANIFEST; +import static io.trino.plugin.iceberg.TestIcebergFileOperations.FileType.METADATA_JSON; +import static io.trino.plugin.iceberg.TestIcebergFileOperations.FileType.METASTORE; +import static io.trino.plugin.iceberg.TestIcebergFileOperations.FileType.SNAPSHOT; +import static io.trino.plugin.iceberg.TestIcebergFileOperations.FileType.STATS; +import static io.trino.plugin.iceberg.TestIcebergFileOperations.FileType.fromFilePath; +import static io.trino.plugin.iceberg.TestIcebergFileOperations.Scope.ALL_FILES; +import static io.trino.plugin.iceberg.TestIcebergFileOperations.Scope.METADATA_FILES; +import static io.trino.testing.MultisetAssertions.assertMultisetsEqual; import static io.trino.testing.TestingNames.randomNameSuffix; import static io.trino.testing.TestingSession.testSessionBuilder; import static java.lang.Math.min; @@ -62,8 +60,8 @@ import static java.util.Objects.requireNonNull; import static java.util.stream.Collectors.toCollection; -@Test(singleThreaded = true) // e.g. trackingFileSystemFactory is shared mutable state -public class TestIcebergMetadataFileOperations +@Test(singleThreaded = true) +public class TestIcebergFileOperations extends AbstractTestQueryFramework { private static final int MAX_PREFIXES_COUNT = 10; @@ -71,7 +69,7 @@ public class TestIcebergMetadataFileOperations private TrackingFileSystemFactory trackingFileSystemFactory; @Override - protected DistributedQueryRunner createQueryRunner() + protected QueryRunner createQueryRunner() throws Exception { Session session = testSessionBuilder() @@ -90,18 +88,13 @@ protected DistributedQueryRunner createQueryRunner() .addCoordinatorProperty("optimizer.experimental-max-prefetched-information-schema-prefixes", Integer.toString(MAX_PREFIXES_COUNT)) .build(); - File baseDir = queryRunner.getCoordinator().getBaseDataDir().resolve("iceberg_data").toFile(); - HiveMetastore metastore = createTestingFileHiveMetastore(baseDir); - - trackingFileSystemFactory = new TrackingFileSystemFactory(new HdfsFileSystemFactory(HDFS_ENVIRONMENT, HDFS_FILE_SYSTEM_STATS)); - queryRunner.installPlugin(new TestingIcebergPlugin( - Optional.of(new TestingIcebergFileMetastoreCatalogModule(metastore)), - Optional.of(trackingFileSystemFactory), - binder -> { - newOptionalBinder(binder, Key.get(boolean.class, AsyncIcebergSplitProducer.class)) - .setBinding().toInstance(false); - })); - queryRunner.createCatalog(ICEBERG_CATALOG, "iceberg"); + Path dataDirectory = queryRunner.getCoordinator().getBaseDataDir().resolve("iceberg_data"); + dataDirectory.toFile().mkdirs(); + trackingFileSystemFactory = new TrackingFileSystemFactory(new LocalFileSystemFactory(dataDirectory)); + queryRunner.installPlugin(new TestingIcebergPlugin(dataDirectory)); + queryRunner.createCatalog(ICEBERG_CATALOG, "iceberg", ImmutableMap.builder() + .put("iceberg.split-manager-threads", "0") + .buildOrThrow()); queryRunner.installPlugin(new TpchPlugin()); queryRunner.createCatalog("tpch", "tpch"); @@ -115,12 +108,30 @@ public void testCreateTable() { assertFileSystemAccesses("CREATE TABLE test_create (id VARCHAR, age INT)", ImmutableMultiset.builder() - .addCopies(new FileOperation(METADATA_JSON, OUTPUT_FILE_CREATE), 1) - .addCopies(new FileOperation(METADATA_JSON, OUTPUT_FILE_LOCATION), 1) - .addCopies(new FileOperation(SNAPSHOT, INPUT_FILE_GET_LENGTH), 1) - .addCopies(new FileOperation(SNAPSHOT, INPUT_FILE_NEW_STREAM), 1) - .addCopies(new FileOperation(SNAPSHOT, OUTPUT_FILE_CREATE_OR_OVERWRITE), 1) - .addCopies(new FileOperation(SNAPSHOT, OUTPUT_FILE_LOCATION), 2) + .add(new FileOperation(METADATA_JSON, OUTPUT_FILE_CREATE)) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_GET_LENGTH)) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_NEW_STREAM)) + .add(new FileOperation(SNAPSHOT, OUTPUT_FILE_CREATE_OR_OVERWRITE)) + .build()); + } + + @Test + public void testCreateOrReplaceTable() + { + assertFileSystemAccesses("CREATE OR REPLACE TABLE test_create_or_replace (id VARCHAR, age INT)", + ImmutableMultiset.builder() + .add(new FileOperation(METADATA_JSON, OUTPUT_FILE_CREATE)) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_GET_LENGTH)) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_NEW_STREAM)) + .add(new FileOperation(SNAPSHOT, OUTPUT_FILE_CREATE_OR_OVERWRITE)) + .build()); + assertFileSystemAccesses("CREATE OR REPLACE TABLE test_create_or_replace (id VARCHAR, age INT)", + ImmutableMultiset.builder() + .add(new FileOperation(METADATA_JSON, OUTPUT_FILE_CREATE)) + .add(new FileOperation(METADATA_JSON, INPUT_FILE_NEW_STREAM)) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_GET_LENGTH)) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_NEW_STREAM)) + .add(new FileOperation(SNAPSHOT, OUTPUT_FILE_CREATE_OR_OVERWRITE)) .build()); } @@ -131,30 +142,53 @@ public void testCreateTableAsSelect() withStatsOnWrite(getSession(), false), "CREATE TABLE test_create_as_select AS SELECT 1 col_name", ImmutableMultiset.builder() - .addCopies(new FileOperation(METADATA_JSON, OUTPUT_FILE_CREATE), 1) - .addCopies(new FileOperation(METADATA_JSON, OUTPUT_FILE_LOCATION), 1) - .addCopies(new FileOperation(SNAPSHOT, INPUT_FILE_GET_LENGTH), 1) - .addCopies(new FileOperation(SNAPSHOT, INPUT_FILE_NEW_STREAM), 1) - .addCopies(new FileOperation(SNAPSHOT, OUTPUT_FILE_CREATE_OR_OVERWRITE), 1) - .addCopies(new FileOperation(SNAPSHOT, OUTPUT_FILE_LOCATION), 2) - .addCopies(new FileOperation(MANIFEST, OUTPUT_FILE_CREATE_OR_OVERWRITE), 1) - .addCopies(new FileOperation(MANIFEST, OUTPUT_FILE_LOCATION), 1) + .add(new FileOperation(METADATA_JSON, OUTPUT_FILE_CREATE)) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_GET_LENGTH)) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_NEW_STREAM)) + .add(new FileOperation(SNAPSHOT, OUTPUT_FILE_CREATE_OR_OVERWRITE)) + .add(new FileOperation(MANIFEST, OUTPUT_FILE_CREATE_OR_OVERWRITE)) .build()); assertFileSystemAccesses( withStatsOnWrite(getSession(), true), "CREATE TABLE test_create_as_select_with_stats AS SELECT 1 col_name", ImmutableMultiset.builder() - .addCopies(new FileOperation(METADATA_JSON, INPUT_FILE_NEW_STREAM), 1) + .add(new FileOperation(METADATA_JSON, INPUT_FILE_NEW_STREAM)) .addCopies(new FileOperation(METADATA_JSON, OUTPUT_FILE_CREATE), 2) // TODO (https://github.com/trinodb/trino/issues/15439): it would be good to publish data and stats in one commit - .addCopies(new FileOperation(METADATA_JSON, OUTPUT_FILE_LOCATION), 2) - .addCopies(new FileOperation(SNAPSHOT, INPUT_FILE_GET_LENGTH), 1) - .addCopies(new FileOperation(SNAPSHOT, INPUT_FILE_NEW_STREAM), 1) - .addCopies(new FileOperation(SNAPSHOT, OUTPUT_FILE_CREATE_OR_OVERWRITE), 1) - .addCopies(new FileOperation(SNAPSHOT, OUTPUT_FILE_LOCATION), 2) - .addCopies(new FileOperation(MANIFEST, OUTPUT_FILE_CREATE_OR_OVERWRITE), 1) - .addCopies(new FileOperation(MANIFEST, OUTPUT_FILE_LOCATION), 1) - .addCopies(new FileOperation(STATS, OUTPUT_FILE_CREATE), 1) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_GET_LENGTH)) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_NEW_STREAM)) + .add(new FileOperation(SNAPSHOT, OUTPUT_FILE_CREATE_OR_OVERWRITE)) + .add(new FileOperation(MANIFEST, OUTPUT_FILE_CREATE_OR_OVERWRITE)) + .add(new FileOperation(STATS, OUTPUT_FILE_CREATE)) + .build()); + } + + @Test + public void testCreateOrReplaceTableAsSelect() + { + assertFileSystemAccesses( + "CREATE OR REPLACE TABLE test_create_or_replace_as_select AS SELECT 1 col_name", + ImmutableMultiset.builder() + .addCopies(new FileOperation(METADATA_JSON, OUTPUT_FILE_CREATE), 2) + .add(new FileOperation(METADATA_JSON, INPUT_FILE_NEW_STREAM)) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_GET_LENGTH)) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_NEW_STREAM)) + .add(new FileOperation(SNAPSHOT, OUTPUT_FILE_CREATE_OR_OVERWRITE)) + .add(new FileOperation(MANIFEST, OUTPUT_FILE_CREATE_OR_OVERWRITE)) + .add(new FileOperation(STATS, OUTPUT_FILE_CREATE)) + .build()); + + assertFileSystemAccesses( + "CREATE OR REPLACE TABLE test_create_or_replace_as_select AS SELECT 1 col_name", + ImmutableMultiset.builder() + .addCopies(new FileOperation(METADATA_JSON, OUTPUT_FILE_CREATE), 2) + .addCopies(new FileOperation(METADATA_JSON, INPUT_FILE_NEW_STREAM), 2) + .addCopies(new FileOperation(SNAPSHOT, INPUT_FILE_NEW_STREAM), 2) + .addCopies(new FileOperation(SNAPSHOT, INPUT_FILE_GET_LENGTH), 2) + .add(new FileOperation(SNAPSHOT, OUTPUT_FILE_CREATE_OR_OVERWRITE)) + .add(new FileOperation(MANIFEST, OUTPUT_FILE_CREATE_OR_OVERWRITE)) + .add(new FileOperation(MANIFEST, INPUT_FILE_NEW_STREAM)) + .add(new FileOperation(STATS, OUTPUT_FILE_CREATE)) .build()); } @@ -164,10 +198,10 @@ public void testSelect() assertUpdate("CREATE TABLE test_select AS SELECT 1 col_name", 1); assertFileSystemAccesses("SELECT * FROM test_select", ImmutableMultiset.builder() - .addCopies(new FileOperation(METADATA_JSON, INPUT_FILE_NEW_STREAM), 1) - .addCopies(new FileOperation(SNAPSHOT, INPUT_FILE_GET_LENGTH), 1) - .addCopies(new FileOperation(SNAPSHOT, INPUT_FILE_NEW_STREAM), 1) - .addCopies(new FileOperation(MANIFEST, INPUT_FILE_NEW_STREAM), 1) + .add(new FileOperation(METADATA_JSON, INPUT_FILE_NEW_STREAM)) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_GET_LENGTH)) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_NEW_STREAM)) + .add(new FileOperation(MANIFEST, INPUT_FILE_NEW_STREAM)) .build()); } @@ -189,32 +223,32 @@ public void testSelectWithLimit(int numberOfFiles) assertFileSystemAccesses("SELECT * FROM test_select_with_limit LIMIT 3", ImmutableMultiset.builder() - .addCopies(new FileOperation(METADATA_JSON, INPUT_FILE_NEW_STREAM), 1) - .addCopies(new FileOperation(SNAPSHOT, INPUT_FILE_GET_LENGTH), 1) - .addCopies(new FileOperation(SNAPSHOT, INPUT_FILE_NEW_STREAM), 1) + .add(new FileOperation(METADATA_JSON, INPUT_FILE_NEW_STREAM)) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_GET_LENGTH)) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_NEW_STREAM)) .addCopies(new FileOperation(MANIFEST, INPUT_FILE_NEW_STREAM), min(icebergManifestPrefetching, numberOfFiles)) .build()); assertFileSystemAccesses("EXPLAIN SELECT * FROM test_select_with_limit LIMIT 3", ImmutableMultiset.builder() - .addCopies(new FileOperation(METADATA_JSON, INPUT_FILE_NEW_STREAM), 1) - .addCopies(new FileOperation(SNAPSHOT, INPUT_FILE_GET_LENGTH), 1) - .addCopies(new FileOperation(SNAPSHOT, INPUT_FILE_NEW_STREAM), 1) + .add(new FileOperation(METADATA_JSON, INPUT_FILE_NEW_STREAM)) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_GET_LENGTH)) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_NEW_STREAM)) .addCopies(new FileOperation(MANIFEST, INPUT_FILE_NEW_STREAM), numberOfFiles) .build()); assertFileSystemAccesses("EXPLAIN ANALYZE SELECT * FROM test_select_with_limit LIMIT 3", ImmutableMultiset.builder() - .addCopies(new FileOperation(METADATA_JSON, INPUT_FILE_NEW_STREAM), 1) - .addCopies(new FileOperation(SNAPSHOT, INPUT_FILE_GET_LENGTH), 1) - .addCopies(new FileOperation(SNAPSHOT, INPUT_FILE_NEW_STREAM), 1) + .add(new FileOperation(METADATA_JSON, INPUT_FILE_NEW_STREAM)) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_GET_LENGTH)) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_NEW_STREAM)) .addCopies(new FileOperation(MANIFEST, INPUT_FILE_NEW_STREAM), numberOfFiles + min(icebergManifestPrefetching, numberOfFiles)) .build()); assertUpdate("DROP TABLE test_select_with_limit"); } - @DataProvider + @DataProvider(name = "testSelectWithLimitDataProvider") public Object[][] testSelectWithLimitDataProvider() { return new Object[][] { @@ -225,6 +259,127 @@ public Object[][] testSelectWithLimitDataProvider() }; } + @Test + public void testReadWholePartition() + { + assertUpdate("DROP TABLE IF EXISTS test_read_part_key"); + + assertUpdate("CREATE TABLE test_read_part_key(key varchar, data varchar) WITH (partitioning=ARRAY['key'])"); + + // Create multiple files per partition + assertUpdate("INSERT INTO test_read_part_key(key, data) VALUES ('p1', '1-abc'), ('p1', '1-def'), ('p2', '2-abc'), ('p2', '2-def')", 4); + assertUpdate("INSERT INTO test_read_part_key(key, data) VALUES ('p1', '1-baz'), ('p2', '2-baz')", 2); + + // Read partition and data columns + assertFileSystemAccesses( + "SELECT key, max(data) FROM test_read_part_key GROUP BY key", + ALL_FILES, + ImmutableMultiset.builder() + .addCopies(new FileOperation(MANIFEST, INPUT_FILE_NEW_STREAM), 2) + .add(new FileOperation(METADATA_JSON, INPUT_FILE_NEW_STREAM)) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_GET_LENGTH)) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_NEW_STREAM)) + .addCopies(new FileOperation(DATA, INPUT_FILE_NEW_STREAM), 4) + .build()); + + // Read partition column only + assertFileSystemAccesses( + "SELECT key, count(*) FROM test_read_part_key GROUP BY key", + ALL_FILES, + ImmutableMultiset.builder() + .addCopies(new FileOperation(MANIFEST, INPUT_FILE_NEW_STREAM), 2) + .add(new FileOperation(METADATA_JSON, INPUT_FILE_NEW_STREAM)) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_GET_LENGTH)) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_NEW_STREAM)) + .build()); + + // Read partition column only, one partition only + assertFileSystemAccesses( + "SELECT count(*) FROM test_read_part_key WHERE key = 'p1'", + ALL_FILES, + ImmutableMultiset.builder() + .addCopies(new FileOperation(MANIFEST, INPUT_FILE_NEW_STREAM), 2) + .add(new FileOperation(METADATA_JSON, INPUT_FILE_NEW_STREAM)) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_GET_LENGTH)) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_NEW_STREAM)) + .build()); + + // Read partition and synthetic columns + assertFileSystemAccesses( + "SELECT count(*), array_agg(\"$path\"), max(\"$file_modified_time\") FROM test_read_part_key GROUP BY key", + ALL_FILES, + ImmutableMultiset.builder() + .addCopies(new FileOperation(MANIFEST, INPUT_FILE_NEW_STREAM), 2) + .add(new FileOperation(METADATA_JSON, INPUT_FILE_NEW_STREAM)) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_GET_LENGTH)) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_NEW_STREAM)) + // TODO return synthetic columns without opening the data files + .addCopies(new FileOperation(DATA, INPUT_FILE_NEW_STREAM), 4) + //.addCopies(new FileOperation(DATA, INPUT_FILE_LAST_MODIFIED), 4) + .build()); + + // Read only row count + assertFileSystemAccesses( + "SELECT count(*) FROM test_read_part_key", + ALL_FILES, + ImmutableMultiset.builder() + .addCopies(new FileOperation(MANIFEST, INPUT_FILE_NEW_STREAM), 2) + .add(new FileOperation(METADATA_JSON, INPUT_FILE_NEW_STREAM)) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_GET_LENGTH)) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_NEW_STREAM)) + .build()); + + assertUpdate("DROP TABLE test_read_part_key"); + } + + @Test + public void testReadWholePartitionSplittableFile() + { + String catalog = getSession().getCatalog().orElseThrow(); + + assertUpdate("DROP TABLE IF EXISTS test_read_whole_splittable_file"); + assertUpdate("CREATE TABLE test_read_whole_splittable_file(key varchar, data varchar) WITH (partitioning=ARRAY['key'])"); + + assertUpdate( + Session.builder(getSession()) + .setSystemProperty(SystemSessionProperties.WRITER_SCALING_MIN_DATA_PROCESSED, "1PB") + .setCatalogSessionProperty(catalog, "parquet_writer_block_size", "1kB") + .setCatalogSessionProperty(catalog, "orc_writer_max_stripe_size", "1kB") + .setCatalogSessionProperty(catalog, "orc_writer_max_stripe_rows", "1000") + .build(), + "INSERT INTO test_read_whole_splittable_file SELECT 'single partition', comment FROM tpch.tiny.orders", 15000); + + Session session = Session.builder(getSession()) + .setCatalogSessionProperty(catalog, IcebergSessionProperties.SPLIT_SIZE, "1kB") + .build(); + + // Read partition column only + assertFileSystemAccesses( + session, + "SELECT key, count(*) FROM test_read_whole_splittable_file GROUP BY key", + ALL_FILES, + ImmutableMultiset.builder() + .add(new FileOperation(MANIFEST, INPUT_FILE_NEW_STREAM)) + .add(new FileOperation(METADATA_JSON, INPUT_FILE_NEW_STREAM)) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_GET_LENGTH)) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_NEW_STREAM)) + .build()); + + // Read only row count + assertFileSystemAccesses( + session, + "SELECT count(*) FROM test_read_whole_splittable_file", + ALL_FILES, + ImmutableMultiset.builder() + .add(new FileOperation(MANIFEST, INPUT_FILE_NEW_STREAM)) + .add(new FileOperation(METADATA_JSON, INPUT_FILE_NEW_STREAM)) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_GET_LENGTH)) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_NEW_STREAM)) + .build()); + + assertUpdate("DROP TABLE test_read_whole_splittable_file"); + } + @Test public void testSelectFromVersionedTable() { @@ -237,29 +392,29 @@ public void testSelectFromVersionedTable() long v3SnapshotId = getLatestSnapshotId(tableName); assertFileSystemAccesses("SELECT * FROM " + tableName + " FOR VERSION AS OF " + v1SnapshotId, ImmutableMultiset.builder() - .addCopies(new FileOperation(METADATA_JSON, INPUT_FILE_NEW_STREAM), 1) - .addCopies(new FileOperation(SNAPSHOT, INPUT_FILE_GET_LENGTH), 1) - .addCopies(new FileOperation(SNAPSHOT, INPUT_FILE_NEW_STREAM), 1) + .add(new FileOperation(METADATA_JSON, INPUT_FILE_NEW_STREAM)) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_GET_LENGTH)) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_NEW_STREAM)) .build()); assertFileSystemAccesses("SELECT * FROM " + tableName + " FOR VERSION AS OF " + v2SnapshotId, ImmutableMultiset.builder() - .addCopies(new FileOperation(METADATA_JSON, INPUT_FILE_NEW_STREAM), 1) - .addCopies(new FileOperation(SNAPSHOT, INPUT_FILE_GET_LENGTH), 1) - .addCopies(new FileOperation(SNAPSHOT, INPUT_FILE_NEW_STREAM), 1) - .addCopies(new FileOperation(MANIFEST, INPUT_FILE_NEW_STREAM), 1) + .add(new FileOperation(METADATA_JSON, INPUT_FILE_NEW_STREAM)) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_GET_LENGTH)) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_NEW_STREAM)) + .add(new FileOperation(MANIFEST, INPUT_FILE_NEW_STREAM)) .build()); assertFileSystemAccesses("SELECT * FROM " + tableName + " FOR VERSION AS OF " + v3SnapshotId, ImmutableMultiset.builder() - .addCopies(new FileOperation(METADATA_JSON, INPUT_FILE_NEW_STREAM), 1) - .addCopies(new FileOperation(SNAPSHOT, INPUT_FILE_GET_LENGTH), 1) - .addCopies(new FileOperation(SNAPSHOT, INPUT_FILE_NEW_STREAM), 1) + .add(new FileOperation(METADATA_JSON, INPUT_FILE_NEW_STREAM)) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_GET_LENGTH)) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_NEW_STREAM)) .addCopies(new FileOperation(MANIFEST, INPUT_FILE_NEW_STREAM), 2) .build()); assertFileSystemAccesses("SELECT * FROM " + tableName, ImmutableMultiset.builder() - .addCopies(new FileOperation(METADATA_JSON, INPUT_FILE_NEW_STREAM), 1) - .addCopies(new FileOperation(SNAPSHOT, INPUT_FILE_GET_LENGTH), 1) - .addCopies(new FileOperation(SNAPSHOT, INPUT_FILE_NEW_STREAM), 1) + .add(new FileOperation(METADATA_JSON, INPUT_FILE_NEW_STREAM)) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_GET_LENGTH)) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_NEW_STREAM)) .addCopies(new FileOperation(MANIFEST, INPUT_FILE_NEW_STREAM), 2) .build()); } @@ -277,29 +432,29 @@ public void testSelectFromVersionedTableWithSchemaEvolution() long v3SnapshotId = getLatestSnapshotId(tableName); assertFileSystemAccesses("SELECT * FROM " + tableName + " FOR VERSION AS OF " + v1SnapshotId, ImmutableMultiset.builder() - .addCopies(new FileOperation(METADATA_JSON, INPUT_FILE_NEW_STREAM), 1) - .addCopies(new FileOperation(SNAPSHOT, INPUT_FILE_GET_LENGTH), 1) - .addCopies(new FileOperation(SNAPSHOT, INPUT_FILE_NEW_STREAM), 1) + .add(new FileOperation(METADATA_JSON, INPUT_FILE_NEW_STREAM)) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_GET_LENGTH)) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_NEW_STREAM)) .build()); assertFileSystemAccesses("SELECT * FROM " + tableName + " FOR VERSION AS OF " + v2SnapshotId, ImmutableMultiset.builder() - .addCopies(new FileOperation(METADATA_JSON, INPUT_FILE_NEW_STREAM), 1) - .addCopies(new FileOperation(SNAPSHOT, INPUT_FILE_GET_LENGTH), 1) - .addCopies(new FileOperation(SNAPSHOT, INPUT_FILE_NEW_STREAM), 1) - .addCopies(new FileOperation(MANIFEST, INPUT_FILE_NEW_STREAM), 1) + .add(new FileOperation(METADATA_JSON, INPUT_FILE_NEW_STREAM)) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_GET_LENGTH)) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_NEW_STREAM)) + .add(new FileOperation(MANIFEST, INPUT_FILE_NEW_STREAM)) .build()); assertFileSystemAccesses("SELECT * FROM " + tableName + " FOR VERSION AS OF " + v3SnapshotId, ImmutableMultiset.builder() - .addCopies(new FileOperation(METADATA_JSON, INPUT_FILE_NEW_STREAM), 1) - .addCopies(new FileOperation(SNAPSHOT, INPUT_FILE_GET_LENGTH), 1) - .addCopies(new FileOperation(SNAPSHOT, INPUT_FILE_NEW_STREAM), 1) + .add(new FileOperation(METADATA_JSON, INPUT_FILE_NEW_STREAM)) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_GET_LENGTH)) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_NEW_STREAM)) .addCopies(new FileOperation(MANIFEST, INPUT_FILE_NEW_STREAM), 2) .build()); assertFileSystemAccesses("SELECT * FROM " + tableName, ImmutableMultiset.builder() - .addCopies(new FileOperation(METADATA_JSON, INPUT_FILE_NEW_STREAM), 1) - .addCopies(new FileOperation(SNAPSHOT, INPUT_FILE_GET_LENGTH), 1) - .addCopies(new FileOperation(SNAPSHOT, INPUT_FILE_NEW_STREAM), 1) + .add(new FileOperation(METADATA_JSON, INPUT_FILE_NEW_STREAM)) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_GET_LENGTH)) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_NEW_STREAM)) .addCopies(new FileOperation(MANIFEST, INPUT_FILE_NEW_STREAM), 2) .build()); } @@ -310,10 +465,10 @@ public void testSelectWithFilter() assertUpdate("CREATE TABLE test_select_with_filter AS SELECT 1 col_name", 1); assertFileSystemAccesses("SELECT * FROM test_select_with_filter WHERE col_name = 1", ImmutableMultiset.builder() - .addCopies(new FileOperation(METADATA_JSON, INPUT_FILE_NEW_STREAM), 1) - .addCopies(new FileOperation(SNAPSHOT, INPUT_FILE_GET_LENGTH), 1) - .addCopies(new FileOperation(SNAPSHOT, INPUT_FILE_NEW_STREAM), 1) - .addCopies(new FileOperation(MANIFEST, INPUT_FILE_NEW_STREAM), 1) + .add(new FileOperation(METADATA_JSON, INPUT_FILE_NEW_STREAM)) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_GET_LENGTH)) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_NEW_STREAM)) + .add(new FileOperation(MANIFEST, INPUT_FILE_NEW_STREAM)) .build()); } @@ -356,10 +511,10 @@ public void testExplainSelect() assertFileSystemAccesses("EXPLAIN SELECT * FROM test_explain", ImmutableMultiset.builder() - .addCopies(new FileOperation(METADATA_JSON, INPUT_FILE_NEW_STREAM), 1) - .addCopies(new FileOperation(SNAPSHOT, INPUT_FILE_GET_LENGTH), 1) - .addCopies(new FileOperation(SNAPSHOT, INPUT_FILE_NEW_STREAM), 1) - .addCopies(new FileOperation(MANIFEST, INPUT_FILE_NEW_STREAM), 1) + .add(new FileOperation(METADATA_JSON, INPUT_FILE_NEW_STREAM)) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_GET_LENGTH)) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_NEW_STREAM)) + .add(new FileOperation(MANIFEST, INPUT_FILE_NEW_STREAM)) .build()); } @@ -370,10 +525,10 @@ public void testShowStatsForTable() assertFileSystemAccesses("SHOW STATS FOR test_show_stats", ImmutableMultiset.builder() - .addCopies(new FileOperation(METADATA_JSON, INPUT_FILE_NEW_STREAM), 1) - .addCopies(new FileOperation(SNAPSHOT, INPUT_FILE_GET_LENGTH), 1) - .addCopies(new FileOperation(SNAPSHOT, INPUT_FILE_NEW_STREAM), 1) - .addCopies(new FileOperation(MANIFEST, INPUT_FILE_NEW_STREAM), 1) + .add(new FileOperation(METADATA_JSON, INPUT_FILE_NEW_STREAM)) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_GET_LENGTH)) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_NEW_STREAM)) + .add(new FileOperation(MANIFEST, INPUT_FILE_NEW_STREAM)) .build()); } @@ -386,10 +541,10 @@ public void testShowStatsForPartitionedTable() assertFileSystemAccesses("SHOW STATS FOR test_show_stats_partitioned", ImmutableMultiset.builder() - .addCopies(new FileOperation(METADATA_JSON, INPUT_FILE_NEW_STREAM), 1) - .addCopies(new FileOperation(SNAPSHOT, INPUT_FILE_GET_LENGTH), 1) - .addCopies(new FileOperation(SNAPSHOT, INPUT_FILE_NEW_STREAM), 1) - .addCopies(new FileOperation(MANIFEST, INPUT_FILE_NEW_STREAM), 1) + .add(new FileOperation(METADATA_JSON, INPUT_FILE_NEW_STREAM)) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_GET_LENGTH)) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_NEW_STREAM)) + .add(new FileOperation(MANIFEST, INPUT_FILE_NEW_STREAM)) .build()); } @@ -400,10 +555,10 @@ public void testShowStatsForTableWithFilter() assertFileSystemAccesses("SHOW STATS FOR (SELECT * FROM test_show_stats_with_filter WHERE age >= 2)", ImmutableMultiset.builder() - .addCopies(new FileOperation(METADATA_JSON, INPUT_FILE_NEW_STREAM), 1) - .addCopies(new FileOperation(SNAPSHOT, INPUT_FILE_GET_LENGTH), 1) - .addCopies(new FileOperation(SNAPSHOT, INPUT_FILE_NEW_STREAM), 1) - .addCopies(new FileOperation(MANIFEST, INPUT_FILE_NEW_STREAM), 1) + .add(new FileOperation(METADATA_JSON, INPUT_FILE_NEW_STREAM)) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_GET_LENGTH)) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_NEW_STREAM)) + .add(new FileOperation(MANIFEST, INPUT_FILE_NEW_STREAM)) .build()); } @@ -416,37 +571,37 @@ public void testPredicateWithVarcharCastToDate() assertFileSystemAccesses("SELECT * FROM test_varchar_as_date_predicate", ImmutableMultiset.builder() - .addCopies(new FileOperation(METADATA_JSON, INPUT_FILE_NEW_STREAM), 1) - .addCopies(new FileOperation(SNAPSHOT, INPUT_FILE_GET_LENGTH), 1) - .addCopies(new FileOperation(SNAPSHOT, INPUT_FILE_NEW_STREAM), 1) + .add(new FileOperation(METADATA_JSON, INPUT_FILE_NEW_STREAM)) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_GET_LENGTH)) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_NEW_STREAM)) .addCopies(new FileOperation(MANIFEST, INPUT_FILE_NEW_STREAM), 2) .build()); // CAST to date and comparison assertFileSystemAccesses("SELECT * FROM test_varchar_as_date_predicate WHERE CAST(a AS date) >= DATE '2005-01-01'", ImmutableMultiset.builder() - .addCopies(new FileOperation(METADATA_JSON, INPUT_FILE_NEW_STREAM), 1) - .addCopies(new FileOperation(SNAPSHOT, INPUT_FILE_GET_LENGTH), 1) - .addCopies(new FileOperation(SNAPSHOT, INPUT_FILE_NEW_STREAM), 1) - .addCopies(new FileOperation(MANIFEST, INPUT_FILE_NEW_STREAM), 1) // fewer than without filter + .add(new FileOperation(METADATA_JSON, INPUT_FILE_NEW_STREAM)) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_GET_LENGTH)) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_NEW_STREAM)) + .add(new FileOperation(MANIFEST, INPUT_FILE_NEW_STREAM)) // fewer than without filter .build()); // CAST to date and BETWEEN assertFileSystemAccesses("SELECT * FROM test_varchar_as_date_predicate WHERE CAST(a AS date) BETWEEN DATE '2005-01-01' AND DATE '2005-12-31'", ImmutableMultiset.builder() - .addCopies(new FileOperation(METADATA_JSON, INPUT_FILE_NEW_STREAM), 1) - .addCopies(new FileOperation(SNAPSHOT, INPUT_FILE_GET_LENGTH), 1) - .addCopies(new FileOperation(SNAPSHOT, INPUT_FILE_NEW_STREAM), 1) - .addCopies(new FileOperation(MANIFEST, INPUT_FILE_NEW_STREAM), 1) // fewer than without filter + .add(new FileOperation(METADATA_JSON, INPUT_FILE_NEW_STREAM)) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_GET_LENGTH)) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_NEW_STREAM)) + .add(new FileOperation(MANIFEST, INPUT_FILE_NEW_STREAM)) // fewer than without filter .build()); // conversion to date as a date function assertFileSystemAccesses("SELECT * FROM test_varchar_as_date_predicate WHERE date(a) >= DATE '2005-01-01'", ImmutableMultiset.builder() - .addCopies(new FileOperation(METADATA_JSON, INPUT_FILE_NEW_STREAM), 1) - .addCopies(new FileOperation(SNAPSHOT, INPUT_FILE_GET_LENGTH), 1) - .addCopies(new FileOperation(SNAPSHOT, INPUT_FILE_NEW_STREAM), 1) - .addCopies(new FileOperation(MANIFEST, INPUT_FILE_NEW_STREAM), 1) // fewer than without filter + .add(new FileOperation(METADATA_JSON, INPUT_FILE_NEW_STREAM)) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_GET_LENGTH)) + .add(new FileOperation(SNAPSHOT, INPUT_FILE_NEW_STREAM)) + .add(new FileOperation(MANIFEST, INPUT_FILE_NEW_STREAM)) // fewer than without filter .build()); assertUpdate("DROP TABLE test_varchar_as_date_predicate"); @@ -495,11 +650,24 @@ public void testInformationSchemaColumns(int tables) assertUpdate(session, "CREATE TABLE test_other_select_i_s_columns" + i + "(id varchar, age integer)"); // won't match the filter } + // Bulk retrieval assertFileSystemAccesses(session, "SELECT * FROM information_schema.columns WHERE table_schema = CURRENT_SCHEMA AND table_name LIKE 'test_select_i_s_columns%'", ImmutableMultiset.builder() .addCopies(new FileOperation(METADATA_JSON, INPUT_FILE_NEW_STREAM), tables * 2) .build()); + // Pointed lookup + assertFileSystemAccesses(session, "SELECT * FROM information_schema.columns WHERE table_schema = CURRENT_SCHEMA AND table_name = 'test_select_i_s_columns0'", + ImmutableMultiset.builder() + .add(new FileOperation(FileType.METADATA_JSON, INPUT_FILE_NEW_STREAM)) + .build()); + + // Pointed lookup via DESCRIBE (which does some additional things before delegating to information_schema.columns) + assertFileSystemAccesses(session, "DESCRIBE test_select_i_s_columns0", + ImmutableMultiset.builder() + .add(new FileOperation(FileType.METADATA_JSON, INPUT_FILE_NEW_STREAM)) + .build()); + for (int i = 0; i < tables; i++) { assertUpdate(session, "DROP TABLE test_select_i_s_columns" + i); assertUpdate(session, "DROP TABLE test_other_select_i_s_columns" + i); @@ -530,10 +698,16 @@ public void testSystemMetadataTableComments(int tables) .addCopies(new FileOperation(METADATA_JSON, INPUT_FILE_NEW_STREAM), tables * 2) .build()); + // Bulk retrieval for two schemas + assertFileSystemAccesses(session, "SELECT * FROM system.metadata.table_comments WHERE schema_name IN (CURRENT_SCHEMA, 'non_existent') AND table_name LIKE 'test_select_s_m_t_comments%'", + ImmutableMultiset.builder() + .addCopies(new FileOperation(METADATA_JSON, INPUT_FILE_NEW_STREAM), tables * 2) + .build()); + // Pointed lookup assertFileSystemAccesses(session, "SELECT * FROM system.metadata.table_comments WHERE schema_name = CURRENT_SCHEMA AND table_name = 'test_select_s_m_t_comments0'", ImmutableMultiset.builder() - .addCopies(new FileOperation(METADATA_JSON, INPUT_FILE_NEW_STREAM), 1) + .add(new FileOperation(METADATA_JSON, INPUT_FILE_NEW_STREAM)) .build()); for (int i = 0; i < tables; i++) { @@ -542,7 +716,7 @@ public void testSystemMetadataTableComments(int tables) } } - @DataProvider + @DataProvider(name = "metadataQueriesTestTableCountDataProvider") public Object[][] metadataQueriesTestTableCountDataProvider() { return new Object[][] { @@ -552,18 +726,90 @@ public Object[][] metadataQueriesTestTableCountDataProvider() }; } + @Test + public void testSystemMetadataMaterializedViews() + { + String schemaName = "test_materialized_views_" + randomNameSuffix(); + assertUpdate("CREATE SCHEMA " + schemaName); + Session session = Session.builder(getSession()) + .setSchema(schemaName) + .build(); + + assertUpdate(session, "CREATE TABLE test_table1 AS SELECT 1 a", 1); + assertUpdate(session, "CREATE TABLE test_table2 AS SELECT 1 a", 1); + + assertUpdate(session, "CREATE MATERIALIZED VIEW mv1 AS SELECT * FROM test_table1 JOIN test_table2 USING (a)"); + assertUpdate(session, "REFRESH MATERIALIZED VIEW mv1", 1); + + assertUpdate(session, "CREATE MATERIALIZED VIEW mv2 AS SELECT count(*) c FROM test_table1 JOIN test_table2 USING (a)"); + assertUpdate(session, "REFRESH MATERIALIZED VIEW mv2", 1); + + // Bulk retrieval + assertFileSystemAccesses(session, "SELECT * FROM system.metadata.materialized_views WHERE schema_name = CURRENT_SCHEMA", + ImmutableMultiset.builder() + .addCopies(new FileOperation(METADATA_JSON, INPUT_FILE_NEW_STREAM), 4) + .build()); + + // Bulk retrieval without selecting freshness + assertFileSystemAccesses( + session, + "SELECT schema_name, name FROM system.metadata.materialized_views WHERE schema_name = CURRENT_SCHEMA", + ImmutableMultiset.of()); + + // Bulk retrieval for two schemas + assertFileSystemAccesses(session, "SELECT * FROM system.metadata.materialized_views WHERE schema_name IN (CURRENT_SCHEMA, 'non_existent')", + ImmutableMultiset.builder() + .addCopies(new FileOperation(METADATA_JSON, INPUT_FILE_NEW_STREAM), 4) + .build()); + + // Pointed lookup + assertFileSystemAccesses(session, "SELECT * FROM system.metadata.materialized_views WHERE schema_name = CURRENT_SCHEMA AND name = 'mv1'", + ImmutableMultiset.builder() + .addCopies(new FileOperation(METADATA_JSON, INPUT_FILE_NEW_STREAM), 3) + .build()); + + // Pointed lookup without selecting freshness + assertFileSystemAccesses( + session, + "SELECT schema_name, name FROM system.metadata.materialized_views WHERE schema_name = CURRENT_SCHEMA AND name = 'mv1'", + ImmutableMultiset.of()); + + assertFileSystemAccesses( + session, + "SELECT * FROM iceberg.information_schema.columns WHERE table_schema = CURRENT_SCHEMA AND table_name = 'mv1'", + ImmutableMultiset.of()); + + assertUpdate("DROP SCHEMA " + schemaName + " CASCADE"); + } + + @Test + public void testShowTables() + { + assertFileSystemAccesses("SHOW TABLES", ImmutableMultiset.of()); + } + private void assertFileSystemAccesses(@Language("SQL") String query, Multiset expectedAccesses) { - assertFileSystemAccesses(getSession(), query, expectedAccesses); + assertFileSystemAccesses(query, METADATA_FILES, expectedAccesses); + } + + private void assertFileSystemAccesses(@Language("SQL") String query, Scope scope, Multiset expectedAccesses) + { + assertFileSystemAccesses(getSession(), query, scope, expectedAccesses); } private void assertFileSystemAccesses(Session session, @Language("SQL") String query, Multiset expectedAccesses) + { + assertFileSystemAccesses(session, query, METADATA_FILES, expectedAccesses); + } + + private void assertFileSystemAccesses(Session session, @Language("SQL") String query, Scope scope, Multiset expectedAccesses) { resetCounts(); getDistributedQueryRunner().executeWithQueryId(session, query); assertMultisetsEqual( getOperations().stream() - .filter(operation -> operation.fileType() != DATA) + .filter(scope) .collect(toImmutableMultiset()), expectedAccesses); } @@ -605,6 +851,25 @@ private record FileOperation(FileType fileType, OperationType operationType) } } + enum Scope + implements Predicate + { + METADATA_FILES { + @Override + public boolean test(FileOperation fileOperation) + { + return fileOperation.fileType() != DATA && fileOperation.fileType() != METASTORE; + } + }, + ALL_FILES { + @Override + public boolean test(FileOperation fileOperation) + { + return fileOperation.fileType() != METASTORE; + } + }, + } + enum FileType { METADATA_JSON, @@ -612,6 +877,7 @@ enum FileType MANIFEST, STATS, DATA, + METASTORE, /**/; public static FileType fromFilePath(String path) @@ -631,6 +897,9 @@ public static FileType fromFilePath(String path) if (path.contains("/data/") && (path.endsWith(".orc") || path.endsWith(".parquet"))) { return DATA; } + if (path.endsWith(".trinoSchema") || path.contains("/.trinoPermissions/")) { + return METASTORE; + } throw new IllegalArgumentException("File not recognized: " + path); } } diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergGcsConnectorSmokeTest.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergGcsConnectorSmokeTest.java index 2c92f12b9d96..11056ace4a02 100644 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergGcsConnectorSmokeTest.java +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergGcsConnectorSmokeTest.java @@ -172,7 +172,7 @@ protected void dropTableFromMetastore(String tableName) HiveMetastore metastore = new BridgingHiveMetastore( testingThriftHiveMetastoreBuilder() .metastoreClient(hiveHadoop.getHiveMetastoreEndpoint()) - .build()); + .build(this::closeAfterClass)); metastore.dropTable(schema, tableName, false); assertThat(metastore.getTable(schema, tableName)).isEmpty(); } @@ -183,7 +183,7 @@ protected String getMetadataLocation(String tableName) HiveMetastore metastore = new BridgingHiveMetastore( testingThriftHiveMetastoreBuilder() .metastoreClient(hiveHadoop.getHiveMetastoreEndpoint()) - .build()); + .build(this::closeAfterClass)); return metastore .getTable(schema, tableName).orElseThrow() .getParameters().get("metadata_location"); diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergGetTableStatisticsOperations.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergGetTableStatisticsOperations.java index 08e56a4e8b3c..e6b7cbbc2032 100644 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergGetTableStatisticsOperations.java +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergGetTableStatisticsOperations.java @@ -23,7 +23,6 @@ import io.trino.metadata.MetadataManager; import io.trino.plugin.hive.metastore.Database; import io.trino.plugin.hive.metastore.HiveMetastore; -import io.trino.plugin.iceberg.catalog.file.TestingIcebergFileMetastoreCatalogModule; import io.trino.plugin.tpch.TpchPlugin; import io.trino.spi.security.PrincipalType; import io.trino.testing.AbstractTestQueryFramework; @@ -42,7 +41,6 @@ import static com.google.common.io.MoreFiles.deleteRecursively; import static com.google.common.io.RecursiveDeleteOption.ALLOW_INSECURE; -import static com.google.inject.util.Modules.EMPTY_MODULE; import static io.trino.execution.querystats.PlanOptimizersStatsCollector.createPlanOptimizersStatsCollector; import static io.trino.plugin.hive.metastore.file.TestingFileHiveMetastore.createTestingFileHiveMetastore; import static io.trino.sql.planner.LogicalPlanner.Stage.OPTIMIZED_AND_VALIDATED; @@ -82,10 +80,11 @@ protected QueryRunner createQueryRunner() localQueryRunner.addFunctions(functions.build()); metastoreDir = Files.createTempDirectory("test_iceberg_get_table_statistics_operations").toFile(); + localQueryRunner.installPlugin(new TestingIcebergPlugin(metastoreDir.toPath())); HiveMetastore metastore = createTestingFileHiveMetastore(metastoreDir); localQueryRunner.createCatalog( "iceberg", - new TestingIcebergConnectorFactory(Optional.of(new TestingIcebergFileMetastoreCatalogModule(metastore)), Optional.empty(), EMPTY_MODULE), + "iceberg", ImmutableMap.of()); Database database = Database.builder() .setDatabaseName("tiny") diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergInputInfo.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergInputInfo.java index 0281aa736dd2..9e52432fafa7 100644 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergInputInfo.java +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergInputInfo.java @@ -22,6 +22,7 @@ import io.trino.tpch.TpchTable; import org.testng.annotations.Test; +import java.util.List; import java.util.Optional; import static io.trino.plugin.iceberg.IcebergQueryRunner.ICEBERG_CATALOG; @@ -40,34 +41,34 @@ protected QueryRunner createQueryRunner() .build(); } - @Test + @Test(enabled = false) public void testInputWithPartitioning() { String tableName = "test_input_info_with_part_" + randomNameSuffix(); assertUpdate("CREATE TABLE " + tableName + " WITH (partitioning = ARRAY['regionkey', 'truncate(name, 1)']) AS SELECT * FROM nation WHERE nationkey < 10", 10); - assertInputInfo(tableName, true, "PARQUET"); + assertInputInfo(tableName, ImmutableList.of("regionkey: identity", "name_trunc: truncate[1]"), "PARQUET", 9); assertUpdate("DROP TABLE " + tableName); } - @Test + @Test(enabled = false) public void testInputWithoutPartitioning() { String tableName = "test_input_info_without_part_" + randomNameSuffix(); assertUpdate("CREATE TABLE " + tableName + " AS SELECT * FROM nation WHERE nationkey < 10", 10); - assertInputInfo(tableName, false, "PARQUET"); + assertInputInfo(tableName, ImmutableList.of(), "PARQUET", 1); assertUpdate("DROP TABLE " + tableName); } - @Test + @Test(enabled = false) public void testInputWithOrcFileFormat() { String tableName = "test_input_info_with_orc_file_format_" + randomNameSuffix(); assertUpdate("CREATE TABLE " + tableName + " WITH (format = 'ORC') AS SELECT * FROM nation WHERE nationkey < 10", 10); - assertInputInfo(tableName, false, "ORC"); + assertInputInfo(tableName, ImmutableList.of(), "ORC", 1); assertUpdate("DROP TABLE " + tableName); } - private void assertInputInfo(String tableName, boolean expectedPartition, String expectedFileFormat) + private void assertInputInfo(String tableName, List partitionFields, String expectedFileFormat, long dataFiles) { inTransaction(session -> { Metadata metadata = getQueryRunner().getMetadata(); @@ -82,8 +83,12 @@ private void assertInputInfo(String tableName, boolean expectedPartition, String IcebergInputInfo icebergInputInfo = (IcebergInputInfo) tableInfo.get(); assertThat(icebergInputInfo).isEqualTo(new IcebergInputInfo( icebergInputInfo.getSnapshotId(), - Optional.of(expectedPartition), - expectedFileFormat)); + partitionFields, + expectedFileFormat, + Optional.of("10"), + Optional.empty(), + Optional.of(String.valueOf(dataFiles)), + Optional.of("0"))); }); } } diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergMaterializedView.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergMaterializedView.java index 59ba415c1dab..c694902ccd80 100644 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergMaterializedView.java +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergMaterializedView.java @@ -14,26 +14,40 @@ package io.trino.plugin.iceberg; import io.trino.Session; +import io.trino.plugin.hive.metastore.HiveMetastore; +import io.trino.plugin.hive.metastore.Table; import io.trino.sql.tree.ExplainType; import io.trino.testing.DistributedQueryRunner; import org.testng.annotations.Test; +import java.io.File; +import java.nio.file.Files; +import java.nio.file.Path; import java.util.Map; import static io.trino.plugin.base.util.Closables.closeAllSuppress; -import static io.trino.plugin.iceberg.IcebergQueryRunner.createIcebergQueryRunner; +import static io.trino.plugin.hive.metastore.file.TestingFileHiveMetastore.createTestingFileHiveMetastore; +import static org.apache.iceberg.BaseMetastoreTableOperations.METADATA_LOCATION_PROP; import static org.assertj.core.api.Assertions.assertThat; public class TestIcebergMaterializedView extends BaseIcebergMaterializedViewTest { private Session secondIceberg; + private String fileMetastoreDirectory; + private HiveMetastore metastore; @Override protected DistributedQueryRunner createQueryRunner() throws Exception { - DistributedQueryRunner queryRunner = createIcebergQueryRunner(); + File metastoreDir = Files.createTempDirectory("test_iceberg_table_smoke_test").toFile(); + metastoreDir.deleteOnExit(); + this.fileMetastoreDirectory = metastoreDir.getAbsolutePath(); + this.metastore = createTestingFileHiveMetastore(metastoreDir); + DistributedQueryRunner queryRunner = IcebergQueryRunner.builder() + .setMetastoreDirectory(metastoreDir) + .build(); try { queryRunner.createCatalog("iceberg2", "iceberg", Map.of( "iceberg.catalog.type", "TESTING_FILE_METASTORE", @@ -56,7 +70,14 @@ protected DistributedQueryRunner createQueryRunner() @Override protected String getSchemaDirectory() { - return getDistributedQueryRunner().getCoordinator().getBaseDataDir().resolve("iceberg_data/tpch").toString(); + return Path.of(fileMetastoreDirectory, "tpch").toString(); + } + + @Override + protected String getStorageMetadataLocation(String materializedViewName) + { + Table table = metastore.getTable("tpch", materializedViewName).orElseThrow(); + return table.getParameters().get(METADATA_LOCATION_PROP); } @Test @@ -83,7 +104,7 @@ AS SELECT sum(value) AS s FROM iceberg2.tpch.common_base_table // After REFRESH, the MV is fresh assertUpdate(defaultIceberg, "REFRESH MATERIALIZED VIEW mv_on_iceberg2", 1); assertThat(getExplainPlan("TABLE mv_on_iceberg2", ExplainType.Type.IO)) - .contains("\"table\" : \"st_") + .contains("\"table\" : \"mv_on_iceberg2$materialized_view_storage") .doesNotContain("common_base_table"); assertThat(query("TABLE mv_on_iceberg2")) .matches("VALUES BIGINT '10'"); @@ -91,7 +112,7 @@ AS SELECT sum(value) AS s FROM iceberg2.tpch.common_base_table // After INSERT to the base table, the MV is still fresh, because it currently does not detect changes to tables in other catalog. assertUpdate(secondIceberg, "INSERT INTO common_base_table VALUES 7", 1); assertThat(getExplainPlan("TABLE mv_on_iceberg2", ExplainType.Type.IO)) - .contains("\"table\" : \"st_") + .contains("\"table\" : \"mv_on_iceberg2$materialized_view_storage") .doesNotContain("common_base_table"); assertThat(query("TABLE mv_on_iceberg2")) .matches("VALUES BIGINT '10'"); @@ -99,7 +120,7 @@ AS SELECT sum(value) AS s FROM iceberg2.tpch.common_base_table // After REFRESH, the MV is fresh again assertUpdate(defaultIceberg, "REFRESH MATERIALIZED VIEW mv_on_iceberg2", 1); assertThat(getExplainPlan("TABLE mv_on_iceberg2", ExplainType.Type.IO)) - .contains("\"table\" : \"st_") + .contains("\"table\" : \"mv_on_iceberg2$materialized_view_storage") .doesNotContain("common_base_table"); assertThat(query("TABLE mv_on_iceberg2")) .matches("VALUES BIGINT '17'"); diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergMergeAppend.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergMergeAppend.java index 9c7435958f5a..812934f585d2 100644 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergMergeAppend.java +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergMergeAppend.java @@ -34,7 +34,8 @@ import java.io.File; -import static io.trino.plugin.hive.metastore.cache.CachingHiveMetastore.memoizeMetastore; +import static com.google.common.util.concurrent.MoreExecutors.directExecutor; +import static io.trino.plugin.hive.metastore.cache.CachingHiveMetastore.createPerTransactionCache; import static io.trino.plugin.hive.metastore.file.TestingFileHiveMetastore.createTestingFileHiveMetastore; import static io.trino.plugin.iceberg.IcebergTestUtils.getFileSystemFactory; import static org.testng.Assert.assertEquals; @@ -54,7 +55,7 @@ protected QueryRunner createQueryRunner() HiveMetastore metastore = createTestingFileHiveMetastore(baseDir); TrinoFileSystemFactory fileSystemFactory = getFileSystemFactory(queryRunner); tableOperationsProvider = new FileMetastoreTableOperationsProvider(fileSystemFactory); - CachingHiveMetastore cachingHiveMetastore = memoizeMetastore(metastore, 1000); + CachingHiveMetastore cachingHiveMetastore = createPerTransactionCache(metastore, 1000); trinoCatalog = new TrinoHiveCatalog( new CatalogName("catalog"), cachingHiveMetastore, @@ -64,7 +65,9 @@ protected QueryRunner createQueryRunner() tableOperationsProvider, false, false, - false); + false, + new IcebergConfig().isHideMaterializedViewStorageTable(), + directExecutor()); return queryRunner; } diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergMetadataListing.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergMetadataListing.java index 5ca8525d54d3..b92a348922d5 100644 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergMetadataListing.java +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergMetadataListing.java @@ -19,7 +19,6 @@ import io.trino.metadata.QualifiedObjectName; import io.trino.plugin.hive.TestingHivePlugin; import io.trino.plugin.hive.metastore.file.FileHiveMetastore; -import io.trino.plugin.iceberg.catalog.file.TestingIcebergFileMetastoreCatalogModule; import io.trino.spi.connector.SchemaTableName; import io.trino.spi.security.Identity; import io.trino.spi.security.SelectedRole; @@ -34,7 +33,6 @@ import java.io.File; import java.util.Optional; -import static com.google.inject.util.Modules.EMPTY_MODULE; import static io.trino.plugin.hive.metastore.file.TestingFileHiveMetastore.createTestingFileHiveMetastore; import static io.trino.spi.security.SelectedRole.Type.ROLE; import static io.trino.testing.TestingSession.testSessionBuilder; @@ -44,7 +42,6 @@ public class TestIcebergMetadataListing extends AbstractTestQueryFramework { private FileHiveMetastore metastore; - private SchemaTableName storageTable; @Override protected DistributedQueryRunner createQueryRunner() @@ -61,9 +58,9 @@ protected DistributedQueryRunner createQueryRunner() metastore = createTestingFileHiveMetastore(baseDir); - queryRunner.installPlugin(new TestingIcebergPlugin(Optional.of(new TestingIcebergFileMetastoreCatalogModule(metastore)), Optional.empty(), EMPTY_MODULE)); + queryRunner.installPlugin(new TestingIcebergPlugin(baseDir.toPath())); queryRunner.createCatalog("iceberg", "iceberg"); - queryRunner.installPlugin(new TestingHivePlugin(metastore)); + queryRunner.installPlugin(new TestingHivePlugin(baseDir.toPath(), metastore)); queryRunner.createCatalog("hive", "hive", ImmutableMap.of("hive.security", "sql-standard")); return queryRunner; @@ -77,7 +74,6 @@ public void setUp() assertQuerySucceeds("CREATE TABLE iceberg.test_schema.iceberg_table2 (_double DOUBLE) WITH (partitioning = ARRAY['_double'])"); assertQuerySucceeds("CREATE MATERIALIZED VIEW iceberg.test_schema.iceberg_materialized_view AS " + "SELECT * FROM iceberg.test_schema.iceberg_table1"); - storageTable = getStorageTable("iceberg", "test_schema", "iceberg_materialized_view"); assertQuerySucceeds("CREATE VIEW iceberg.test_schema.iceberg_view AS SELECT * FROM iceberg.test_schema.iceberg_table1"); assertQuerySucceeds("CREATE TABLE hive.test_schema.hive_table (_double DOUBLE)"); @@ -99,12 +95,12 @@ public void tearDown() @Test public void testTableListing() { - assertThat(metastore.getAllTables("test_schema")) + assertThat(metastore.getTables("test_schema")) + .extracting(table -> table.tableName().getTableName()) .containsExactlyInAnyOrder( "iceberg_table1", "iceberg_table2", "iceberg_materialized_view", - storageTable.getTableName(), "iceberg_view", "hive_table", "hive_view"); @@ -115,7 +111,6 @@ public void testTableListing() "'iceberg_table1', " + "'iceberg_table2', " + "'iceberg_materialized_view', " + - "'" + storageTable.getTableName() + "', " + "'iceberg_view', " + "'hive_table', " + "'hive_view'"); @@ -133,8 +128,6 @@ public void testTableColumnListing() "('iceberg_table2', '_double'), " + "('iceberg_materialized_view', '_string'), " + "('iceberg_materialized_view', '_integer'), " + - "('" + storageTable.getTableName() + "', '_string'), " + - "('" + storageTable.getTableName() + "', '_integer'), " + "('iceberg_view', '_string'), " + "('iceberg_view', '_integer'), " + "('hive_view', '_double')"); diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergMetastoreAccessOperations.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergMetastoreAccessOperations.java index 0706f08e4e7b..ab928923a95b 100644 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergMetastoreAccessOperations.java +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergMetastoreAccessOperations.java @@ -28,20 +28,25 @@ import java.io.File; import java.util.Optional; -import static com.google.inject.util.Modules.EMPTY_MODULE; -import static io.trino.plugin.hive.metastore.CountingAccessHiveMetastore.Method.CREATE_TABLE; -import static io.trino.plugin.hive.metastore.CountingAccessHiveMetastore.Method.DROP_TABLE; -import static io.trino.plugin.hive.metastore.CountingAccessHiveMetastore.Method.GET_ALL_TABLES_FROM_DATABASE; -import static io.trino.plugin.hive.metastore.CountingAccessHiveMetastore.Method.GET_DATABASE; -import static io.trino.plugin.hive.metastore.CountingAccessHiveMetastore.Method.GET_TABLE; -import static io.trino.plugin.hive.metastore.CountingAccessHiveMetastore.Method.GET_TABLE_WITH_PARAMETER; -import static io.trino.plugin.hive.metastore.CountingAccessHiveMetastore.Method.REPLACE_TABLE; +import static io.trino.plugin.hive.metastore.MetastoreMethod.CREATE_TABLE; +import static io.trino.plugin.hive.metastore.MetastoreMethod.DROP_TABLE; +import static io.trino.plugin.hive.metastore.MetastoreMethod.GET_ALL_DATABASES; +import static io.trino.plugin.hive.metastore.MetastoreMethod.GET_ALL_TABLES_FROM_DATABASE; +import static io.trino.plugin.hive.metastore.MetastoreMethod.GET_DATABASE; +import static io.trino.plugin.hive.metastore.MetastoreMethod.GET_TABLE; +import static io.trino.plugin.hive.metastore.MetastoreMethod.GET_TABLE_WITH_PARAMETER; +import static io.trino.plugin.hive.metastore.MetastoreMethod.REPLACE_TABLE; import static io.trino.plugin.hive.metastore.file.TestingFileHiveMetastore.createTestingFileHiveMetastore; import static io.trino.plugin.iceberg.IcebergSessionProperties.COLLECT_EXTENDED_STATISTICS_ON_WRITE; +import static io.trino.plugin.iceberg.TableType.ALL_ENTRIES; +import static io.trino.plugin.iceberg.TableType.ALL_MANIFESTS; import static io.trino.plugin.iceberg.TableType.DATA; +import static io.trino.plugin.iceberg.TableType.ENTRIES; import static io.trino.plugin.iceberg.TableType.FILES; import static io.trino.plugin.iceberg.TableType.HISTORY; import static io.trino.plugin.iceberg.TableType.MANIFESTS; +import static io.trino.plugin.iceberg.TableType.MATERIALIZED_VIEW_STORAGE; +import static io.trino.plugin.iceberg.TableType.METADATA_LOG_ENTRIES; import static io.trino.plugin.iceberg.TableType.PARTITIONS; import static io.trino.plugin.iceberg.TableType.PROPERTIES; import static io.trino.plugin.iceberg.TableType.REFS; @@ -72,7 +77,7 @@ protected DistributedQueryRunner createQueryRunner() File baseDir = queryRunner.getCoordinator().getBaseDataDir().resolve("iceberg_data").toFile(); metastore = new CountingAccessHiveMetastore(createTestingFileHiveMetastore(baseDir)); - queryRunner.installPlugin(new TestingIcebergPlugin(Optional.of(new TestingIcebergFileMetastoreCatalogModule(metastore)), Optional.empty(), EMPTY_MODULE)); + queryRunner.installPlugin(new TestingIcebergPlugin(baseDir.toPath(), Optional.of(new TestingIcebergFileMetastoreCatalogModule(metastore)))); queryRunner.createCatalog("iceberg", "iceberg"); queryRunner.execute("CREATE SCHEMA test_schema"); @@ -182,7 +187,7 @@ public void testSelectFromMaterializedView() assertMetastoreInvocations("SELECT * FROM test_select_mview_view", ImmutableMultiset.builder() - .addCopies(GET_TABLE, 3) + .addCopies(GET_TABLE, 2) .build()); } @@ -194,7 +199,7 @@ public void testSelectFromMaterializedViewWithFilter() assertMetastoreInvocations("SELECT * FROM test_select_mview_where_view WHERE age = 2", ImmutableMultiset.builder() - .addCopies(GET_TABLE, 3) + .addCopies(GET_TABLE, 2) .build()); } @@ -206,7 +211,7 @@ public void testRefreshMaterializedView() assertMetastoreInvocations("REFRESH MATERIALIZED VIEW test_refresh_mview_view", ImmutableMultiset.builder() - .addCopies(GET_TABLE, 6) + .addCopies(GET_TABLE, 2) .addCopies(REPLACE_TABLE, 1) .build()); } @@ -308,9 +313,12 @@ public void testSelectSystemTable() .addCopies(GET_TABLE, 1) .build()); + assertQueryFails("SELECT * FROM \"test_select_snapshots$materialized_view_storage\"", + "Table 'test_schema.test_select_snapshots\\$materialized_view_storage' not found"); + // This test should get updated if a new system table is added. assertThat(TableType.values()) - .containsExactly(DATA, HISTORY, SNAPSHOTS, MANIFESTS, PARTITIONS, FILES, PROPERTIES, REFS); + .containsExactly(DATA, HISTORY, METADATA_LOG_ENTRIES, SNAPSHOTS, ALL_MANIFESTS, MANIFESTS, PARTITIONS, FILES, ALL_ENTRIES, ENTRIES, PROPERTIES, REFS, MATERIALIZED_VIEW_STORAGE); } @Test @@ -344,11 +352,24 @@ public void testInformationSchemaColumns(int tables) assertUpdate(session, "CREATE TABLE test_other_select_i_s_columns" + i + "(id varchar, age integer)"); // won't match the filter } + // Bulk retrieval assertMetastoreInvocations(session, "SELECT * FROM information_schema.columns WHERE table_schema = CURRENT_SCHEMA AND table_name LIKE 'test_select_i_s_columns%'", ImmutableMultiset.builder() .add(GET_ALL_TABLES_FROM_DATABASE) .addCopies(GET_TABLE, tables * 2) - .addCopies(GET_TABLE_WITH_PARAMETER, 2) + .build()); + + // Pointed lookup + assertMetastoreInvocations(session, "SELECT * FROM information_schema.columns WHERE table_schema = CURRENT_SCHEMA AND table_name = 'test_select_i_s_columns0'", + ImmutableMultiset.builder() + .add(GET_TABLE) + .build()); + + // Pointed lookup via DESCRIBE (which does some additional things before delegating to information_schema.columns) + assertMetastoreInvocations(session, "DESCRIBE test_select_i_s_columns0", + ImmutableMultiset.builder() + .add(GET_DATABASE) + .add(GET_TABLE) .build()); for (int i = 0; i < tables; i++) { @@ -405,6 +426,71 @@ public Object[][] metadataQueriesTestTableCountDataProvider() }; } + @Test + public void testSystemMetadataMaterializedViews() + { + String schemaName = "test_materialized_views_" + randomNameSuffix(); + assertUpdate("CREATE SCHEMA " + schemaName); + Session session = Session.builder(getSession()) + .setSchema(schemaName) + .build(); + + assertUpdate(session, "CREATE TABLE test_table1 AS SELECT 1 a", 1); + assertUpdate(session, "CREATE TABLE test_table2 AS SELECT 1 a", 1); + + assertUpdate(session, "CREATE MATERIALIZED VIEW mv1 AS SELECT * FROM test_table1 JOIN test_table2 USING (a)"); + assertUpdate(session, "REFRESH MATERIALIZED VIEW mv1", 1); + + assertUpdate(session, "CREATE MATERIALIZED VIEW mv2 AS SELECT count(*) c FROM test_table1 JOIN test_table2 USING (a)"); + assertUpdate(session, "REFRESH MATERIALIZED VIEW mv2", 1); + + // Bulk retrieval + assertMetastoreInvocations(session, "SELECT * FROM system.metadata.materialized_views WHERE schema_name = CURRENT_SCHEMA", + ImmutableMultiset.builder() + .add(GET_ALL_TABLES_FROM_DATABASE) + .addCopies(GET_TABLE, 4) + .build()); + + // Bulk retrieval without selecting freshness + assertMetastoreInvocations(session, "SELECT schema_name, name FROM system.metadata.materialized_views WHERE schema_name = CURRENT_SCHEMA", + ImmutableMultiset.builder() + .add(GET_ALL_TABLES_FROM_DATABASE) + .addCopies(GET_TABLE, 4) + .build()); + + // Bulk retrieval for two schemas + assertMetastoreInvocations(session, "SELECT * FROM system.metadata.materialized_views WHERE schema_name IN (CURRENT_SCHEMA, 'non_existent')", + ImmutableMultiset.builder() + .add(GET_ALL_DATABASES) + .addCopies(GET_ALL_TABLES_FROM_DATABASE, 2) + .addCopies(GET_TABLE, 4) + .build()); + + // Pointed lookup + assertMetastoreInvocations(session, "SELECT * FROM system.metadata.materialized_views WHERE schema_name = CURRENT_SCHEMA AND name = 'mv1'", + ImmutableMultiset.builder() + .addCopies(GET_TABLE, 3) + .build()); + + // Pointed lookup without selecting freshness + assertMetastoreInvocations(session, "SELECT schema_name, name FROM system.metadata.materialized_views WHERE schema_name = CURRENT_SCHEMA AND name = 'mv1'", + ImmutableMultiset.builder() + .addCopies(GET_TABLE, 3) + .build()); + + assertUpdate("DROP SCHEMA " + schemaName + " CASCADE"); + } + + @Test + public void testShowTables() + { + assertMetastoreInvocations("SHOW TABLES", + ImmutableMultiset.builder() + .add(GET_DATABASE) + .add(GET_ALL_TABLES_FROM_DATABASE) + .build()); + } + private void assertMetastoreInvocations(@Language("SQL") String query, Multiset expectedInvocations) { assertMetastoreInvocations(getSession(), query, expectedInvocations); diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergMigrateProcedure.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergMigrateProcedure.java index 4fcf249e78da..674696439d7d 100644 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergMigrateProcedure.java +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergMigrateProcedure.java @@ -45,7 +45,7 @@ protected QueryRunner createQueryRunner() { dataDirectory = Files.createTempDirectory("_test_hidden"); DistributedQueryRunner queryRunner = IcebergQueryRunner.builder().setMetastoreDirectory(dataDirectory.toFile()).build(); - queryRunner.installPlugin(new TestingHivePlugin()); + queryRunner.installPlugin(new TestingHivePlugin(Files.createTempDirectory(null))); queryRunner.createCatalog("hive", "hive", ImmutableMap.builder() .put("hive.metastore", "file") .put("hive.metastore.catalog.dir", dataDirectory.toString()) diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergNodeLocalDynamicSplitPruning.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergNodeLocalDynamicSplitPruning.java index 019732d1af00..1149c8bf6f76 100644 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergNodeLocalDynamicSplitPruning.java +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergNodeLocalDynamicSplitPruning.java @@ -35,6 +35,7 @@ import io.trino.plugin.hive.orc.OrcWriterConfig; import io.trino.plugin.hive.parquet.ParquetReaderConfig; import io.trino.plugin.hive.parquet.ParquetWriterConfig; +import io.trino.plugin.iceberg.catalog.rest.DefaultIcebergFileSystemFactory; import io.trino.spi.Page; import io.trino.spi.SplitWeight; import io.trino.spi.block.BlockBuilder; @@ -43,7 +44,11 @@ import io.trino.spi.connector.ConnectorPageSource; import io.trino.spi.connector.DynamicFilter; import io.trino.spi.predicate.Domain; +import io.trino.spi.predicate.Range; import io.trino.spi.predicate.TupleDomain; +import io.trino.spi.predicate.ValueSet; +import io.trino.spi.type.DecimalType; +import io.trino.spi.type.SqlDecimal; import io.trino.spi.type.Type; import io.trino.testing.TestingConnectorSession; import org.apache.iceberg.PartitionSpec; @@ -54,7 +59,9 @@ import org.testng.annotations.Test; import java.io.IOException; +import java.math.BigDecimal; import java.nio.file.Files; +import java.time.LocalDate; import java.util.List; import java.util.Map; import java.util.Optional; @@ -69,15 +76,16 @@ import static io.trino.plugin.hive.HiveType.HIVE_STRING; import static io.trino.plugin.iceberg.ColumnIdentity.TypeCategory.PRIMITIVE; import static io.trino.plugin.iceberg.IcebergFileFormat.ORC; +import static io.trino.plugin.iceberg.TypeConverter.toOrcType; +import static io.trino.spi.type.DateType.DATE; +import static io.trino.spi.type.Decimals.writeShortDecimal; import static io.trino.spi.type.IntegerType.INTEGER; import static io.trino.spi.type.VarcharType.VARCHAR; import static io.trino.testing.TestingHandles.TEST_CATALOG_HANDLE; import static io.trino.type.InternalTypeManager.TESTING_TYPE_MANAGER; import static java.util.concurrent.CompletableFuture.completedFuture; import static org.apache.iceberg.types.Types.NestedField.optional; -import static org.testng.Assert.assertEquals; -import static org.testng.Assert.assertNotNull; -import static org.testng.Assert.assertNull; +import static org.assertj.core.api.Assertions.assertThat; public class TestIcebergNodeLocalDynamicSplitPruning { @@ -85,11 +93,11 @@ public class TestIcebergNodeLocalDynamicSplitPruning private static final String TABLE_NAME = "test"; private static final Column KEY_COLUMN = new Column("a_integer", HIVE_INT, Optional.empty()); private static final ColumnIdentity KEY_COLUMN_IDENTITY = new ColumnIdentity(1, KEY_COLUMN.getName(), PRIMITIVE, ImmutableList.of()); - private static final IcebergColumnHandle KEY_ICEBERG_COLUMN_HANDLE = new IcebergColumnHandle(KEY_COLUMN_IDENTITY, INTEGER, ImmutableList.of(), INTEGER, Optional.empty()); + private static final IcebergColumnHandle KEY_ICEBERG_COLUMN_HANDLE = new IcebergColumnHandle(KEY_COLUMN_IDENTITY, INTEGER, ImmutableList.of(), INTEGER, false, Optional.empty()); private static final int KEY_COLUMN_VALUE = 42; private static final Column DATA_COLUMN = new Column("a_varchar", HIVE_STRING, Optional.empty()); private static final ColumnIdentity DATA_COLUMN_IDENTITY = new ColumnIdentity(2, DATA_COLUMN.getName(), PRIMITIVE, ImmutableList.of()); - private static final IcebergColumnHandle DATA_ICEBERG_COLUMN_HANDLE = new IcebergColumnHandle(DATA_COLUMN_IDENTITY, VARCHAR, ImmutableList.of(), VARCHAR, Optional.empty()); + private static final IcebergColumnHandle DATA_ICEBERG_COLUMN_HANDLE = new IcebergColumnHandle(DATA_COLUMN_IDENTITY, VARCHAR, ImmutableList.of(), VARCHAR, false, Optional.empty()); private static final String DATA_COLUMN_VALUE = "hello world"; private static final Schema TABLE_SCHEMA = new Schema( optional(KEY_COLUMN_IDENTITY.getId(), KEY_COLUMN.getName(), Types.IntegerType.get()), @@ -100,9 +108,22 @@ public class TestIcebergNodeLocalDynamicSplitPruning private static final ParquetWriterConfig PARQUET_WRITER_CONFIG = new ParquetWriterConfig(); @Test - public void testDynamicSplitPruning() + public void testDynamicSplitPruningOnUnpartitionedTable() throws IOException { + String tableName = "unpartitioned_table"; + String keyColumnName = "a_integer"; + ColumnIdentity keyColumnIdentity = new ColumnIdentity(1, keyColumnName, PRIMITIVE, ImmutableList.of()); + IcebergColumnHandle keyColumnHandle = new IcebergColumnHandle(keyColumnIdentity, INTEGER, ImmutableList.of(), INTEGER, true, Optional.empty()); + int keyColumnValue = 42; + String dataColumnName = "a_varchar"; + ColumnIdentity dataColumnIdentity = new ColumnIdentity(2, dataColumnName, PRIMITIVE, ImmutableList.of()); + IcebergColumnHandle dataColumnHandle = new IcebergColumnHandle(dataColumnIdentity, VARCHAR, ImmutableList.of(), VARCHAR, true, Optional.empty()); + String dataColumnValue = "hello world"; + Schema tableSchema = new Schema( + optional(keyColumnIdentity.getId(), keyColumnName, Types.IntegerType.get()), + optional(dataColumnIdentity.getId(), dataColumnName, Types.StringType.get())); + IcebergConfig icebergConfig = new IcebergConfig(); HiveTransactionHandle transaction = new HiveTransactionHandle(false); try (TempFile file = new TempFile()) { @@ -110,19 +131,399 @@ public void testDynamicSplitPruning() TrinoOutputFile outputFile = new LocalOutputFile(file.file()); TrinoInputFile inputFile = new LocalInputFile(file.file()); - writeOrcContent(outputFile); + List columnNames = ImmutableList.of(keyColumnName, dataColumnName); + List types = ImmutableList.of(INTEGER, VARCHAR); + + try (OrcWriter writer = new OrcWriter( + OutputStreamOrcDataSink.create(outputFile), + columnNames, + types, + toOrcType(tableSchema), + NONE, + new OrcWriterOptions(), + ImmutableMap.of(), + true, + OrcWriteValidation.OrcWriteValidationMode.BOTH, + new OrcWriterStats())) { + BlockBuilder keyBuilder = INTEGER.createBlockBuilder(null, 1); + INTEGER.writeLong(keyBuilder, keyColumnValue); + BlockBuilder dataBuilder = VARCHAR.createBlockBuilder(null, 1); + VARCHAR.writeString(dataBuilder, dataColumnValue); + writer.write(new Page(keyBuilder.build(), dataBuilder.build())); + } + + IcebergSplit split = new IcebergSplit( + inputFile.toString(), + 0, + inputFile.length(), + inputFile.length(), + -1, // invalid; normally known + ORC, + PartitionSpecParser.toJson(PartitionSpec.unpartitioned()), + PartitionData.toJson(new PartitionData(new Object[] {})), + ImmutableList.of(), + SplitWeight.standard(), + TupleDomain.all(), + ImmutableMap.of(), + 0); - try (ConnectorPageSource emptyPageSource = createTestingPageSource(transaction, icebergConfig, inputFile, getDynamicFilter(getTupleDomainForSplitPruning()))) { - assertNull(emptyPageSource.getNextPage()); + String tablePath = inputFile.location().fileName(); + TableHandle tableHandle = new TableHandle( + TEST_CATALOG_HANDLE, + new IcebergTableHandle( + CatalogHandle.fromId("iceberg:NORMAL:v12345"), + "test_schema", + tableName, + TableType.DATA, + Optional.empty(), + SchemaParser.toJson(tableSchema), + Optional.of(PartitionSpecParser.toJson(PartitionSpec.unpartitioned())), + 2, + TupleDomain.withColumnDomains(ImmutableMap.of(keyColumnHandle, Domain.singleValue(INTEGER, (long) keyColumnValue))), + TupleDomain.all(), + OptionalLong.empty(), + ImmutableSet.of(keyColumnHandle), + Optional.empty(), + tablePath, + ImmutableMap.of(), + Optional.empty(), + false, + Optional.empty(), + ImmutableSet.of(), + Optional.of(false)), + transaction); + + TupleDomain splitPruningPredicate = TupleDomain.withColumnDomains( + ImmutableMap.of( + keyColumnHandle, + Domain.singleValue(INTEGER, 1L))); + try (ConnectorPageSource emptyPageSource = createTestingPageSource(transaction, icebergConfig, split, tableHandle, ImmutableList.of(keyColumnHandle, dataColumnHandle), getDynamicFilter(splitPruningPredicate))) { + assertThat(emptyPageSource.getNextPage()).isNull(); } - try (ConnectorPageSource nonEmptyPageSource = createTestingPageSource(transaction, icebergConfig, inputFile, getDynamicFilter(getNonSelectiveTupleDomain()))) { + TupleDomain nonSelectivePredicate = TupleDomain.withColumnDomains( + ImmutableMap.of( + keyColumnHandle, + Domain.singleValue(INTEGER, (long) keyColumnValue))); + try (ConnectorPageSource nonEmptyPageSource = createTestingPageSource(transaction, icebergConfig, split, tableHandle, ImmutableList.of(keyColumnHandle, dataColumnHandle), getDynamicFilter(nonSelectivePredicate))) { Page page = nonEmptyPageSource.getNextPage(); - assertNotNull(page); - assertEquals(page.getBlock(0).getPositionCount(), 1); - assertEquals(page.getBlock(0).getInt(0, 0), KEY_COLUMN_VALUE); - assertEquals(page.getBlock(1).getPositionCount(), 1); - assertEquals(page.getBlock(1).getSlice(0, 0, page.getBlock(1).getSliceLength(0)).toStringUtf8(), DATA_COLUMN_VALUE); + assertThat(page).isNotNull(); + assertThat(page.getPositionCount()).isEqualTo(1); + assertThat(page.getBlock(0).getInt(0, 0)).isEqualTo(keyColumnValue); + assertThat(page.getBlock(1).getSlice(0, 0, page.getBlock(1).getSliceLength(0)).toStringUtf8()).isEqualTo(dataColumnValue); + } + } + } + + @Test + public void testDynamicSplitPruningWithExplicitPartitionFilter() + throws IOException + { + String tableName = "sales_table"; + String dateColumnName = "date"; + ColumnIdentity dateColumnIdentity = new ColumnIdentity(1, dateColumnName, PRIMITIVE, ImmutableList.of()); + IcebergColumnHandle dateColumnHandle = new IcebergColumnHandle(dateColumnIdentity, DATE, ImmutableList.of(), DATE, true, Optional.empty()); + long dateColumnValue = LocalDate.of(2023, 1, 10).toEpochDay(); + String receiptColumnName = "receipt"; + ColumnIdentity receiptColumnIdentity = new ColumnIdentity(2, receiptColumnName, PRIMITIVE, ImmutableList.of()); + IcebergColumnHandle receiptColumnHandle = new IcebergColumnHandle(receiptColumnIdentity, VARCHAR, ImmutableList.of(), VARCHAR, true, Optional.empty()); + String receiptColumnValue = "#12345"; + String amountColumnName = "amount"; + ColumnIdentity amountColumnIdentity = new ColumnIdentity(3, amountColumnName, PRIMITIVE, ImmutableList.of()); + DecimalType amountColumnType = DecimalType.createDecimalType(10, 2); + IcebergColumnHandle amountColumnHandle = new IcebergColumnHandle(amountColumnIdentity, amountColumnType, ImmutableList.of(), amountColumnType, true, Optional.empty()); + BigDecimal amountColumnValue = new BigDecimal("1234567.65"); + Schema tableSchema = new Schema( + optional(dateColumnIdentity.getId(), dateColumnName, Types.DateType.get()), + optional(receiptColumnIdentity.getId(), receiptColumnName, Types.StringType.get()), + optional(amountColumnIdentity.getId(), amountColumnName, Types.DecimalType.of(10, 2))); + PartitionSpec partitionSpec = PartitionSpec.builderFor(tableSchema) + .identity(dateColumnName) + .build(); + + IcebergConfig icebergConfig = new IcebergConfig(); + HiveTransactionHandle transaction = new HiveTransactionHandle(false); + try (TempFile file = new TempFile()) { + Files.delete(file.path()); + + TrinoOutputFile outputFile = new LocalOutputFile(file.file()); + TrinoInputFile inputFile = new LocalInputFile(file.file()); + List columnNames = ImmutableList.of(dateColumnName, receiptColumnName, amountColumnName); + List types = ImmutableList.of(DATE, VARCHAR, amountColumnType); + + try (OrcWriter writer = new OrcWriter( + OutputStreamOrcDataSink.create(outputFile), + columnNames, + types, + toOrcType(tableSchema), + NONE, + new OrcWriterOptions(), + ImmutableMap.of(), + true, + OrcWriteValidation.OrcWriteValidationMode.BOTH, + new OrcWriterStats())) { + BlockBuilder dateBuilder = DATE.createBlockBuilder(null, 1); + DATE.writeLong(dateBuilder, dateColumnValue); + BlockBuilder receiptBuilder = VARCHAR.createBlockBuilder(null, 1); + VARCHAR.writeString(receiptBuilder, receiptColumnValue); + BlockBuilder amountBuilder = amountColumnType.createBlockBuilder(null, 1); + writeShortDecimal(amountBuilder, amountColumnValue.unscaledValue().longValueExact()); + writer.write(new Page(dateBuilder.build(), receiptBuilder.build(), amountBuilder.build())); + } + + IcebergSplit split = new IcebergSplit( + inputFile.toString(), + 0, + inputFile.length(), + inputFile.length(), + -1, // invalid; normally known + ORC, + PartitionSpecParser.toJson(partitionSpec), + PartitionData.toJson(new PartitionData(new Object[] {dateColumnValue})), + ImmutableList.of(), + SplitWeight.standard(), + TupleDomain.all(), + ImmutableMap.of(), + 0); + + String tablePath = inputFile.location().fileName(); + TableHandle tableHandle = new TableHandle( + TEST_CATALOG_HANDLE, + new IcebergTableHandle( + CatalogHandle.fromId("iceberg:NORMAL:v12345"), + "test_schema", + tableName, + TableType.DATA, + Optional.empty(), + SchemaParser.toJson(tableSchema), + Optional.of(PartitionSpecParser.toJson(partitionSpec)), + 2, + TupleDomain.all(), + TupleDomain.all(), + OptionalLong.empty(), + ImmutableSet.of(dateColumnHandle), + Optional.empty(), + tablePath, + ImmutableMap.of(), + Optional.empty(), + false, + Optional.empty(), + ImmutableSet.of(), + Optional.of(false)), + transaction); + + // Simulate situations where the dynamic filter (e.g.: while performing a JOIN with another table) reduces considerably + // the amount of data to be processed from the current table + + TupleDomain differentDatePredicate = TupleDomain.withColumnDomains( + ImmutableMap.of( + dateColumnHandle, + Domain.singleValue(DATE, LocalDate.of(2023, 2, 2).toEpochDay()))); + TupleDomain nonOverlappingDatePredicate = TupleDomain.withColumnDomains( + ImmutableMap.of( + dateColumnHandle, + Domain.create(ValueSet.ofRanges(Range.greaterThanOrEqual(DATE, LocalDate.of(2023, 2, 2).toEpochDay())), true))); + for (TupleDomain partitionPredicate : List.of(differentDatePredicate, nonOverlappingDatePredicate)) { + try (ConnectorPageSource emptyPageSource = createTestingPageSource( + transaction, + icebergConfig, + split, + tableHandle, + ImmutableList.of(dateColumnHandle, receiptColumnHandle, amountColumnHandle), + getDynamicFilter(partitionPredicate))) { + assertThat(emptyPageSource.getNextPage()).isNull(); + } + } + + TupleDomain sameDatePredicate = TupleDomain.withColumnDomains( + ImmutableMap.of( + dateColumnHandle, + Domain.singleValue(DATE, dateColumnValue))); + TupleDomain overlappingDatePredicate = TupleDomain.withColumnDomains( + ImmutableMap.of( + dateColumnHandle, + Domain.create(ValueSet.ofRanges(Range.range(DATE, LocalDate.of(2023, 1, 1).toEpochDay(), true, LocalDate.of(2023, 2, 1).toEpochDay(), false)), true))); + for (TupleDomain partitionPredicate : List.of(sameDatePredicate, overlappingDatePredicate)) { + try (ConnectorPageSource nonEmptyPageSource = createTestingPageSource( + transaction, + icebergConfig, + split, + tableHandle, + ImmutableList.of(dateColumnHandle, receiptColumnHandle, amountColumnHandle), + getDynamicFilter(partitionPredicate))) { + Page page = nonEmptyPageSource.getNextPage(); + assertThat(page).isNotNull(); + assertThat(page.getPositionCount()).isEqualTo(1); + assertThat(page.getBlock(0).getInt(0, 0)).isEqualTo(dateColumnValue); + assertThat(page.getBlock(1).getSlice(0, 0, page.getBlock(1).getSliceLength(0)).toStringUtf8()).isEqualTo(receiptColumnValue); + assertThat(((SqlDecimal) amountColumnType.getObjectValue(null, page.getBlock(2), 0)).toBigDecimal()).isEqualTo(amountColumnValue); + } + } + } + } + + @Test + public void testDynamicSplitPruningWithExplicitPartitionFilterPartitionEvolution() + throws IOException + { + String tableName = "sales_table"; + String yearColumnName = "year"; + ColumnIdentity yearColumnIdentity = new ColumnIdentity(1, yearColumnName, PRIMITIVE, ImmutableList.of()); + IcebergColumnHandle yearColumnHandle = new IcebergColumnHandle(yearColumnIdentity, INTEGER, ImmutableList.of(), INTEGER, true, Optional.empty()); + long yearColumnValue = 2023L; + String monthColumnName = "month"; + ColumnIdentity monthColumnIdentity = new ColumnIdentity(2, monthColumnName, PRIMITIVE, ImmutableList.of()); + IcebergColumnHandle monthColumnHandle = new IcebergColumnHandle(monthColumnIdentity, INTEGER, ImmutableList.of(), INTEGER, true, Optional.empty()); + long monthColumnValue = 1L; + String receiptColumnName = "receipt"; + ColumnIdentity receiptColumnIdentity = new ColumnIdentity(3, receiptColumnName, PRIMITIVE, ImmutableList.of()); + IcebergColumnHandle receiptColumnHandle = new IcebergColumnHandle(receiptColumnIdentity, VARCHAR, ImmutableList.of(), VARCHAR, true, Optional.empty()); + String receiptColumnValue = "#12345"; + String amountColumnName = "amount"; + ColumnIdentity amountColumnIdentity = new ColumnIdentity(4, amountColumnName, PRIMITIVE, ImmutableList.of()); + DecimalType amountColumnType = DecimalType.createDecimalType(10, 2); + IcebergColumnHandle amountColumnHandle = new IcebergColumnHandle(amountColumnIdentity, amountColumnType, ImmutableList.of(), amountColumnType, true, Optional.empty()); + BigDecimal amountColumnValue = new BigDecimal("1234567.65"); + Schema tableSchema = new Schema( + optional(yearColumnIdentity.getId(), yearColumnName, Types.IntegerType.get()), + optional(monthColumnIdentity.getId(), monthColumnName, Types.IntegerType.get()), + optional(receiptColumnIdentity.getId(), receiptColumnName, Types.StringType.get()), + optional(amountColumnIdentity.getId(), amountColumnName, Types.DecimalType.of(10, 2))); + PartitionSpec partitionSpec = PartitionSpec.builderFor(tableSchema) + .identity(yearColumnName) + .build(); + IcebergConfig icebergConfig = new IcebergConfig(); + HiveTransactionHandle transaction = new HiveTransactionHandle(false); + try (TempFile file = new TempFile()) { + Files.delete(file.path()); + + TrinoOutputFile outputFile = new LocalOutputFile(file.file()); + TrinoInputFile inputFile = new LocalInputFile(file.file()); + List columnNames = ImmutableList.of(yearColumnName, monthColumnName, receiptColumnName, amountColumnName); + List types = ImmutableList.of(INTEGER, INTEGER, VARCHAR, amountColumnType); + + try (OrcWriter writer = new OrcWriter( + OutputStreamOrcDataSink.create(outputFile), + columnNames, + types, + toOrcType(tableSchema), + NONE, + new OrcWriterOptions(), + ImmutableMap.of(), + true, + OrcWriteValidation.OrcWriteValidationMode.BOTH, + new OrcWriterStats())) { + BlockBuilder yearBuilder = INTEGER.createBlockBuilder(null, 1); + INTEGER.writeLong(yearBuilder, yearColumnValue); + BlockBuilder monthBuilder = INTEGER.createBlockBuilder(null, 1); + INTEGER.writeLong(monthBuilder, monthColumnValue); + BlockBuilder receiptBuilder = VARCHAR.createBlockBuilder(null, 1); + VARCHAR.writeString(receiptBuilder, receiptColumnValue); + BlockBuilder amountBuilder = amountColumnType.createBlockBuilder(null, 1); + writeShortDecimal(amountBuilder, amountColumnValue.unscaledValue().longValueExact()); + writer.write(new Page(yearBuilder.build(), monthBuilder.build(), receiptBuilder.build(), amountBuilder.build())); + } + + IcebergSplit split = new IcebergSplit( + inputFile.toString(), + 0, + inputFile.length(), + inputFile.length(), + -1, // invalid; normally known + ORC, + PartitionSpecParser.toJson(partitionSpec), + PartitionData.toJson(new PartitionData(new Object[] {yearColumnValue})), + ImmutableList.of(), + SplitWeight.standard(), + TupleDomain.all(), + ImmutableMap.of(), + 0); + + String tablePath = inputFile.location().fileName(); + // Simulate the situation where `month` column is added at a later phase as partitioning column + // in addition to the `year` column, which leads to use it as unenforced predicate in the table handle + // after applying the filter + TableHandle tableHandle = new TableHandle( + TEST_CATALOG_HANDLE, + new IcebergTableHandle( + CatalogHandle.fromId("iceberg:NORMAL:v12345"), + "test_schema", + tableName, + TableType.DATA, + Optional.empty(), + SchemaParser.toJson(tableSchema), + Optional.of(PartitionSpecParser.toJson(partitionSpec)), + 2, + TupleDomain.withColumnDomains( + ImmutableMap.of( + yearColumnHandle, + Domain.create(ValueSet.ofRanges(Range.range(INTEGER, 2023L, true, 2024L, true)), true))), + TupleDomain.withColumnDomains( + ImmutableMap.of( + monthColumnHandle, + Domain.create(ValueSet.ofRanges(Range.range(INTEGER, 1L, true, 12L, true)), true))), + OptionalLong.empty(), + ImmutableSet.of(yearColumnHandle, monthColumnHandle, receiptColumnHandle, amountColumnHandle), + Optional.empty(), + tablePath, + ImmutableMap.of(), + Optional.empty(), + false, + Optional.empty(), + ImmutableSet.of(), + Optional.of(false)), + transaction); + + // Simulate situations where the dynamic filter (e.g.: while performing a JOIN with another table) reduces considerably + // the amount of data to be processed from the current table + TupleDomain differentYearPredicate = TupleDomain.withColumnDomains( + ImmutableMap.of( + yearColumnHandle, + Domain.singleValue(INTEGER, 2024L))); + TupleDomain sameYearAndDifferentMonthPredicate = TupleDomain.withColumnDomains( + ImmutableMap.of( + yearColumnHandle, + Domain.singleValue(INTEGER, 2023L), + monthColumnHandle, + Domain.singleValue(INTEGER, 2L))); + for (TupleDomain partitionPredicate : List.of(differentYearPredicate, sameYearAndDifferentMonthPredicate)) { + try (ConnectorPageSource emptyPageSource = createTestingPageSource( + transaction, + icebergConfig, + split, + tableHandle, + ImmutableList.of(yearColumnHandle, monthColumnHandle, receiptColumnHandle, amountColumnHandle), + getDynamicFilter(partitionPredicate))) { + assertThat(emptyPageSource.getNextPage()).isNull(); + } + } + + TupleDomain sameYearPredicate = TupleDomain.withColumnDomains( + ImmutableMap.of( + yearColumnHandle, + Domain.singleValue(INTEGER, 2023L))); + TupleDomain sameYearAndMonthPredicate = TupleDomain.withColumnDomains( + ImmutableMap.of( + yearColumnHandle, + Domain.singleValue(INTEGER, 2023L), + monthColumnHandle, + Domain.singleValue(INTEGER, 1L))); + for (TupleDomain partitionPredicate : List.of(sameYearPredicate, sameYearAndMonthPredicate)) { + try (ConnectorPageSource nonEmptyPageSource = createTestingPageSource( + transaction, + icebergConfig, + split, + tableHandle, + ImmutableList.of(yearColumnHandle, monthColumnHandle, receiptColumnHandle, amountColumnHandle), + getDynamicFilter(partitionPredicate))) { + Page page = nonEmptyPageSource.getNextPage(); + assertThat(page).isNotNull(); + assertThat(page.getPositionCount()).isEqualTo(1); + assertThat(page.getBlock(0).getInt(0, 0)).isEqualTo(2023L); + assertThat(page.getBlock(1).getInt(0, 0)).isEqualTo(1L); + assertThat(page.getBlock(2).getSlice(0, 0, page.getBlock(2).getSliceLength(0)).toStringUtf8()).isEqualTo(receiptColumnValue); + assertThat(((SqlDecimal) amountColumnType.getObjectValue(null, page.getBlock(3), 0)).toBigDecimal()).isEqualTo(amountColumnValue); + } } } } @@ -137,7 +538,7 @@ private static void writeOrcContent(TrinoOutputFile outputFile) OutputStreamOrcDataSink.create(outputFile), columnNames, types, - TypeConverter.toOrcType(TABLE_SCHEMA), + toOrcType(TABLE_SCHEMA), NONE, new OrcWriterOptions(), ImmutableMap.of(), @@ -152,76 +553,30 @@ private static void writeOrcContent(TrinoOutputFile outputFile) } } - private static ConnectorPageSource createTestingPageSource(HiveTransactionHandle transaction, IcebergConfig icebergConfig, TrinoInputFile inputFile, DynamicFilter dynamicFilter) - throws IOException + private static ConnectorPageSource createTestingPageSource( + HiveTransactionHandle transaction, + IcebergConfig icebergConfig, + IcebergSplit split, + TableHandle tableHandle, + List columns, + DynamicFilter dynamicFilter) { - IcebergSplit split = new IcebergSplit( - inputFile.toString(), - 0, - inputFile.length(), - inputFile.length(), - ORC, - PartitionSpecParser.toJson(PartitionSpec.unpartitioned()), - PartitionData.toJson(new PartitionData(new Object[] {})), - ImmutableList.of(), - SplitWeight.standard()); - - String tablePath = inputFile.location().fileName(); - TableHandle tableHandle = new TableHandle( - TEST_CATALOG_HANDLE, - new IcebergTableHandle( - CatalogHandle.fromId("iceberg:NORMAL:v12345"), - SCHEMA_NAME, - TABLE_NAME, - TableType.DATA, - Optional.empty(), - SchemaParser.toJson(TABLE_SCHEMA), - Optional.of(PartitionSpecParser.toJson(PartitionSpec.unpartitioned())), - 2, - TupleDomain.withColumnDomains(ImmutableMap.of(KEY_ICEBERG_COLUMN_HANDLE, Domain.singleValue(INTEGER, (long) KEY_COLUMN_VALUE))), - TupleDomain.all(), - OptionalLong.empty(), - ImmutableSet.of(KEY_ICEBERG_COLUMN_HANDLE), - Optional.empty(), - tablePath, - ImmutableMap.of(), - false, - Optional.empty()), - transaction); - FileFormatDataSourceStats stats = new FileFormatDataSourceStats(); - IcebergPageSourceProvider provider = new IcebergPageSourceProvider( - new HdfsFileSystemFactory(HDFS_ENVIRONMENT, HDFS_FILE_SYSTEM_STATS), + IcebergPageSourceProviderFactory factory = new IcebergPageSourceProviderFactory( + new DefaultIcebergFileSystemFactory(new HdfsFileSystemFactory(HDFS_ENVIRONMENT, HDFS_FILE_SYSTEM_STATS)), stats, ORC_READER_CONFIG, PARQUET_READER_CONFIG, TESTING_TYPE_MANAGER); - - return provider.createPageSource( + return factory.createPageSourceProvider().createPageSource( transaction, getSession(icebergConfig), split, tableHandle.getConnectorHandle(), - ImmutableList.of(KEY_ICEBERG_COLUMN_HANDLE, DATA_ICEBERG_COLUMN_HANDLE), + columns, dynamicFilter); } - private static TupleDomain getTupleDomainForSplitPruning() - { - return TupleDomain.withColumnDomains( - ImmutableMap.of( - KEY_ICEBERG_COLUMN_HANDLE, - Domain.singleValue(INTEGER, 1L))); - } - - private static TupleDomain getNonSelectiveTupleDomain() - { - return TupleDomain.withColumnDomains( - ImmutableMap.of( - KEY_ICEBERG_COLUMN_HANDLE, - Domain.singleValue(INTEGER, (long) KEY_COLUMN_VALUE))); - } - private static TestingConnectorSession getSession(IcebergConfig icebergConfig) { return TestingConnectorSession.builder() diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergOrcMetricsCollection.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergOrcMetricsCollection.java index 0432d970640c..e48a53f87709 100644 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergOrcMetricsCollection.java +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergOrcMetricsCollection.java @@ -43,12 +43,12 @@ import java.util.Map; import java.util.Optional; -import static com.google.inject.util.Modules.EMPTY_MODULE; +import static com.google.common.util.concurrent.MoreExecutors.directExecutor; import static io.trino.SystemSessionProperties.MAX_DRIVERS_PER_TASK; import static io.trino.SystemSessionProperties.TASK_CONCURRENCY; import static io.trino.SystemSessionProperties.TASK_PARTITIONED_WRITER_COUNT; import static io.trino.SystemSessionProperties.TASK_WRITER_COUNT; -import static io.trino.plugin.hive.metastore.cache.CachingHiveMetastore.memoizeMetastore; +import static io.trino.plugin.hive.metastore.cache.CachingHiveMetastore.createPerTransactionCache; import static io.trino.plugin.hive.metastore.file.TestingFileHiveMetastore.createTestingFileHiveMetastore; import static io.trino.plugin.iceberg.DataFileRecord.toDataFileRecord; import static io.trino.plugin.iceberg.IcebergQueryRunner.ICEBERG_CATALOG; @@ -83,12 +83,12 @@ protected QueryRunner createQueryRunner() File baseDir = queryRunner.getCoordinator().getBaseDataDir().resolve("iceberg_data").toFile(); HiveMetastore metastore = createTestingFileHiveMetastore(baseDir); - queryRunner.installPlugin(new TestingIcebergPlugin(Optional.of(new TestingIcebergFileMetastoreCatalogModule(metastore)), Optional.empty(), EMPTY_MODULE)); + queryRunner.installPlugin(new TestingIcebergPlugin(baseDir.toPath(), Optional.of(new TestingIcebergFileMetastoreCatalogModule(metastore)))); queryRunner.createCatalog(ICEBERG_CATALOG, "iceberg", ImmutableMap.of("iceberg.file-format", "ORC")); TrinoFileSystemFactory fileSystemFactory = getFileSystemFactory(queryRunner); tableOperationsProvider = new FileMetastoreTableOperationsProvider(fileSystemFactory); - CachingHiveMetastore cachingHiveMetastore = memoizeMetastore(metastore, 1000); + CachingHiveMetastore cachingHiveMetastore = createPerTransactionCache(metastore, 1000); trinoCatalog = new TrinoHiveCatalog( new CatalogName("catalog"), cachingHiveMetastore, @@ -98,7 +98,9 @@ protected QueryRunner createQueryRunner() tableOperationsProvider, false, false, - false); + false, + new IcebergConfig().isHideMaterializedViewStorageTable(), + directExecutor()); queryRunner.installPlugin(new TpchPlugin()); queryRunner.createCatalog("tpch", "tpch"); diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergParquetConnectorTest.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergParquetConnectorTest.java index 16633c9f091f..312ff45aa1fc 100644 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergParquetConnectorTest.java +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergParquetConnectorTest.java @@ -13,6 +13,7 @@ */ package io.trino.plugin.iceberg; +import io.trino.filesystem.Location; import io.trino.testing.MaterializedResult; import io.trino.testing.sql.TestTable; import org.testng.annotations.Test; @@ -24,7 +25,6 @@ import static io.trino.plugin.iceberg.IcebergFileFormat.PARQUET; import static io.trino.plugin.iceberg.IcebergTestUtils.checkParquetFileSorting; import static io.trino.plugin.iceberg.IcebergTestUtils.withSmallRowGroups; -import static org.assertj.core.api.Assertions.assertThatThrownBy; import static org.testng.Assert.assertEquals; public class TestIcebergParquetConnectorTest @@ -45,7 +45,9 @@ protected boolean supportsIcebergFileStatistics(String typeName) protected boolean supportsRowGroupStatistics(String typeName) { return !(typeName.equalsIgnoreCase("varbinary") || + typeName.equalsIgnoreCase("time") || typeName.equalsIgnoreCase("time(6)") || + typeName.equalsIgnoreCase("timestamp(3) with time zone") || typeName.equalsIgnoreCase("timestamp(6) with time zone")); } @@ -84,18 +86,9 @@ protected Optional filterSetColumnTypesDataProvider(SetColum return super.filterSetColumnTypesDataProvider(setup); } - @Override - public void testDropAmbiguousRowFieldCaseSensitivity() - { - // TODO https://github.com/trinodb/trino/issues/16273 The connector can't read row types having ambiguous field names in Parquet files. e.g. row(X int, x int) - assertThatThrownBy(super::testDropAmbiguousRowFieldCaseSensitivity) - .hasMessageContaining("Error opening Iceberg split") - .hasStackTraceContaining("Multiple entries with same key"); - } - @Override protected boolean isFileSorted(String path, String sortColumnName) { - return checkParquetFileSorting(path, sortColumnName); + return checkParquetFileSorting(fileSystem.newInputFile(Location.of(path)), sortColumnName); } } diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergParquetWithBloomFilters.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergParquetWithBloomFilters.java index 54b3c936b1d1..8a2815303446 100644 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergParquetWithBloomFilters.java +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergParquetWithBloomFilters.java @@ -21,6 +21,7 @@ import io.trino.testing.DistributedQueryRunner; import io.trino.testing.QueryRunner; +import java.nio.file.Files; import java.nio.file.Path; import java.util.List; @@ -39,7 +40,7 @@ protected QueryRunner createQueryRunner() dataDirectory = queryRunner.getCoordinator().getBaseDataDir().resolve("iceberg_data"); // create hive catalog - queryRunner.installPlugin(new TestingHivePlugin()); + queryRunner.installPlugin(new TestingHivePlugin(Files.createTempDirectory(null))); queryRunner.createCatalog("hive", "hive", ImmutableMap.builder() .put("hive.metastore", "file") .put("hive.metastore.catalog.dir", dataDirectory.toString()) diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergPlugin.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergPlugin.java index 8395130d2174..036664c3b3c7 100644 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergPlugin.java +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergPlugin.java @@ -256,6 +256,25 @@ public void testRestCatalog() .shutdown(); } + @Test + public void testRestCatalogValidations() + { + ConnectorFactory factory = getConnectorFactory(); + + assertThatThrownBy(() -> factory.create( + "test", + Map.of( + "iceberg.catalog.type", "rest", + "iceberg.register-table-procedure.enabled", "true", + "iceberg.rest-catalog.uri", "https://foo:1234", + "iceberg.rest-catalog.vended-credentials-enabled", "true", + "bootstrap.quiet", "true"), + new TestingConnectorContext()) + .shutdown()) + .isInstanceOf(ApplicationConfigurationException.class) + .hasMessageContaining("Using the `register_table` procedure with vended credentials is currently not supported"); + } + @Test public void testJdbcCatalog() { diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergProjectionPushdownPlans.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergProjectionPushdownPlans.java index 672ae53495eb..c9c2f1f9cff5 100644 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergProjectionPushdownPlans.java +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergProjectionPushdownPlans.java @@ -43,7 +43,6 @@ import static com.google.common.base.Predicates.equalTo; import static com.google.common.io.MoreFiles.deleteRecursively; import static com.google.common.io.RecursiveDeleteOption.ALLOW_INSECURE; -import static com.google.inject.util.Modules.EMPTY_MODULE; import static io.trino.plugin.hive.metastore.file.TestingFileHiveMetastore.createTestingFileHiveMetastore; import static io.trino.spi.type.BigintType.BIGINT; import static io.trino.sql.planner.assertions.PlanMatchPattern.any; @@ -89,7 +88,7 @@ protected LocalQueryRunner createLocalQueryRunner() queryRunner.createCatalog( CATALOG, - new TestingIcebergConnectorFactory(Optional.of(new TestingIcebergFileMetastoreCatalogModule(metastore)), Optional.empty(), EMPTY_MODULE), + new TestingIcebergConnectorFactory(metastoreDir.toPath(), Optional.of(new TestingIcebergFileMetastoreCatalogModule(metastore))), ImmutableMap.of()); Database database = Database.builder() @@ -159,12 +158,14 @@ public void testDereferencePushdown() column0Handle.getType(), ImmutableList.of(column0Handle.getColumnIdentity().getChildren().get(0).getId()), BIGINT, + false, Optional.empty()); IcebergColumnHandle columnY = new IcebergColumnHandle( column0Handle.getColumnIdentity(), column0Handle.getType(), ImmutableList.of(column0Handle.getColumnIdentity().getChildren().get(1).getId()), BIGINT, + false, Optional.empty()); // Simple Projection pushdown diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergSplitSource.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergSplitSource.java index af02fe6ec1ef..7d3db069909a 100644 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergSplitSource.java +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergSplitSource.java @@ -18,15 +18,24 @@ import com.google.common.collect.ImmutableSet; import io.airlift.units.Duration; import io.trino.filesystem.TrinoFileSystemFactory; +import io.trino.filesystem.cache.DefaultCachingHostAddressProvider; import io.trino.plugin.base.CatalogName; import io.trino.plugin.hive.TrinoViewHiveMetastore; import io.trino.plugin.hive.metastore.HiveMetastore; import io.trino.plugin.hive.metastore.cache.CachingHiveMetastore; +import io.trino.plugin.hive.orc.OrcReaderConfig; +import io.trino.plugin.hive.orc.OrcWriterConfig; +import io.trino.plugin.hive.parquet.ParquetReaderConfig; +import io.trino.plugin.hive.parquet.ParquetWriterConfig; import io.trino.plugin.iceberg.catalog.TrinoCatalog; import io.trino.plugin.iceberg.catalog.file.FileMetastoreTableOperationsProvider; import io.trino.plugin.iceberg.catalog.hms.TrinoHiveCatalog; +import io.trino.plugin.iceberg.catalog.rest.DefaultIcebergFileSystemFactory; +import io.trino.plugin.iceberg.fileio.ForwardingFileIo; +import io.trino.spi.SplitWeight; import io.trino.spi.connector.CatalogHandle; import io.trino.spi.connector.ColumnHandle; +import io.trino.spi.connector.ConnectorSession; import io.trino.spi.connector.DynamicFilter; import io.trino.spi.connector.SchemaTableName; import io.trino.spi.predicate.Domain; @@ -38,43 +47,66 @@ import io.trino.testing.AbstractTestQueryFramework; import io.trino.testing.DistributedQueryRunner; import io.trino.testing.QueryRunner; +import io.trino.testing.TestingConnectorSession; +import org.apache.iceberg.PartitionSpec; import org.apache.iceberg.PartitionSpecParser; import org.apache.iceberg.SchemaParser; import org.apache.iceberg.Table; +import org.apache.iceberg.TableProperties; +import org.apache.iceberg.data.GenericRecord; +import org.apache.iceberg.data.Record; +import org.apache.iceberg.data.parquet.GenericParquetWriter; +import org.apache.iceberg.deletes.PositionDelete; +import org.apache.iceberg.deletes.PositionDeleteWriter; +import org.apache.iceberg.io.FileIO; +import org.apache.iceberg.parquet.Parquet; import org.apache.iceberg.types.Conversions; import org.apache.iceberg.types.Type; import org.apache.iceberg.types.Types; import org.testng.annotations.AfterClass; import org.testng.annotations.Test; +import java.io.Closeable; import java.io.File; import java.io.IOException; import java.nio.ByteBuffer; import java.nio.file.Files; +import java.util.List; import java.util.Map; import java.util.Optional; import java.util.OptionalLong; import java.util.Set; +import java.util.UUID; import java.util.concurrent.CompletableFuture; -import java.util.concurrent.TimeUnit; import static com.google.common.io.MoreFiles.deleteRecursively; import static com.google.common.io.RecursiveDeleteOption.ALLOW_INSECURE; -import static io.trino.plugin.hive.metastore.cache.CachingHiveMetastore.memoizeMetastore; +import static com.google.common.util.concurrent.MoreExecutors.directExecutor; +import static com.google.common.util.concurrent.MoreExecutors.newDirectExecutorService; +import static io.trino.plugin.hive.metastore.cache.CachingHiveMetastore.createPerTransactionCache; import static io.trino.plugin.hive.metastore.file.TestingFileHiveMetastore.createTestingFileHiveMetastore; +import static io.trino.plugin.iceberg.IcebergSplitSource.createFileStatisticsDomain; import static io.trino.plugin.iceberg.IcebergTestUtils.getFileSystemFactory; +import static io.trino.plugin.iceberg.util.EqualityDeleteUtils.writeEqualityDeleteForTable; import static io.trino.spi.connector.Constraint.alwaysTrue; import static io.trino.spi.type.BigintType.BIGINT; -import static io.trino.testing.TestingConnectorSession.SESSION; import static io.trino.tpch.TpchTable.NATION; import static java.util.concurrent.TimeUnit.SECONDS; import static org.assertj.core.api.Assertions.assertThat; -import static org.testng.Assert.assertFalse; -import static org.testng.Assert.assertTrue; public class TestIcebergSplitSource extends AbstractTestQueryFramework { + private static final ConnectorSession SESSION = TestingConnectorSession.builder() + .setPropertyMetadata(new IcebergSessionProperties( + new IcebergConfig(), + new OrcReaderConfig(), + new OrcWriterConfig(), + new ParquetReaderConfig(), + new ParquetWriterConfig()) + .getSessionProperties()) + .build(); + private File metastoreDir; private TrinoFileSystemFactory fileSystemFactory; private TrinoCatalog catalog; @@ -93,7 +125,7 @@ protected QueryRunner createQueryRunner() .build(); this.fileSystemFactory = getFileSystemFactory(queryRunner); - CachingHiveMetastore cachingHiveMetastore = memoizeMetastore(metastore, 1000); + CachingHiveMetastore cachingHiveMetastore = createPerTransactionCache(metastore, 1000); this.catalog = new TrinoHiveCatalog( new CatalogName("hive"), cachingHiveMetastore, @@ -103,7 +135,9 @@ protected QueryRunner createQueryRunner() new FileMetastoreTableOperationsProvider(fileSystemFactory), false, false, - false); + false, + new IcebergConfig().isHideMaterializedViewStorageTable(), + directExecutor()); return queryRunner; } @@ -122,29 +156,14 @@ public void testIncompleteDynamicFilterTimeout() long startMillis = System.currentTimeMillis(); SchemaTableName schemaTableName = new SchemaTableName("tpch", "nation"); Table nationTable = catalog.loadTable(SESSION, schemaTableName); - IcebergTableHandle tableHandle = new IcebergTableHandle( - CatalogHandle.fromId("iceberg:NORMAL:v12345"), - schemaTableName.getSchemaName(), - schemaTableName.getTableName(), - TableType.DATA, - Optional.empty(), - SchemaParser.toJson(nationTable.schema()), - Optional.of(PartitionSpecParser.toJson(nationTable.spec())), - 1, - TupleDomain.all(), - TupleDomain.all(), - OptionalLong.empty(), - ImmutableSet.of(), - Optional.empty(), - nationTable.location(), - nationTable.properties(), - false, - Optional.empty()); + IcebergTableHandle tableHandle = createTableHandle(schemaTableName, nationTable, TupleDomain.all()); + CompletableFuture isBlocked = new CompletableFuture<>(); try (IcebergSplitSource splitSource = new IcebergSplitSource( - fileSystemFactory, + new DefaultIcebergFileSystemFactory(fileSystemFactory), SESSION, tableHandle, + nationTable, nationTable.newScan(), Optional.empty(), new DynamicFilter() @@ -158,14 +177,7 @@ public Set getColumnsCovered() @Override public CompletableFuture isBlocked() { - return CompletableFuture.runAsync(() -> { - try { - TimeUnit.HOURS.sleep(1); - } - catch (InterruptedException e) { - throw new IllegalStateException(e); - } - }); + return isBlocked; } @Override @@ -190,7 +202,9 @@ public TupleDomain getCurrentPredicate() alwaysTrue(), new TestingTypeManager(), false, - new IcebergConfig().getMinimumAssignedSplitWeight())) { + new IcebergConfig().getMinimumAssignedSplitWeight(), + new DefaultCachingHostAddressProvider(), + newDirectExecutorService())) { ImmutableList.Builder splits = ImmutableList.builder(); while (!splitSource.isFinished()) { splitSource.getNextBatch(100).get() @@ -200,11 +214,14 @@ public TupleDomain getCurrentPredicate() .forEach(splits::add); } assertThat(splits.build().size()).isGreaterThan(0); - assertTrue(splitSource.isFinished()); + assertThat(splitSource.isFinished()).isTrue(); assertThat(System.currentTimeMillis() - startMillis) .as("IcebergSplitSource failed to wait for dynamicFilteringWaitTimeout") .isGreaterThanOrEqualTo(2000); } + finally { + isBlocked.complete(null); + } } @Test @@ -215,19 +232,20 @@ public void testBigintPartitionPruning() BIGINT, ImmutableList.of(), BIGINT, + true, Optional.empty()); - assertFalse(IcebergSplitSource.partitionMatchesPredicate( + assertThat(IcebergSplitSource.partitionMatchesPredicate( ImmutableSet.of(bigintColumn), () -> ImmutableMap.of(bigintColumn, NullableValue.of(BIGINT, 1000L)), - TupleDomain.fromFixedValues(ImmutableMap.of(bigintColumn, NullableValue.of(BIGINT, 100L))))); - assertTrue(IcebergSplitSource.partitionMatchesPredicate( + TupleDomain.fromFixedValues(ImmutableMap.of(bigintColumn, NullableValue.of(BIGINT, 100L))))).isFalse(); + assertThat(IcebergSplitSource.partitionMatchesPredicate( ImmutableSet.of(bigintColumn), () -> ImmutableMap.of(bigintColumn, NullableValue.of(BIGINT, 1000L)), - TupleDomain.fromFixedValues(ImmutableMap.of(bigintColumn, NullableValue.of(BIGINT, 1000L))))); - assertFalse(IcebergSplitSource.partitionMatchesPredicate( + TupleDomain.fromFixedValues(ImmutableMap.of(bigintColumn, NullableValue.of(BIGINT, 1000L))))).isTrue(); + assertThat(IcebergSplitSource.partitionMatchesPredicate( ImmutableSet.of(bigintColumn), () -> ImmutableMap.of(bigintColumn, NullableValue.of(BIGINT, 1000L)), - TupleDomain.fromFixedValues(ImmutableMap.of(bigintColumn, NullableValue.asNull(BIGINT))))); + TupleDomain.fromFixedValues(ImmutableMap.of(bigintColumn, NullableValue.asNull(BIGINT))))).isFalse(); } @Test @@ -238,97 +256,18 @@ public void testBigintStatisticsPruning() BIGINT, ImmutableList.of(), BIGINT, + true, Optional.empty()); + Map primitiveTypes = ImmutableMap.of(1, Types.LongType.get()); Map lowerBound = ImmutableMap.of(1, Conversions.toByteBuffer(Types.LongType.get(), 1000L)); Map upperBound = ImmutableMap.of(1, Conversions.toByteBuffer(Types.LongType.get(), 2000L)); + TupleDomain domainLowerUpperBound = TupleDomain.withColumnDomains( + ImmutableMap.of(bigintColumn, Domain.create(ValueSet.ofRanges(Range.range(BIGINT, 1000L, true, 2000L, true)), false))); + List predicatedColumns = ImmutableList.of(bigintColumn); - assertFalse(IcebergSplitSource.fileMatchesPredicate( - primitiveTypes, - TupleDomain.fromFixedValues(ImmutableMap.of(bigintColumn, NullableValue.of(BIGINT, 0L))), - lowerBound, - upperBound, - ImmutableMap.of(1, 0L))); - assertTrue(IcebergSplitSource.fileMatchesPredicate( - primitiveTypes, - TupleDomain.fromFixedValues(ImmutableMap.of(bigintColumn, NullableValue.of(BIGINT, 1000L))), - lowerBound, - upperBound, - ImmutableMap.of(1, 0L))); - assertTrue(IcebergSplitSource.fileMatchesPredicate( - primitiveTypes, - TupleDomain.fromFixedValues(ImmutableMap.of(bigintColumn, NullableValue.of(BIGINT, 1500L))), - lowerBound, - upperBound, - ImmutableMap.of(1, 0L))); - assertTrue(IcebergSplitSource.fileMatchesPredicate( - primitiveTypes, - TupleDomain.fromFixedValues(ImmutableMap.of(bigintColumn, NullableValue.of(BIGINT, 2000L))), - lowerBound, - upperBound, - ImmutableMap.of(1, 0L))); - assertFalse(IcebergSplitSource.fileMatchesPredicate( - primitiveTypes, - TupleDomain.fromFixedValues(ImmutableMap.of(bigintColumn, NullableValue.of(BIGINT, 3000L))), - lowerBound, - upperBound, - ImmutableMap.of(1, 0L))); - - Domain outsideStatisticsRangeAllowNulls = Domain.create(ValueSet.ofRanges(Range.range(BIGINT, 0L, true, 100L, true)), true); - assertFalse(IcebergSplitSource.fileMatchesPredicate( - primitiveTypes, - TupleDomain.withColumnDomains(ImmutableMap.of(bigintColumn, outsideStatisticsRangeAllowNulls)), - lowerBound, - upperBound, - ImmutableMap.of(1, 0L))); - assertTrue(IcebergSplitSource.fileMatchesPredicate( - primitiveTypes, - TupleDomain.withColumnDomains(ImmutableMap.of(bigintColumn, outsideStatisticsRangeAllowNulls)), - lowerBound, - upperBound, - ImmutableMap.of(1, 1L))); - - Domain outsideStatisticsRangeNoNulls = Domain.create(ValueSet.ofRanges(Range.range(BIGINT, 0L, true, 100L, true)), false); - assertFalse(IcebergSplitSource.fileMatchesPredicate( - primitiveTypes, - TupleDomain.withColumnDomains(ImmutableMap.of(bigintColumn, outsideStatisticsRangeNoNulls)), - lowerBound, - upperBound, - ImmutableMap.of(1, 0L))); - assertFalse(IcebergSplitSource.fileMatchesPredicate( - primitiveTypes, - TupleDomain.withColumnDomains(ImmutableMap.of(bigintColumn, outsideStatisticsRangeNoNulls)), - lowerBound, - upperBound, - ImmutableMap.of(1, 1L))); - - Domain insideStatisticsRange = Domain.create(ValueSet.ofRanges(Range.range(BIGINT, 1001L, true, 1002L, true)), false); - assertTrue(IcebergSplitSource.fileMatchesPredicate( - primitiveTypes, - TupleDomain.withColumnDomains(ImmutableMap.of(bigintColumn, insideStatisticsRange)), - lowerBound, - upperBound, - ImmutableMap.of(1, 0L))); - assertTrue(IcebergSplitSource.fileMatchesPredicate( - primitiveTypes, - TupleDomain.withColumnDomains(ImmutableMap.of(bigintColumn, insideStatisticsRange)), - lowerBound, - upperBound, - ImmutableMap.of(1, 1L))); - - Domain overlappingStatisticsRange = Domain.create(ValueSet.ofRanges(Range.range(BIGINT, 990L, true, 1010L, true)), false); - assertTrue(IcebergSplitSource.fileMatchesPredicate( - primitiveTypes, - TupleDomain.withColumnDomains(ImmutableMap.of(bigintColumn, overlappingStatisticsRange)), - lowerBound, - upperBound, - ImmutableMap.of(1, 0L))); - assertTrue(IcebergSplitSource.fileMatchesPredicate( - primitiveTypes, - TupleDomain.withColumnDomains(ImmutableMap.of(bigintColumn, overlappingStatisticsRange)), - lowerBound, - upperBound, - ImmutableMap.of(1, 1L))); + assertThat(createFileStatisticsDomain(primitiveTypes, lowerBound, upperBound, ImmutableMap.of(1, 0L), predicatedColumns)) + .isEqualTo(domainLowerUpperBound); } @Test @@ -339,50 +278,144 @@ public void testNullStatisticsMaps() BIGINT, ImmutableList.of(), BIGINT, + true, Optional.empty()); Map primitiveTypes = ImmutableMap.of(1, Types.LongType.get()); Map lowerBound = ImmutableMap.of(1, Conversions.toByteBuffer(Types.LongType.get(), -1000L)); Map upperBound = ImmutableMap.of(1, Conversions.toByteBuffer(Types.LongType.get(), 2000L)); - TupleDomain domainOfZero = TupleDomain.fromFixedValues(ImmutableMap.of(bigintColumn, NullableValue.of(BIGINT, 0L))); - - assertTrue(IcebergSplitSource.fileMatchesPredicate( - primitiveTypes, - domainOfZero, - null, - upperBound, - ImmutableMap.of(1, 0L))); - assertTrue(IcebergSplitSource.fileMatchesPredicate( - primitiveTypes, - domainOfZero, - ImmutableMap.of(), - upperBound, - ImmutableMap.of(1, 0L))); - - assertTrue(IcebergSplitSource.fileMatchesPredicate( - primitiveTypes, - domainOfZero, - lowerBound, - null, - ImmutableMap.of(1, 0L))); - assertTrue(IcebergSplitSource.fileMatchesPredicate( - primitiveTypes, - domainOfZero, - lowerBound, - ImmutableMap.of(), - ImmutableMap.of(1, 0L))); - - TupleDomain onlyNull = TupleDomain.withColumnDomains(ImmutableMap.of(bigintColumn, Domain.onlyNull(BIGINT))); - assertTrue(IcebergSplitSource.fileMatchesPredicate( - primitiveTypes, - onlyNull, - ImmutableMap.of(), - ImmutableMap.of(), - null)); - assertTrue(IcebergSplitSource.fileMatchesPredicate( - primitiveTypes, - onlyNull, - ImmutableMap.of(), - ImmutableMap.of(), - ImmutableMap.of())); + TupleDomain domainLessThanUpperBound = TupleDomain.withColumnDomains( + ImmutableMap.of(bigintColumn, Domain.create(ValueSet.ofRanges(Range.lessThanOrEqual(BIGINT, 2000L)), false))); + List predicatedColumns = ImmutableList.of(bigintColumn); + + assertThat(createFileStatisticsDomain(primitiveTypes, null, upperBound, ImmutableMap.of(1, 0L), predicatedColumns)) + .isEqualTo(domainLessThanUpperBound); + assertThat(createFileStatisticsDomain(primitiveTypes, ImmutableMap.of(), upperBound, ImmutableMap.of(1, 0L), predicatedColumns)) + .isEqualTo(domainLessThanUpperBound); + + TupleDomain domainGreaterThanLessBound = TupleDomain.withColumnDomains( + ImmutableMap.of(bigintColumn, Domain.create(ValueSet.ofRanges(Range.greaterThanOrEqual(BIGINT, -1000L)), false))); + assertThat(createFileStatisticsDomain(primitiveTypes, lowerBound, null, ImmutableMap.of(1, 0L), predicatedColumns)) + .isEqualTo(domainGreaterThanLessBound); + assertThat(createFileStatisticsDomain(primitiveTypes, lowerBound, ImmutableMap.of(), ImmutableMap.of(1, 0L), predicatedColumns)) + .isEqualTo(domainGreaterThanLessBound); + + assertThat(createFileStatisticsDomain(primitiveTypes, ImmutableMap.of(), ImmutableMap.of(), null, predicatedColumns)) + .isEqualTo(TupleDomain.all()); + assertThat(createFileStatisticsDomain(primitiveTypes, ImmutableMap.of(), ImmutableMap.of(), ImmutableMap.of(), predicatedColumns)) + .isEqualTo(TupleDomain.all()); + assertThat(createFileStatisticsDomain(primitiveTypes, ImmutableMap.of(), ImmutableMap.of(), ImmutableMap.of(1, 1L), predicatedColumns)) + .isEqualTo(TupleDomain.all()); + + assertThat(createFileStatisticsDomain(primitiveTypes, ImmutableMap.of(), ImmutableMap.of(), ImmutableMap.of(1, 0L), predicatedColumns)) + .isEqualTo(TupleDomain.withColumnDomains(ImmutableMap.of(bigintColumn, Domain.notNull(BIGINT)))); + } + + @Test + public void testSplitWeight() + throws Exception + { + SchemaTableName schemaTableName = new SchemaTableName("tpch", "nation"); + Table nationTable = catalog.loadTable(SESSION, schemaTableName); + // Decrease target split size so that changes in split weight are significant enough to be detected + nationTable.updateProperties() + .set(TableProperties.SPLIT_SIZE, "10000") + .commit(); + IcebergTableHandle tableHandle = createTableHandle(schemaTableName, nationTable, TupleDomain.all()); + + IcebergSplit split = generateSplit(nationTable, tableHandle, DynamicFilter.EMPTY); + SplitWeight weightWithoutDelete = split.getSplitWeight(); + + String dataFilePath = (String) computeActual("SELECT file_path FROM \"" + schemaTableName.getTableName() + "$files\" LIMIT 1").getOnlyValue(); + + // Write position delete file + FileIO fileIo = new ForwardingFileIo(fileSystemFactory.create(SESSION)); + PositionDeleteWriter writer = Parquet.writeDeletes(fileIo.newOutputFile("local:///delete_file_" + UUID.randomUUID())) + .createWriterFunc(GenericParquetWriter::buildWriter) + .forTable(nationTable) + .overwrite() + .rowSchema(nationTable.schema()) + .withSpec(PartitionSpec.unpartitioned()) + .buildPositionWriter(); + PositionDelete positionDelete = PositionDelete.create(); + PositionDelete record = positionDelete.set(dataFilePath, 0, GenericRecord.create(nationTable.schema())); + try (Closeable ignored = writer) { + writer.write(record); + } + nationTable.newRowDelta().addDeletes(writer.toDeleteFile()).commit(); + + split = generateSplit(nationTable, tableHandle, DynamicFilter.EMPTY); + SplitWeight splitWeightWithPositionDelete = split.getSplitWeight(); + assertThat(splitWeightWithPositionDelete.getRawValue()).isGreaterThan(weightWithoutDelete.getRawValue()); + + // Write equality delete file + writeEqualityDeleteForTable( + nationTable, + fileSystemFactory, + Optional.of(nationTable.spec()), + Optional.of(new PartitionData(new Long[] {1L})), + ImmutableMap.of("regionkey", 1L), + Optional.empty()); + + split = generateSplit(nationTable, tableHandle, DynamicFilter.EMPTY); + assertThat(split.getSplitWeight().getRawValue()).isGreaterThan(splitWeightWithPositionDelete.getRawValue()); + } + + private IcebergSplit generateSplit(Table nationTable, IcebergTableHandle tableHandle, DynamicFilter dynamicFilter) + throws Exception + { + try (IcebergSplitSource splitSource = new IcebergSplitSource( + new DefaultIcebergFileSystemFactory(fileSystemFactory), + SESSION, + tableHandle, + nationTable, + nationTable.newScan(), + Optional.empty(), + dynamicFilter, + new Duration(0, SECONDS), + alwaysTrue(), + new TestingTypeManager(), + false, + 0, + new DefaultCachingHostAddressProvider(), + newDirectExecutorService())) { + ImmutableList.Builder builder = ImmutableList.builder(); + while (!splitSource.isFinished()) { + splitSource.getNextBatch(100).get() + .getSplits() + .stream() + .map(IcebergSplit.class::cast) + .forEach(builder::add); + } + List splits = builder.build(); + assertThat(splits).hasSize(1); + assertThat(splitSource.isFinished()).isTrue(); + + return splits.get(0); + } + } + + private static IcebergTableHandle createTableHandle(SchemaTableName schemaTableName, Table nationTable, TupleDomain unenforcedPredicate) + { + return new IcebergTableHandle( + CatalogHandle.fromId("iceberg:NORMAL:v12345"), + schemaTableName.getSchemaName(), + schemaTableName.getTableName(), + TableType.DATA, + Optional.empty(), + SchemaParser.toJson(nationTable.schema()), + Optional.of(PartitionSpecParser.toJson(nationTable.spec())), + 1, + unenforcedPredicate, + TupleDomain.all(), + OptionalLong.empty(), + ImmutableSet.of(), + Optional.empty(), + nationTable.location(), + nationTable.properties(), + Optional.empty(), + false, + Optional.empty(), + ImmutableSet.of(), + Optional.of(false)); } } diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergTableName.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergTableName.java index 54293b89ea14..b88dc1f118d7 100644 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergTableName.java +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergTableName.java @@ -17,9 +17,8 @@ import java.util.Optional; -import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED; -import static io.trino.testing.assertions.TrinoExceptionAssert.assertTrinoExceptionThrownBy; import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertFalse; import static org.testng.Assert.assertTrue; @@ -33,10 +32,10 @@ public void testParse() assertParseNameAndType("abc$history", "abc", TableType.HISTORY); assertParseNameAndType("abc$snapshots", "abc", TableType.SNAPSHOTS); - assertNoValidTableType("abc$data"); + assertInvalid("abc$data"); assertInvalid("abc@123", "Invalid Iceberg table name: abc@123"); assertInvalid("abc@xyz", "Invalid Iceberg table name: abc@xyz"); - assertNoValidTableType("abc$what"); + assertInvalid("abc$what"); assertInvalid("abc@123$data@456", "Invalid Iceberg table name: abc@123$data@456"); assertInvalid("abc@123$snapshots", "Invalid Iceberg table name: abc@123$snapshots"); assertInvalid("abc$snapshots@456", "Invalid Iceberg table name: abc$snapshots@456"); @@ -83,20 +82,22 @@ public void testTableNameWithType() private static void assertInvalid(String inputName, String message) { - assertTrinoExceptionThrownBy(() -> IcebergTableName.tableTypeFrom(inputName)) - .hasErrorCode(NOT_SUPPORTED) + assertThat(IcebergTableName.isIcebergTableName(inputName)).isFalse(); + + assertThatThrownBy(() -> IcebergTableName.tableTypeFrom(inputName)) + .isInstanceOf(IllegalArgumentException.class) .hasMessage(message); } - private static void assertNoValidTableType(String inputName) + private static void assertInvalid(String inputName) { - assertThat(IcebergTableName.tableTypeFrom(inputName)) - .isEmpty(); + assertInvalid(inputName, "Invalid Iceberg table name: " + inputName); } private static void assertParseNameAndType(String inputName, String tableName, TableType tableType) { assertEquals(IcebergTableName.tableNameFrom(inputName), tableName); + assertThat(IcebergTableName.tableNameFrom(inputName)).isEqualTo(tableName); assertEquals(IcebergTableName.tableTypeFrom(inputName), Optional.of(tableType)); } } diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergV2.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergV2.java index 44d840408ace..1a0f945123ab 100644 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergV2.java +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergV2.java @@ -15,20 +15,12 @@ import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableSet; import io.trino.Session; import io.trino.filesystem.TrinoFileSystemFactory; -import io.trino.plugin.base.CatalogName; -import io.trino.plugin.base.util.Closables; -import io.trino.plugin.blackhole.BlackHolePlugin; -import io.trino.plugin.hive.TrinoViewHiveMetastore; +import io.trino.plugin.hive.TestingHivePlugin; import io.trino.plugin.hive.metastore.HiveMetastore; -import io.trino.plugin.hive.metastore.cache.CachingHiveMetastore; -import io.trino.plugin.iceberg.catalog.IcebergTableOperationsProvider; -import io.trino.plugin.iceberg.catalog.TrinoCatalog; -import io.trino.plugin.iceberg.catalog.file.FileMetastoreTableOperationsProvider; -import io.trino.plugin.iceberg.catalog.hms.TrinoHiveCatalog; import io.trino.plugin.iceberg.fileio.ForwardingFileIo; -import io.trino.spi.connector.SchemaTableName; import io.trino.spi.predicate.Domain; import io.trino.spi.predicate.Range; import io.trino.spi.predicate.TupleDomain; @@ -40,7 +32,6 @@ import io.trino.testing.DistributedQueryRunner; import io.trino.testing.QueryRunner; import io.trino.testing.sql.TestTable; -import org.apache.hadoop.fs.Path; import org.apache.iceberg.BaseTable; import org.apache.iceberg.DataFile; import org.apache.iceberg.DataFiles; @@ -56,7 +47,6 @@ import org.apache.iceberg.data.GenericRecord; import org.apache.iceberg.data.Record; import org.apache.iceberg.data.parquet.GenericParquetWriter; -import org.apache.iceberg.deletes.EqualityDeleteWriter; import org.apache.iceberg.deletes.PositionDelete; import org.apache.iceberg.deletes.PositionDeleteWriter; import org.apache.iceberg.io.FileIO; @@ -66,30 +56,30 @@ import org.testng.annotations.Test; import java.io.Closeable; -import java.io.File; import java.io.IOException; import java.nio.ByteBuffer; -import java.nio.file.Files; import java.util.ArrayList; import java.util.List; import java.util.Map; import java.util.Optional; +import java.util.Set; import java.util.UUID; import java.util.concurrent.CyclicBarrier; import java.util.concurrent.ExecutorService; import java.util.concurrent.Future; +import java.util.regex.Matcher; +import java.util.regex.Pattern; import java.util.stream.Collectors; import java.util.stream.IntStream; import static com.google.common.base.Verify.verify; import static com.google.common.collect.ImmutableList.toImmutableList; +import static com.google.common.collect.ImmutableSet.toImmutableSet; import static com.google.common.collect.Iterables.getOnlyElement; -import static com.google.common.io.MoreFiles.deleteRecursively; -import static com.google.common.io.RecursiveDeleteOption.ALLOW_INSECURE; -import static io.trino.plugin.hive.metastore.cache.CachingHiveMetastore.memoizeMetastore; -import static io.trino.plugin.hive.metastore.file.TestingFileHiveMetastore.createTestingFileHiveMetastore; import static io.trino.plugin.iceberg.IcebergTestUtils.getFileSystemFactory; -import static io.trino.plugin.iceberg.IcebergUtil.loadIcebergTable; +import static io.trino.plugin.iceberg.IcebergTestUtils.getHiveMetastore; +import static io.trino.plugin.iceberg.util.EqualityDeleteUtils.writeEqualityDeleteForTable; +import static io.trino.plugin.iceberg.util.EqualityDeleteUtils.writeEqualityDeleteForTableWithSchema; import static io.trino.spi.type.IntegerType.INTEGER; import static io.trino.testing.TestingConnectorSession.SESSION; import static io.trino.testing.TestingNames.randomNameSuffix; @@ -97,8 +87,10 @@ import static java.lang.String.format; import static java.nio.ByteOrder.LITTLE_ENDIAN; import static java.nio.charset.StandardCharsets.UTF_8; +import static java.util.Map.entry; import static java.util.concurrent.Executors.newFixedThreadPool; import static java.util.concurrent.TimeUnit.SECONDS; +import static org.apache.iceberg.FileContent.EQUALITY_DELETES; import static org.apache.iceberg.FileFormat.ORC; import static org.apache.iceberg.TableProperties.SPLIT_SIZE; import static org.assertj.core.api.Assertions.assertThat; @@ -110,31 +102,23 @@ public class TestIcebergV2 extends AbstractTestQueryFramework { private HiveMetastore metastore; - private java.nio.file.Path tempDir; - private File metastoreDir; private TrinoFileSystemFactory fileSystemFactory; @Override protected QueryRunner createQueryRunner() throws Exception { - tempDir = Files.createTempDirectory("test_iceberg_v2"); - metastoreDir = tempDir.resolve("iceberg_data").toFile(); - metastore = createTestingFileHiveMetastore(metastoreDir); - DistributedQueryRunner queryRunner = IcebergQueryRunner.builder() .setInitialTables(NATION) - .setMetastoreDirectory(metastoreDir) .build(); - try { - queryRunner.installPlugin(new BlackHolePlugin()); - queryRunner.createCatalog("blackhole", "blackhole"); - } - catch (RuntimeException e) { - Closables.closeAllSuppress(e, queryRunner); - throw e; - } + metastore = getHiveMetastore(queryRunner); + fileSystemFactory = getFileSystemFactory(queryRunner); + + queryRunner.installPlugin(new TestingHivePlugin(queryRunner.getCoordinator().getBaseDataDir().resolve("iceberg_data"))); + queryRunner.createCatalog("hive", "hive", ImmutableMap.builder() + .put("hive.security", "allow-all") + .buildOrThrow()); return queryRunner; } @@ -149,7 +133,6 @@ public void initFileSystemFactory() public void tearDown() throws IOException { - deleteRecursively(tempDir, ALLOW_INSECURE); } @Test @@ -193,12 +176,9 @@ public void testV2TableWithPositionDelete() String dataFilePath = (String) computeActual("SELECT file_path FROM \"" + tableName + "$files\" LIMIT 1").getOnlyValue(); - Path metadataDir = new Path(metastoreDir.toURI()); - String deleteFileName = "delete_file_" + UUID.randomUUID(); FileIO fileIo = new ForwardingFileIo(fileSystemFactory.create(SESSION)); - Path path = new Path(metadataDir, deleteFileName); - PositionDeleteWriter writer = Parquet.writeDeletes(fileIo.newOutputFile(path.toString())) + PositionDeleteWriter writer = Parquet.writeDeletes(fileIo.newOutputFile("local:///delete_file_" + UUID.randomUUID())) .createWriterFunc(GenericParquetWriter::buildWriter) .forTable(icebergTable) .overwrite() @@ -223,10 +203,15 @@ public void testV2TableWithEqualityDelete() String tableName = "test_v2_equality_delete" + randomNameSuffix(); assertUpdate("CREATE TABLE " + tableName + " AS SELECT * FROM tpch.tiny.nation", 25); Table icebergTable = loadTable(tableName); - writeEqualityDeleteToNationTable(icebergTable, Optional.of(icebergTable.spec()), Optional.of(new PartitionData(new Long[]{1L}))); + writeEqualityDeleteToNationTable(icebergTable, Optional.of(icebergTable.spec()), Optional.of(new PartitionData(new Long[] {1L}))); assertQuery("SELECT * FROM " + tableName, "SELECT * FROM nation WHERE regionkey != 1"); // nationkey is before the equality delete column in the table schema, comment is after assertQuery("SELECT nationkey, comment FROM " + tableName, "SELECT nationkey, comment FROM nation WHERE regionkey != 1"); + + assertUpdate("INSERT INTO " + tableName + " SELECT * FROM tpch.tiny.nation", 25); + writeEqualityDeleteToNationTable(icebergTable, Optional.of(icebergTable.spec()), Optional.of(new PartitionData(new Long[] {2L})), ImmutableMap.of("regionkey", 2L)); + // the equality delete file is applied to 2 data files + assertQuery("SELECT count(*) FROM \"" + tableName + "$files\" WHERE content = " + EQUALITY_DELETES.id(), "VALUES 2"); } @Test @@ -594,7 +579,7 @@ public void testDeletingEntireFileWithMultipleSplits() long initialSnapshotId = (long) computeScalar("SELECT snapshot_id FROM \"" + tableName + "$snapshots\" ORDER BY committed_at DESC FETCH FIRST 1 ROW WITH TIES"); assertUpdate("DELETE FROM " + tableName + " WHERE regionkey < 10", 25); long parentSnapshotId = (long) computeScalar("SELECT parent_id FROM \"" + tableName + "$snapshots\" ORDER BY committed_at DESC FETCH FIRST 1 ROW WITH TIES"); - assertEquals(initialSnapshotId, parentSnapshotId); + assertThat(initialSnapshotId).isEqualTo(parentSnapshotId); assertThat(query("SELECT * FROM " + tableName)).returnsEmptyResult(); assertThat(this.loadTable(tableName).newScan().planFiles()).hasSize(1); } @@ -636,7 +621,7 @@ public void testFilesTable() throws Exception { String tableName = "test_files_table_" + randomNameSuffix(); - String tableLocation = metastoreDir.getPath() + "/" + tableName; + String tableLocation = "local:///" + tableName; assertUpdate("CREATE TABLE " + tableName + " WITH (location = '" + tableLocation + "', format_version = 2) AS SELECT * FROM tpch.tiny.nation", 25); BaseTable table = loadTable(tableName); Metrics metrics = new Metrics( @@ -657,8 +642,6 @@ public void testFilesTable() .withEncryptionKeyMetadata(ByteBuffer.wrap("Trino".getBytes(UTF_8))) .build(); table.newAppend().appendFile(dataFile).commit(); - // TODO Currently, Trino does not include equality delete files stats in the $files table. - // Once it is fixed by https://github.com/trinodb/trino/pull/16232, include equality delete output in the test. writeEqualityDeleteToNationTable(table); assertQuery( "SELECT " + @@ -680,14 +663,14 @@ public void testFilesTable() (0, 'PARQUET', 25L, - JSON '{"1":141,"2":220,"3":99,"4":807}', + JSON '{"1":137,"2":216,"3":91,"4":801}', JSON '{"1":25,"2":25,"3":25,"4":25}', jSON '{"1":0,"2":0,"3":0,"4":0}', jSON '{}', JSON '{"1":"0","2":"ALGERIA","3":"0","4":" haggle. careful"}', JSON '{"1":"24","2":"VIETNAM","3":"4","4":"y final packaget"}', null, - null, + ARRAY[4L], null), (0, 'ORC', @@ -700,7 +683,19 @@ public void testFilesTable() JSON '{"1":"4"}', X'54 72 69 6e 6f', ARRAY[4L], - null) + null), + (2, + 'PARQUET', + 1L, + JSON '{"3":49}', + JSON '{"3":1}', + JSON '{"3":0}', + JSON '{}', + JSON '{"3":"1"}', + JSON '{"3":"1"}', + null, + ARRAY[4], + ARRAY[3]) """); } @@ -714,7 +709,7 @@ public void testStatsFilePruning() Optional snapshotId = Optional.of((long) computeScalar("SELECT snapshot_id FROM \"" + testTable.getName() + "$snapshots\" ORDER BY committed_at DESC FETCH FIRST 1 ROW WITH TIES")); TypeManager typeManager = new TestingTypeManager(); Table table = loadTable(testTable.getName()); - TableStatistics withNoFilter = TableStatisticsReader.makeTableStatistics(typeManager, table, snapshotId, TupleDomain.all(), TupleDomain.all(), true); + TableStatistics withNoFilter = TableStatisticsReader.makeTableStatistics(typeManager, table, snapshotId, TupleDomain.all(), TupleDomain.all(), ImmutableSet.of(), true, fileSystemFactory.create(SESSION)); assertEquals(withNoFilter.getRowCount().getValue(), 4.0); TableStatistics withPartitionFilter = TableStatisticsReader.makeTableStatistics( @@ -722,10 +717,12 @@ public void testStatsFilePruning() table, snapshotId, TupleDomain.withColumnDomains(ImmutableMap.of( - new IcebergColumnHandle(ColumnIdentity.primitiveColumnIdentity(1, "b"), INTEGER, ImmutableList.of(), INTEGER, Optional.empty()), + new IcebergColumnHandle(ColumnIdentity.primitiveColumnIdentity(1, "b"), INTEGER, ImmutableList.of(), INTEGER, true, Optional.empty()), Domain.singleValue(INTEGER, 10L))), TupleDomain.all(), - true); + ImmutableSet.of(), + true, + fileSystemFactory.create(SESSION)); assertEquals(withPartitionFilter.getRowCount().getValue(), 3.0); TableStatistics withUnenforcedFilter = TableStatisticsReader.makeTableStatistics( @@ -734,9 +731,11 @@ public void testStatsFilePruning() snapshotId, TupleDomain.all(), TupleDomain.withColumnDomains(ImmutableMap.of( - new IcebergColumnHandle(ColumnIdentity.primitiveColumnIdentity(0, "a"), INTEGER, ImmutableList.of(), INTEGER, Optional.empty()), + new IcebergColumnHandle(ColumnIdentity.primitiveColumnIdentity(0, "a"), INTEGER, ImmutableList.of(), INTEGER, true, Optional.empty()), Domain.create(ValueSet.ofRanges(Range.greaterThan(INTEGER, 100L)), true))), - true); + ImmutableSet.of(), + true, + fileSystemFactory.create(SESSION)); assertEquals(withUnenforcedFilter.getRowCount().getValue(), 2.0); } } @@ -796,36 +795,37 @@ private void writeEqualityDeleteToNationTable(Table icebergTable, Optional partitionSpec, Optional partitionData, Map overwriteValues) + private void writeEqualityDeleteToNationTable( + Table icebergTable, + Optional partitionSpec, + Optional partitionData, + Map overwriteValues) throws Exception { - Path metadataDir = new Path(metastoreDir.toURI()); - String deleteFileName = "delete_file_" + UUID.randomUUID(); - FileIO fileIo = new ForwardingFileIo(fileSystemFactory.create(SESSION)); - - Schema deleteRowSchema = icebergTable.schema().select(overwriteValues.keySet()); - List equalityFieldIds = overwriteValues.keySet().stream() - .map(name -> deleteRowSchema.findField(name).fieldId()) - .collect(toImmutableList()); - Parquet.DeleteWriteBuilder writerBuilder = Parquet.writeDeletes(fileIo.newOutputFile(new Path(metadataDir, deleteFileName).toString())) - .forTable(icebergTable) - .rowSchema(deleteRowSchema) - .createWriterFunc(GenericParquetWriter::buildWriter) - .equalityFieldIds(equalityFieldIds) - .overwrite(); - if (partitionSpec.isPresent() && partitionData.isPresent()) { - writerBuilder = writerBuilder - .withSpec(partitionSpec.get()) - .withPartition(partitionData.get()); - } - EqualityDeleteWriter writer = writerBuilder.buildEqualityWriter(); + writeEqualityDeleteToNationTableWithDeleteColumns(icebergTable, partitionSpec, partitionData, overwriteValues, Optional.empty()); + } - Record dataDelete = GenericRecord.create(deleteRowSchema); - try (Closeable ignored = writer) { - writer.write(dataDelete.copy(overwriteValues)); - } + private void writeEqualityDeleteToNationTableWithDeleteColumns( + Table icebergTable, + Optional partitionSpec, + Optional partitionData, + Map overwriteValues, + Optional> deleteFileColumns) + throws Exception + { + writeEqualityDeleteForTable(icebergTable, fileSystemFactory, partitionSpec, partitionData, overwriteValues, deleteFileColumns); + } - icebergTable.newRowDelta().addDeletes(writer.toDeleteFile()).commit(); + private void writeEqualityDeleteToNationTableWithDeleteColumns( + Table icebergTable, + Optional partitionSpec, + Optional partitionData, + Map overwriteValues, + Schema deleteRowSchema, + List equalityDeleteFieldIds) + throws Exception + { + writeEqualityDeleteForTableWithSchema(icebergTable, fileSystemFactory, partitionSpec, partitionData, deleteRowSchema, equalityDeleteFieldIds, overwriteValues); } private Table updateTableToV2(String tableName) @@ -838,21 +838,59 @@ private Table updateTableToV2(String tableName) return table; } + @Test + public void testEnvironmentContext() + { + try (TestTable table = new TestTable(getQueryRunner()::execute, "test_environment_context", "(x int)")) { + Table icebergTable = loadTable(table.getName()); + assertThat(icebergTable.currentSnapshot().summary()) + .contains(entry("engine-name", "trino"), entry("engine-version", "testversion")); + } + } + + private void testHighlyNestedFieldPartitioningWithTimestampTransform(String partitioning, String partitionDirectoryRegex, Set expectedPartitionDirectories) + { + String tableName = "test_highly_nested_field_partitioning_with_timestamp_transform_" + randomNameSuffix(); + assertUpdate("CREATE TABLE " + tableName + " (id INTEGER, grandparent ROW(parent ROW(ts TIMESTAMP(6), a INT), b INT)) WITH (partitioning = " + partitioning + ")"); + assertUpdate( + "INSERT INTO " + tableName + " VALUES " + + "(1, ROW(ROW(TIMESTAMP '2021-01-01 01:01:01.111111', 1), 1)), " + + "(2, ROW(ROW(TIMESTAMP '2022-02-02 02:02:02.222222', 2), 2)), " + + "(3, ROW(ROW(TIMESTAMP '2023-03-03 03:03:03.333333', 3), 3)), " + + "(4, ROW(ROW(TIMESTAMP '2022-02-02 02:04:04.444444', 4), 4))", + 4); + + assertThat(loadTable(tableName).newScan().planFiles()).hasSize(3); + Set partitionedDirectories = computeActual("SELECT file_path FROM \"" + tableName + "$files\"") + .getMaterializedRows().stream() + .map(entry -> extractPartitionFolder((String) entry.getField(0), partitionDirectoryRegex)) + .flatMap(Optional::stream) + .collect(toImmutableSet()); + + assertThat(partitionedDirectories).isEqualTo(expectedPartitionDirectories); + + assertQuery("SELECT id, grandparent.parent.ts, grandparent.parent.a, grandparent.b FROM " + tableName, "VALUES " + + "(1, '2021-01-01 01:01:01.111111', 1, 1), " + + "(2, '2022-02-02 02:02:02.222222', 2, 2), " + + "(3, '2023-03-03 03:03:03.333333', 3, 3), " + + "(4, '2022-02-02 02:04:04.444444', 4, 4)"); + + assertUpdate("DROP TABLE " + tableName); + } + + private Optional extractPartitionFolder(String file, String regex) + { + Pattern pattern = Pattern.compile(regex); + Matcher matcher = pattern.matcher(file); + if (matcher.matches()) { + return Optional.of(matcher.group(1)); + } + return Optional.empty(); + } + private BaseTable loadTable(String tableName) { - IcebergTableOperationsProvider tableOperationsProvider = new FileMetastoreTableOperationsProvider(fileSystemFactory); - CachingHiveMetastore cachingHiveMetastore = memoizeMetastore(metastore, 1000); - TrinoCatalog catalog = new TrinoHiveCatalog( - new CatalogName("hive"), - cachingHiveMetastore, - new TrinoViewHiveMetastore(cachingHiveMetastore, false, "trino-version", "test"), - fileSystemFactory, - new TestingTypeManager(), - tableOperationsProvider, - false, - false, - false); - return (BaseTable) loadIcebergTable(catalog, tableOperationsProvider, SESSION, new SchemaTableName("tpch", tableName)); + return IcebergTestUtils.loadTable(tableName, metastore, fileSystemFactory, "hive", "tpch"); } private List getActiveFiles(String tableName) diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestMetadataQueryOptimization.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestMetadataQueryOptimization.java index 032c26f41658..a204e7f94a52 100644 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestMetadataQueryOptimization.java +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestMetadataQueryOptimization.java @@ -35,7 +35,6 @@ import static com.google.common.io.MoreFiles.deleteRecursively; import static com.google.common.io.RecursiveDeleteOption.ALLOW_INSECURE; -import static com.google.inject.util.Modules.EMPTY_MODULE; import static io.trino.SystemSessionProperties.TASK_PARTITIONED_WRITER_COUNT; import static io.trino.plugin.hive.metastore.file.TestingFileHiveMetastore.createTestingFileHiveMetastore; import static io.trino.sql.planner.assertions.PlanMatchPattern.anyTree; @@ -75,7 +74,7 @@ protected LocalQueryRunner createLocalQueryRunner() queryRunner.createCatalog( ICEBERG_CATALOG, - new TestingIcebergConnectorFactory(Optional.of(new TestingIcebergFileMetastoreCatalogModule(metastore)), Optional.empty(), EMPTY_MODULE), + new TestingIcebergConnectorFactory(baseDir.toPath(), Optional.of(new TestingIcebergFileMetastoreCatalogModule(metastore))), ImmutableMap.of()); Database database = Database.builder() @@ -107,8 +106,8 @@ public void testOptimization() anyTree(values( ImmutableList.of("b", "c"), ImmutableList.of( - ImmutableList.of(new LongLiteral("6"), new LongLiteral("7")), - ImmutableList.of(new LongLiteral("9"), new LongLiteral("10")))))); + ImmutableList.of(new LongLiteral("9"), new LongLiteral("10")), + ImmutableList.of(new LongLiteral("6"), new LongLiteral("7")))))); assertPlan( format("SELECT DISTINCT b, c FROM %s WHERE b > 7", testTable), diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestPartitionFields.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestPartitionFields.java index 4deb541c1259..fe0ef58184bd 100644 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestPartitionFields.java +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestPartitionFields.java @@ -109,7 +109,7 @@ private static void assertInvalid(String value, String message) private static PartitionSpec parseField(String value) { - return partitionSpec(builder -> parsePartitionField(builder, value)); + return partitionSpec(builder -> parsePartitionField(builder, value, "")); } private static PartitionSpec partitionSpec(Consumer consumer) diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestSharedHiveMetastore.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestSharedHiveMetastore.java index c3a3240fd2d1..1e3235779942 100644 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestSharedHiveMetastore.java +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestSharedHiveMetastore.java @@ -25,7 +25,6 @@ import java.nio.file.Path; -import static io.trino.plugin.hive.metastore.file.TestingFileHiveMetastore.createTestingFileHiveMetastore; import static io.trino.plugin.iceberg.IcebergQueryRunner.ICEBERG_CATALOG; import static io.trino.plugin.tpch.TpchMetadata.TINY_SCHEMA_NAME; import static io.trino.testing.QueryAssertions.copyTpchTables; @@ -74,7 +73,7 @@ protected QueryRunner createQueryRunner() "hive.metastore.catalog.dir", dataDirectory.toString(), "iceberg.hive-catalog-name", "hive")); - queryRunner.installPlugin(new TestingHivePlugin(createTestingFileHiveMetastore(dataDirectory.toFile()))); + queryRunner.installPlugin(new TestingHivePlugin(dataDirectory)); queryRunner.createCatalog(HIVE_CATALOG, "hive", ImmutableMap.of("hive.allow-drop-table", "true")); queryRunner.createCatalog( "hive_with_redirections", @@ -101,10 +100,10 @@ protected String getExpectedHiveCreateSchema(String catalogName) { String expectedHiveCreateSchema = "CREATE SCHEMA %s.%s\n" + "WITH (\n" + - " location = 'file:%s/%s'\n" + + " location = 'local:///%s'\n" + ")"; - return format(expectedHiveCreateSchema, catalogName, schema, dataDirectory, schema); + return format(expectedHiveCreateSchema, catalogName, schema, schema); } @Override diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestSortFieldUtils.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestSortFieldUtils.java index 42e5c524f380..b73047315c1e 100644 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestSortFieldUtils.java +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestSortFieldUtils.java @@ -51,14 +51,14 @@ public void testParse() // uppercase assertParse("ORDER_KEY ASC NULLS LAST", sortOrder(builder -> builder.asc("order_key", NullOrder.NULLS_LAST))); assertParse("ORDER_KEY DESC NULLS FIRST", sortOrder(builder -> builder.desc("order_key", NullOrder.NULLS_FIRST))); - assertDoesNotParse("\"ORDER_KEY\" ASC NULLS LAST", "Uppercase characters in identifier '\"ORDER_KEY\"' are not supported."); - assertDoesNotParse("\"ORDER_KEY\" DESC NULLS FIRST", "Uppercase characters in identifier '\"ORDER_KEY\"' are not supported."); + assertDoesNotParse("\"ORDER_KEY\" ASC NULLS LAST", "Cannot find field 'ORDER_KEY' .*"); + assertDoesNotParse("\"ORDER_KEY\" DESC NULLS FIRST", "Cannot find field 'ORDER_KEY' .*"); // mixed case assertParse("OrDER_keY Asc NullS LAst", sortOrder(builder -> builder.asc("order_key", NullOrder.NULLS_LAST))); assertParse("OrDER_keY Desc NullS FIrsT", sortOrder(builder -> builder.desc("order_key", NullOrder.NULLS_FIRST))); - assertDoesNotParse("\"OrDER_keY\" Asc NullS LAst", "Uppercase characters in identifier '\"OrDER_keY\"' are not supported."); - assertDoesNotParse("\"OrDER_keY\" Desc NullS FIrsT", "Uppercase characters in identifier '\"OrDER_keY\"' are not supported."); + assertDoesNotParse("\"OrDER_keY\" Asc NullS LAst", "Cannot find field 'OrDER_keY' .*"); + assertDoesNotParse("\"OrDER_keY\" Desc NullS FIrsT", "Cannot find field 'OrDER_keY' .*"); assertParse("comment", sortOrder(builder -> builder.asc("comment"))); assertParse("\"comment\"", sortOrder(builder -> builder.asc("comment"))); @@ -106,13 +106,13 @@ private static void assertParse(@Language("SQL") String value, SortOrder expecte private static void assertDoesNotParse(@Language("SQL") String value) { - assertDoesNotParse(value, "Unable to parse sort field: [%s]".formatted(value)); + assertDoesNotParse(value, "\\QUnable to parse sort field: [%s]".formatted(value)); } private static void assertDoesNotParse(@Language("SQL") String value, String expectedMessage) { assertThatThrownBy(() -> parseField(value)) - .hasMessage(expectedMessage); + .hasMessageMatching(expectedMessage); } private static SortOrder parseField(String value) diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestStructLikeWrapperWithFieldIdToIndex.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestStructLikeWrapperWithFieldIdToIndex.java index 09ad20afc743..73b8dadf5439 100644 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestStructLikeWrapperWithFieldIdToIndex.java +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestStructLikeWrapperWithFieldIdToIndex.java @@ -36,8 +36,8 @@ public void testStructLikeWrapperWithFieldIdToIndexEquals() NestedField.optional(1001, "level", IntegerType.get())); PartitionData firstPartitionData = PartitionData.fromJson("{\"partitionValues\":[\"ERROR\",\"449245\"]}", new Type[] {StringType.get(), IntegerType.get()}); PartitionData secondPartitionData = PartitionData.fromJson("{\"partitionValues\":[\"449245\",\"ERROR\"]}", new Type[] {IntegerType.get(), StringType.get()}); - PartitionTable.StructLikeWrapperWithFieldIdToIndex first = new PartitionTable.StructLikeWrapperWithFieldIdToIndex(StructLikeWrapper.forType(firstStructType).set(firstPartitionData), firstStructType); - PartitionTable.StructLikeWrapperWithFieldIdToIndex second = new PartitionTable.StructLikeWrapperWithFieldIdToIndex(StructLikeWrapper.forType(secondStructType).set(secondPartitionData), secondStructType); + StructLikeWrapperWithFieldIdToIndex first = new StructLikeWrapperWithFieldIdToIndex(StructLikeWrapper.forType(firstStructType).set(firstPartitionData), firstStructType); + StructLikeWrapperWithFieldIdToIndex second = new StructLikeWrapperWithFieldIdToIndex(StructLikeWrapper.forType(secondStructType).set(secondPartitionData), secondStructType); assertThat(first).isNotEqualTo(second); } } diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestingIcebergConnectorFactory.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestingIcebergConnectorFactory.java index fbc412b548c2..84edaa76bc6d 100644 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestingIcebergConnectorFactory.java +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestingIcebergConnectorFactory.java @@ -13,30 +13,47 @@ */ package io.trino.plugin.iceberg; +import com.google.common.collect.ImmutableMap; import com.google.inject.Module; import io.trino.filesystem.TrinoFileSystemFactory; +import io.trino.filesystem.local.LocalFileSystemFactory; +import io.trino.plugin.hive.metastore.file.FileHiveMetastoreConfig; import io.trino.spi.connector.Connector; import io.trino.spi.connector.ConnectorContext; import io.trino.spi.connector.ConnectorFactory; +import java.nio.file.Path; import java.util.Map; import java.util.Optional; -import static io.trino.plugin.iceberg.InternalIcebergConnectorFactory.createConnector; +import static com.google.inject.multibindings.MapBinder.newMapBinder; +import static io.airlift.configuration.ConfigBinder.configBinder; +import static io.trino.plugin.iceberg.IcebergConnectorFactory.createConnector; import static java.util.Objects.requireNonNull; public class TestingIcebergConnectorFactory implements ConnectorFactory { private final Optional icebergCatalogModule; - private final Optional fileSystemFactory; private final Module module; - public TestingIcebergConnectorFactory(Optional icebergCatalogModule, Optional fileSystemFactory, Module module) + public TestingIcebergConnectorFactory(Path localFileSystemRootPath) { + this(localFileSystemRootPath, Optional.empty()); + } + + @Deprecated + public TestingIcebergConnectorFactory( + Path localFileSystemRootPath, + Optional icebergCatalogModule) + { + boolean ignored = localFileSystemRootPath.toFile().mkdirs(); this.icebergCatalogModule = requireNonNull(icebergCatalogModule, "icebergCatalogModule is null"); - this.fileSystemFactory = requireNonNull(fileSystemFactory, "fileSystemFactory is null"); - this.module = requireNonNull(module, "module is null"); + this.module = binder -> { + newMapBinder(binder, String.class, TrinoFileSystemFactory.class) + .addBinding("local").toInstance(new LocalFileSystemFactory(localFileSystemRootPath)); + configBinder(binder).bindConfigDefaults(FileHiveMetastoreConfig.class, config -> config.setCatalogDirectory("local:///")); + }; } @Override @@ -48,6 +65,12 @@ public String getName() @Override public Connector create(String catalogName, Map config, ConnectorContext context) { - return createConnector(catalogName, config, context, module, icebergCatalogModule, fileSystemFactory); + if (!config.containsKey("iceberg.catalog.type")) { + config = ImmutableMap.builder() + .putAll(config) + .put("iceberg.catalog.type", "TESTING_FILE_METASTORE") + .buildOrThrow(); + } + return createConnector(catalogName, config, context, module, icebergCatalogModule); } } diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestingIcebergPlugin.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestingIcebergPlugin.java index 7ed7cb4c18c8..fa94c7db61c3 100644 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestingIcebergPlugin.java +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestingIcebergPlugin.java @@ -15,9 +15,9 @@ import com.google.common.collect.ImmutableList; import com.google.inject.Module; -import io.trino.filesystem.TrinoFileSystemFactory; import io.trino.spi.connector.ConnectorFactory; +import java.nio.file.Path; import java.util.List; import java.util.Optional; @@ -27,15 +27,19 @@ public class TestingIcebergPlugin extends IcebergPlugin { + private final Path localFileSystemRootPath; private final Optional icebergCatalogModule; - private final Optional fileSystemFactory; - private final Module module; - public TestingIcebergPlugin(Optional icebergCatalogModule, Optional fileSystemFactory, Module module) + public TestingIcebergPlugin(Path localFileSystemRootPath) { + this(localFileSystemRootPath, Optional.empty()); + } + + @Deprecated + public TestingIcebergPlugin(Path localFileSystemRootPath, Optional icebergCatalogModule) + { + this.localFileSystemRootPath = requireNonNull(localFileSystemRootPath, "localFileSystemRootPath is null"); this.icebergCatalogModule = requireNonNull(icebergCatalogModule, "icebergCatalogModule is null"); - this.fileSystemFactory = requireNonNull(fileSystemFactory, "fileSystemFactory is null"); - this.module = requireNonNull(module, "module is null"); } @Override @@ -44,6 +48,6 @@ public Iterable getConnectorFactories() List connectorFactories = ImmutableList.copyOf(super.getConnectorFactories()); verify(connectorFactories.size() == 1, "Unexpected connector factories: %s", connectorFactories); - return ImmutableList.of(new TestingIcebergConnectorFactory(icebergCatalogModule, fileSystemFactory, module)); + return ImmutableList.of(new TestingIcebergConnectorFactory(localFileSystemRootPath, icebergCatalogModule)); } } diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/BaseTrinoCatalogTest.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/BaseTrinoCatalogTest.java index da32418fad8d..770aae90da89 100644 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/BaseTrinoCatalogTest.java +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/BaseTrinoCatalogTest.java @@ -16,12 +16,17 @@ import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import io.airlift.log.Logger; +import io.trino.plugin.base.util.AutoCloseableCloser; import io.trino.plugin.hive.NodeVersion; +import io.trino.plugin.hive.metastore.TableInfo; +import io.trino.plugin.hive.metastore.TableInfo.ExtendedRelationType; import io.trino.plugin.iceberg.CommitTaskData; +import io.trino.plugin.iceberg.IcebergFileFormat; import io.trino.plugin.iceberg.IcebergMetadata; import io.trino.plugin.iceberg.TableStatisticsWriter; import io.trino.spi.TrinoException; import io.trino.spi.connector.CatalogHandle; +import io.trino.spi.connector.ConnectorMaterializedViewDefinition; import io.trino.spi.connector.ConnectorMetadata; import io.trino.spi.connector.ConnectorSession; import io.trino.spi.connector.ConnectorViewDefinition; @@ -42,14 +47,24 @@ import java.io.UncheckedIOException; import java.nio.file.Files; import java.nio.file.Path; +import java.time.Duration; import java.util.HashMap; import java.util.Map; import java.util.Optional; +import static com.google.common.util.concurrent.MoreExecutors.directExecutor; +import static com.google.common.util.concurrent.MoreExecutors.newDirectExecutorService; import static io.airlift.json.JsonCodec.jsonCodec; import static io.trino.plugin.hive.HiveErrorCode.HIVE_DATABASE_LOCATION_ERROR; +import static io.trino.plugin.hive.metastore.TableInfo.ExtendedRelationType.TABLE; +import static io.trino.plugin.hive.metastore.TableInfo.ExtendedRelationType.TRINO_MATERIALIZED_VIEW; +import static io.trino.plugin.hive.metastore.TableInfo.ExtendedRelationType.TRINO_VIEW; import static io.trino.plugin.iceberg.IcebergSchemaProperties.LOCATION_PROPERTY; +import static io.trino.plugin.iceberg.IcebergTableProperties.FILE_FORMAT_PROPERTY; +import static io.trino.plugin.iceberg.IcebergTableProperties.FORMAT_VERSION_PROPERTY; import static io.trino.plugin.iceberg.IcebergUtil.quotedTableName; +import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED; +import static io.trino.spi.type.BigintType.BIGINT; import static io.trino.sql.planner.TestingPlannerContext.PLANNER_CONTEXT; import static io.trino.testing.TestingConnectorSession.SESSION; import static io.trino.testing.TestingNames.randomNameSuffix; @@ -112,10 +127,15 @@ public void testNonLowercaseNamespace() CatalogHandle.fromId("iceberg:NORMAL:v12345"), jsonCodec(CommitTaskData.class), catalog, - connectorIdentity -> { + (connectorIdentity, fileIoProperties) -> { throw new UnsupportedOperationException(); }, - new TableStatisticsWriter(new NodeVersion("test-version"))); + new TableStatisticsWriter(new NodeVersion("test-version")), + Optional.empty(), + false, + ignore -> false, + newDirectExecutorService(), + directExecutor()); assertThat(icebergMetadata.schemaExists(SESSION, namespace)).as("icebergMetadata.schemaExists(namespace)") .isFalse(); assertThat(icebergMetadata.schemaExists(SESSION, schema)).as("icebergMetadata.schemaExists(schema)") @@ -147,11 +167,11 @@ public void testCreateTable() new Schema(Types.NestedField.of(1, true, "col1", Types.LongType.get())), PartitionSpec.unpartitioned(), SortOrder.unsorted(), - tableLocation, + Optional.of(tableLocation), tableProperties) .commitTransaction(); - assertThat(catalog.listTables(SESSION, Optional.of(namespace))).contains(schemaTableName); - assertThat(catalog.listTables(SESSION, Optional.empty())).contains(schemaTableName); + assertThat(catalog.listTables(SESSION, Optional.of(namespace))).contains(new TableInfo(schemaTableName, TABLE)); + assertThat(catalog.listTables(SESSION, Optional.empty())).contains(new TableInfo(schemaTableName, TABLE)); Table icebergTable = catalog.loadTable(SESSION, schemaTableName); assertEquals(icebergTable.name(), quotedTableName(schemaTableName)); @@ -163,8 +183,8 @@ public void testCreateTable() assertThat(icebergTable.properties()).containsAllEntriesOf(tableProperties); catalog.dropTable(SESSION, schemaTableName); - assertThat(catalog.listTables(SESSION, Optional.of(namespace))).doesNotContain(schemaTableName); - assertThat(catalog.listTables(SESSION, Optional.empty())).doesNotContain(schemaTableName); + assertThat(catalog.listTables(SESSION, Optional.of(namespace))).doesNotContain(new TableInfo(schemaTableName, TABLE)); + assertThat(catalog.listTables(SESSION, Optional.empty())).doesNotContain(new TableInfo(schemaTableName, TABLE)); } finally { try { @@ -208,11 +228,11 @@ public void testCreateWithSortTable() tableSchema, PartitionSpec.unpartitioned(), sortOrder, - tableLocation, + Optional.of(tableLocation), ImmutableMap.of()) .commitTransaction(); - assertThat(catalog.listTables(SESSION, Optional.of(namespace))).contains(schemaTableName); - assertThat(catalog.listTables(SESSION, Optional.empty())).contains(schemaTableName); + assertThat(catalog.listTables(SESSION, Optional.of(namespace))).contains(new TableInfo(schemaTableName, TABLE)); + assertThat(catalog.listTables(SESSION, Optional.empty())).contains(new TableInfo(schemaTableName, TABLE)); Table icebergTable = catalog.loadTable(SESSION, schemaTableName); assertEquals(icebergTable.name(), quotedTableName(schemaTableName)); @@ -263,23 +283,23 @@ public void testRenameTable() new Schema(Types.NestedField.of(1, true, "col1", Types.LongType.get())), PartitionSpec.unpartitioned(), SortOrder.unsorted(), - arbitraryTableLocation(catalog, SESSION, sourceSchemaTableName), + Optional.of(arbitraryTableLocation(catalog, SESSION, sourceSchemaTableName)), ImmutableMap.of()) .commitTransaction(); - assertThat(catalog.listTables(SESSION, Optional.of(namespace))).contains(sourceSchemaTableName); + assertThat(catalog.listTables(SESSION, Optional.of(namespace))).contains(new TableInfo(sourceSchemaTableName, TABLE)); // Rename within the same schema SchemaTableName targetSchemaTableName = new SchemaTableName(sourceSchemaTableName.getSchemaName(), "newTableName"); catalog.renameTable(SESSION, sourceSchemaTableName, targetSchemaTableName); - assertThat(catalog.listTables(SESSION, Optional.of(namespace))).doesNotContain(sourceSchemaTableName); - assertThat(catalog.listTables(SESSION, Optional.of(namespace))).contains(targetSchemaTableName); + assertThat(catalog.listTables(SESSION, Optional.empty()).stream().map(TableInfo::tableName).toList()).doesNotContain(sourceSchemaTableName); + assertThat(catalog.listTables(SESSION, Optional.of(namespace))).contains(new TableInfo(targetSchemaTableName, TABLE)); // Move to a different schema sourceSchemaTableName = targetSchemaTableName; targetSchemaTableName = new SchemaTableName(targetNamespace, sourceSchemaTableName.getTableName()); catalog.renameTable(SESSION, sourceSchemaTableName, targetSchemaTableName); - assertThat(catalog.listTables(SESSION, Optional.of(namespace))).doesNotContain(sourceSchemaTableName); - assertThat(catalog.listTables(SESSION, Optional.of(targetNamespace))).contains(targetSchemaTableName); + assertThat(catalog.listTables(SESSION, Optional.of(namespace)).stream().map(TableInfo::tableName).toList()).doesNotContain(sourceSchemaTableName); + assertThat(catalog.listTables(SESSION, Optional.of(targetNamespace))).contains(new TableInfo(targetSchemaTableName, TABLE)); catalog.dropTable(SESSION, targetSchemaTableName); } @@ -352,7 +372,7 @@ public void testView() Optional.empty(), Optional.empty(), ImmutableList.of( - new ConnectorViewDefinition.ViewColumn("name", VarcharType.createVarcharType(25).getTypeId(), Optional.empty())), + new ConnectorViewDefinition.ViewColumn("name", VarcharType.createUnboundedVarcharType().getTypeId(), Optional.empty())), Optional.empty(), Optional.of(SESSION.getUser()), false); @@ -361,25 +381,23 @@ public void testView() catalog.createNamespace(SESSION, namespace, ImmutableMap.of(), new TrinoPrincipal(PrincipalType.USER, SESSION.getUser())); catalog.createView(SESSION, schemaTableName, viewDefinition, false); - assertThat(catalog.listTables(SESSION, Optional.of(namespace))).contains(schemaTableName); - assertThat(catalog.listViews(SESSION, Optional.of(namespace))).contains(schemaTableName); + assertThat(catalog.listTables(SESSION, Optional.of(namespace)).stream()).contains(new TableInfo(schemaTableName, getViewType())); Map views = catalog.getViews(SESSION, Optional.of(schemaTableName.getSchemaName())); - assertEquals(views.size(), 1); + assertThat(views).hasSize(1); assertViewDefinition(views.get(schemaTableName), viewDefinition); assertViewDefinition(catalog.getView(SESSION, schemaTableName).orElseThrow(), viewDefinition); catalog.renameView(SESSION, schemaTableName, renamedSchemaTableName); - assertThat(catalog.listTables(SESSION, Optional.of(namespace))).doesNotContain(schemaTableName); - assertThat(catalog.listViews(SESSION, Optional.of(namespace))).doesNotContain(schemaTableName); + assertThat(catalog.listTables(SESSION, Optional.of(namespace)).stream().map(TableInfo::tableName).toList()).doesNotContain(schemaTableName); views = catalog.getViews(SESSION, Optional.of(schemaTableName.getSchemaName())); - assertEquals(views.size(), 1); + assertThat(views).hasSize(1); assertViewDefinition(views.get(renamedSchemaTableName), viewDefinition); assertViewDefinition(catalog.getView(SESSION, renamedSchemaTableName).orElseThrow(), viewDefinition); assertThat(catalog.getView(SESSION, schemaTableName)).isEmpty(); catalog.dropView(SESSION, renamedSchemaTableName); - assertThat(catalog.listTables(SESSION, Optional.of(namespace))) + assertThat(catalog.listTables(SESSION, Optional.empty()).stream().map(TableInfo::tableName).toList()) .doesNotContain(renamedSchemaTableName); } finally { @@ -392,6 +410,145 @@ public void testView() } } + protected ExtendedRelationType getViewType() + { + return TRINO_VIEW; + } + + @Test + public void testListTables() + throws Exception + { + TrinoCatalog catalog = createTrinoCatalog(false); + TrinoPrincipal principal = new TrinoPrincipal(PrincipalType.USER, SESSION.getUser()); + try (AutoCloseableCloser closer = AutoCloseableCloser.create()) { + String ns1 = "ns1"; + String ns2 = "ns2"; + + catalog.createNamespace(SESSION, ns1, defaultNamespaceProperties(ns1), principal); + catalog.createNamespace(SESSION, ns2, defaultNamespaceProperties(ns2), principal); + SchemaTableName table1 = new SchemaTableName(ns1, "t1"); + SchemaTableName table2 = new SchemaTableName(ns2, "t2"); + catalog.newCreateTableTransaction( + SESSION, + table1, + new Schema(Types.NestedField.of(1, true, "col1", Types.LongType.get())), + PartitionSpec.unpartitioned(), + SortOrder.unsorted(), + Optional.of(arbitraryTableLocation(catalog, SESSION, table1)), + ImmutableMap.of()) + .commitTransaction(); + closer.register(() -> catalog.dropTable(SESSION, table1)); + + catalog.newCreateTableTransaction( + SESSION, + table2, + new Schema(Types.NestedField.of(1, true, "col1", Types.LongType.get())), + PartitionSpec.unpartitioned(), + SortOrder.unsorted(), + Optional.of(arbitraryTableLocation(catalog, SESSION, table2)), + ImmutableMap.of()) + .commitTransaction(); + closer.register(() -> catalog.dropTable(SESSION, table2)); + + ImmutableList.Builder allTables = ImmutableList.builder() + .add(new TableInfo(table1, TABLE)) + .add(new TableInfo(table2, TABLE)); + + ImmutableList.Builder icebergTables = ImmutableList.builder() + .add(table1) + .add(table2); + SchemaTableName view = new SchemaTableName(ns2, "view"); + try { + catalog.createView( + SESSION, + view, + new ConnectorViewDefinition( + "SELECT name FROM local.tiny.nation", + Optional.empty(), + Optional.empty(), + ImmutableList.of( + new ConnectorViewDefinition.ViewColumn("name", VarcharType.createUnboundedVarcharType().getTypeId(), Optional.empty())), + Optional.empty(), + Optional.of(SESSION.getUser()), + false), + false); + closer.register(() -> catalog.dropView(SESSION, view)); + allTables.add(new TableInfo(view, getViewType())); + } + catch (TrinoException e) { + assertThat(e.getErrorCode()).isEqualTo(NOT_SUPPORTED.toErrorCode()); + } + + try { + SchemaTableName materializedView = new SchemaTableName(ns2, "mv"); + createMaterializedView( + SESSION, + catalog, + materializedView, + someMaterializedView(), + ImmutableMap.of( + FILE_FORMAT_PROPERTY, IcebergFileFormat.PARQUET, + FORMAT_VERSION_PROPERTY, 1), + false, + false); + closer.register(() -> catalog.dropMaterializedView(SESSION, materializedView)); + allTables.add(new TableInfo(materializedView, TRINO_MATERIALIZED_VIEW)); + } + catch (TrinoException e) { + assertThat(e.getErrorCode()).isEqualTo(NOT_SUPPORTED.toErrorCode()); + } + + createExternalIcebergTable(catalog, ns2, closer).ifPresent(table -> { + allTables.add(new TableInfo(table, TABLE)); + icebergTables.add(table); + }); + createExternalNonIcebergTable(catalog, ns2, closer).ifPresent(table -> { + allTables.add(new TableInfo(table, TABLE)); + }); + + // No namespace provided, all tables across all namespaces should be returned + assertThat(catalog.listTables(SESSION, Optional.empty())).containsAll(allTables.build()); + assertThat(catalog.listIcebergTables(SESSION, Optional.empty())).containsAll(icebergTables.build()); + // Namespace is provided and exists + assertThat(catalog.listTables(SESSION, Optional.of(ns1))).containsExactly(new TableInfo(table1, TABLE)); + assertThat(catalog.listIcebergTables(SESSION, Optional.of(ns1))).containsExactly(table1); + // Namespace is provided and does not exist + assertThat(catalog.listTables(SESSION, Optional.of("non_existing"))).isEmpty(); + assertThat(catalog.listIcebergTables(SESSION, Optional.of("non_existing"))).isEmpty(); + } + } + + protected void createMaterializedView( + ConnectorSession session, + TrinoCatalog catalog, + SchemaTableName materializedView, + ConnectorMaterializedViewDefinition materializedViewDefinition, + Map properties, + boolean replace, + boolean ignoreExisting) + { + catalog.createMaterializedView( + session, + materializedView, + materializedViewDefinition, + properties, + replace, + ignoreExisting); + } + + protected Optional createExternalIcebergTable(TrinoCatalog catalog, String namespace, AutoCloseableCloser closer) + throws Exception + { + return Optional.empty(); + } + + protected Optional createExternalNonIcebergTable(TrinoCatalog catalog, String namespace, AutoCloseableCloser closer) + throws Exception + { + return Optional.empty(); + } + private String arbitraryTableLocation(TrinoCatalog catalog, ConnectorSession session, SchemaTableName schemaTableName) throws Exception { @@ -426,4 +583,18 @@ private void assertViewColumnDefinition(ConnectorViewDefinition.ViewColumn actua assertEquals(actualViewColumn.getName(), expectedViewColumn.getName()); assertEquals(actualViewColumn.getType(), expectedViewColumn.getType()); } + + private static ConnectorMaterializedViewDefinition someMaterializedView() + { + return new ConnectorMaterializedViewDefinition( + "select 1", + Optional.empty(), + Optional.empty(), + Optional.empty(), + ImmutableList.of(new ConnectorMaterializedViewDefinition.Column("test", BIGINT.getTypeId(), Optional.empty())), + Optional.of(Duration.ZERO), + Optional.empty(), + Optional.of("owner"), + ImmutableMap.of()); + } } diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/file/TestIcebergFileMetastoreCreateTableFailure.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/file/TestIcebergFileMetastoreCreateTableFailure.java index a4e0927cebc3..93d2279187eb 100644 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/file/TestIcebergFileMetastoreCreateTableFailure.java +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/file/TestIcebergFileMetastoreCreateTableFailure.java @@ -14,6 +14,7 @@ package io.trino.plugin.iceberg.catalog.file; import io.trino.Session; +import io.trino.filesystem.local.LocalFileSystemFactory; import io.trino.plugin.hive.NodeVersion; import io.trino.plugin.hive.metastore.HiveMetastore; import io.trino.plugin.hive.metastore.HiveMetastoreConfig; @@ -35,8 +36,6 @@ import static com.google.common.io.MoreFiles.deleteRecursively; import static com.google.common.io.RecursiveDeleteOption.ALLOW_INSECURE; -import static com.google.inject.util.Modules.EMPTY_MODULE; -import static io.trino.plugin.hive.HiveTestUtils.HDFS_ENVIRONMENT; import static io.trino.testing.TestingNames.randomNameSuffix; import static io.trino.testing.TestingSession.testSessionBuilder; import static org.assertj.core.api.Assertions.assertThat; @@ -61,10 +60,10 @@ protected DistributedQueryRunner createQueryRunner() // Using FileHiveMetastore as approximation of HMS this.metastore = new FileHiveMetastore( new NodeVersion("testversion"), - HDFS_ENVIRONMENT, + new LocalFileSystemFactory(Path.of(dataDirectory.toString())), new HiveMetastoreConfig().isHideDeltaLakeTables(), new FileHiveMetastoreConfig() - .setCatalogDirectory(dataDirectory.toString())) + .setCatalogDirectory("local://")) { @Override public synchronized void createTable(Table table, PrincipalPrivileges principalPrivileges) @@ -79,7 +78,7 @@ public synchronized void createTable(Table table, PrincipalPrivileges principalP .build(); DistributedQueryRunner queryRunner = DistributedQueryRunner.builder(session).build(); - queryRunner.installPlugin(new TestingIcebergPlugin(Optional.of(new TestingIcebergFileMetastoreCatalogModule(metastore)), Optional.empty(), EMPTY_MODULE)); + queryRunner.installPlugin(new TestingIcebergPlugin(dataDirectory, Optional.of(new TestingIcebergFileMetastoreCatalogModule(metastore)))); queryRunner.createCatalog(ICEBERG_CATALOG, "iceberg"); queryRunner.execute("CREATE SCHEMA " + SCHEMA_NAME); diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/file/TestIcebergFileMetastoreTableOperationsInsertFailure.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/file/TestIcebergFileMetastoreTableOperationsInsertFailure.java index 5da898650dba..a985dde1bedf 100644 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/file/TestIcebergFileMetastoreTableOperationsInsertFailure.java +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/file/TestIcebergFileMetastoreTableOperationsInsertFailure.java @@ -15,6 +15,7 @@ import com.google.common.collect.ImmutableMap; import io.trino.Session; +import io.trino.filesystem.local.LocalFileSystemFactory; import io.trino.metadata.InternalFunctionBundle; import io.trino.plugin.hive.NodeVersion; import io.trino.plugin.hive.metastore.Database; @@ -35,12 +36,11 @@ import java.io.File; import java.nio.file.Files; +import java.util.Map; import java.util.Optional; import static com.google.common.io.MoreFiles.deleteRecursively; import static com.google.common.io.RecursiveDeleteOption.ALLOW_INSECURE; -import static com.google.inject.util.Modules.EMPTY_MODULE; -import static io.trino.plugin.hive.HiveTestUtils.HDFS_ENVIRONMENT; import static io.trino.testing.TestingSession.testSessionBuilder; import static java.lang.String.format; import static org.assertj.core.api.Assertions.assertThatThrownBy; @@ -66,16 +66,15 @@ protected LocalQueryRunner createQueryRunner() HiveMetastore metastore = new FileHiveMetastore( new NodeVersion("testversion"), - HDFS_ENVIRONMENT, + new LocalFileSystemFactory(baseDir.toPath()), new HiveMetastoreConfig().isHideDeltaLakeTables(), new FileHiveMetastoreConfig() - .setCatalogDirectory(baseDir.toURI().toString()) - .setMetastoreUser("test")) + .setCatalogDirectory("local://")) { @Override - public synchronized void replaceTable(String databaseName, String tableName, Table newTable, PrincipalPrivileges principalPrivileges) + public synchronized void replaceTable(String databaseName, String tableName, Table newTable, PrincipalPrivileges principalPrivileges, Map environmentContext) { - super.replaceTable(databaseName, tableName, newTable, principalPrivileges); + super.replaceTable(databaseName, tableName, newTable, principalPrivileges, environmentContext); throw new RuntimeException("Test-simulated metastore timeout exception"); } }; @@ -87,7 +86,7 @@ public synchronized void replaceTable(String databaseName, String tableName, Tab queryRunner.createCatalog( ICEBERG_CATALOG, - new TestingIcebergConnectorFactory(Optional.of(new TestingIcebergFileMetastoreCatalogModule(metastore)), Optional.empty(), EMPTY_MODULE), + new TestingIcebergConnectorFactory(baseDir.toPath(), Optional.of(new TestingIcebergFileMetastoreCatalogModule(metastore))), ImmutableMap.of()); Database database = Database.builder() diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/file/TestTrinoHiveCatalogWithFileMetastore.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/file/TestTrinoHiveCatalogWithFileMetastore.java index 816bba04112a..373b9fe750ee 100644 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/file/TestTrinoHiveCatalogWithFileMetastore.java +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/file/TestTrinoHiveCatalogWithFileMetastore.java @@ -18,6 +18,7 @@ import io.trino.plugin.hive.TrinoViewHiveMetastore; import io.trino.plugin.hive.metastore.HiveMetastore; import io.trino.plugin.hive.metastore.cache.CachingHiveMetastore; +import io.trino.plugin.iceberg.IcebergConfig; import io.trino.plugin.iceberg.catalog.BaseTrinoCatalogTest; import io.trino.plugin.iceberg.catalog.TrinoCatalog; import io.trino.plugin.iceberg.catalog.hms.TrinoHiveCatalog; @@ -32,8 +33,9 @@ import static com.google.common.io.MoreFiles.deleteRecursively; import static com.google.common.io.RecursiveDeleteOption.ALLOW_INSECURE; +import static com.google.common.util.concurrent.MoreExecutors.directExecutor; import static io.trino.plugin.hive.HiveTestUtils.HDFS_FILE_SYSTEM_FACTORY; -import static io.trino.plugin.hive.metastore.cache.CachingHiveMetastore.memoizeMetastore; +import static io.trino.plugin.hive.metastore.cache.CachingHiveMetastore.createPerTransactionCache; import static io.trino.plugin.hive.metastore.file.TestingFileHiveMetastore.createTestingFileHiveMetastore; public class TestTrinoHiveCatalogWithFileMetastore @@ -62,7 +64,7 @@ public void tearDown() protected TrinoCatalog createTrinoCatalog(boolean useUniqueTableLocations) { TrinoFileSystemFactory fileSystemFactory = HDFS_FILE_SYSTEM_FACTORY; - CachingHiveMetastore cachingHiveMetastore = memoizeMetastore(metastore, 1000); + CachingHiveMetastore cachingHiveMetastore = createPerTransactionCache(metastore, 1000); return new TrinoHiveCatalog( new CatalogName("catalog"), cachingHiveMetastore, @@ -72,6 +74,8 @@ protected TrinoCatalog createTrinoCatalog(boolean useUniqueTableLocations) new FileMetastoreTableOperationsProvider(fileSystemFactory), useUniqueTableLocations, false, - false); + false, + new IcebergConfig().isHideMaterializedViewStorageTable(), + directExecutor()); } } diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/glue/GlueMetastoreMethod.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/glue/GlueMetastoreMethod.java index b9ef1975c10f..7e69a37a32de 100644 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/glue/GlueMetastoreMethod.java +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/glue/GlueMetastoreMethod.java @@ -14,7 +14,7 @@ package io.trino.plugin.iceberg.catalog.glue; import com.google.common.math.DoubleMath; -import io.trino.plugin.hive.aws.AwsApiCallStats; +import io.trino.plugin.hive.metastore.glue.AwsApiCallStats; import io.trino.plugin.hive.metastore.glue.GlueMetastoreStats; import java.math.RoundingMode; diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/glue/TestIcebergGlueCatalogAccessOperations.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/glue/TestIcebergGlueCatalogAccessOperations.java index ba7d3f290a0a..141994b91efb 100644 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/glue/TestIcebergGlueCatalogAccessOperations.java +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/glue/TestIcebergGlueCatalogAccessOperations.java @@ -14,7 +14,6 @@ package io.trino.plugin.iceberg.catalog.glue; import com.google.common.collect.HashMultiset; -import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableMultiset; import com.google.common.collect.Multiset; import com.google.inject.Binder; @@ -25,10 +24,11 @@ import io.airlift.log.Logger; import io.trino.Session; import io.trino.filesystem.TrackingFileSystemFactory; -import io.trino.filesystem.hdfs.HdfsFileSystemFactory; import io.trino.plugin.hive.metastore.glue.GlueMetastoreStats; +import io.trino.plugin.iceberg.IcebergConnector; +import io.trino.plugin.iceberg.IcebergQueryRunner; +import io.trino.plugin.iceberg.SchemaInitializer; import io.trino.plugin.iceberg.TableType; -import io.trino.plugin.iceberg.TestingIcebergPlugin; import io.trino.plugin.iceberg.catalog.TrinoCatalogFactory; import io.trino.spi.NodeManager; import io.trino.testing.AbstractTestQueryFramework; @@ -38,10 +38,8 @@ import org.testng.annotations.AfterClass; import org.testng.annotations.Test; -import java.io.File; import java.lang.annotation.Retention; import java.lang.annotation.Target; -import java.nio.file.Files; import java.util.Arrays; import java.util.List; import java.util.Map; @@ -50,18 +48,16 @@ import java.util.function.Function; import static com.google.common.base.Preconditions.checkState; -import static com.google.common.base.Verify.verifyNotNull; import static com.google.common.collect.ImmutableMap.toImmutableMap; import static com.google.common.collect.ImmutableMultiset.toImmutableMultiset; import static io.trino.filesystem.TrackingFileSystemFactory.OperationType.INPUT_FILE_NEW_STREAM; -import static io.trino.plugin.hive.HiveTestUtils.HDFS_ENVIRONMENT; -import static io.trino.plugin.hive.HiveTestUtils.HDFS_FILE_SYSTEM_STATS; import static io.trino.plugin.hive.util.MultisetAssertions.assertMultisetsEqual; import static io.trino.plugin.iceberg.IcebergSessionProperties.COLLECT_EXTENDED_STATISTICS_ON_WRITE; import static io.trino.plugin.iceberg.TableType.DATA; import static io.trino.plugin.iceberg.TableType.FILES; import static io.trino.plugin.iceberg.TableType.HISTORY; import static io.trino.plugin.iceberg.TableType.MANIFESTS; +import static io.trino.plugin.iceberg.TableType.MATERIALIZED_VIEW_STORAGE; import static io.trino.plugin.iceberg.TableType.PARTITIONS; import static io.trino.plugin.iceberg.TableType.PROPERTIES; import static io.trino.plugin.iceberg.TableType.REFS; @@ -109,26 +105,13 @@ public class TestIcebergGlueCatalogAccessOperations protected QueryRunner createQueryRunner() throws Exception { - File tmp = Files.createTempDirectory("test_iceberg").toFile(); - DistributedQueryRunner queryRunner = DistributedQueryRunner.builder(testSession) + DistributedQueryRunner queryRunner = IcebergQueryRunner.builder(testSchema) .addCoordinatorProperty("optimizer.experimental-max-prefetched-information-schema-prefixes", Integer.toString(MAX_PREFIXES_COUNT)) + .addIcebergProperty("iceberg.catalog.type", "glue") + .addIcebergProperty("hive.metastore.glue.default-warehouse-dir", "local:///glue") + .setSchemaInitializer(SchemaInitializer.builder().withSchemaName(testSchema).build()) .build(); - - trackingFileSystemFactory = new TrackingFileSystemFactory(new HdfsFileSystemFactory(HDFS_ENVIRONMENT, HDFS_FILE_SYSTEM_STATS)); - - AtomicReference glueStatsReference = new AtomicReference<>(); - queryRunner.installPlugin(new TestingIcebergPlugin( - Optional.empty(), - Optional.of(trackingFileSystemFactory), - new StealStatsModule(glueStatsReference))); - queryRunner.createCatalog("iceberg", "iceberg", - ImmutableMap.of( - "iceberg.catalog.type", "glue", - "hive.metastore.glue.default-warehouse-dir", tmp.getAbsolutePath())); - - queryRunner.execute("CREATE SCHEMA " + testSchema); - - glueStats = verifyNotNull(glueStatsReference.get(), "glueStatsReference not set"); + glueStats = ((IcebergConnector) queryRunner.getCoordinator().getConnector("iceberg")).getInjector().getInstance(GlueMetastoreStats.class); return queryRunner; } @@ -282,7 +265,7 @@ public void testSelectFromMaterializedView() assertGlueMetastoreApiInvocations("SELECT * FROM test_select_mview_view", ImmutableMultiset.builder() - .addCopies(GET_TABLE, 3) + .addCopies(GET_TABLE, 2) .build()); } finally { @@ -300,7 +283,7 @@ public void testSelectFromMaterializedViewWithFilter() assertGlueMetastoreApiInvocations("SELECT * FROM test_select_mview_where_view WHERE age = 2", ImmutableMultiset.builder() - .addCopies(GET_TABLE, 3) + .addCopies(GET_TABLE, 2) .build()); } finally { @@ -318,7 +301,7 @@ public void testRefreshMaterializedView() assertGlueMetastoreApiInvocations("REFRESH MATERIALIZED VIEW test_refresh_mview_view", ImmutableMultiset.builder() - .addCopies(GET_TABLE, 5) + .addCopies(GET_TABLE, 4) .addCopies(UPDATE_TABLE, 1) .build()); } @@ -458,9 +441,12 @@ public void testSelectSystemTable() .addCopies(GET_TABLE, 1) .build()); + assertQueryFails("SELECT * FROM \"test_select_snapshots$materialized_view_storage\"", + "Table '" + testSchema + ".test_select_snapshots\\$materialized_view_storage' not found"); + // This test should get updated if a new system table is added. assertThat(TableType.values()) - .containsExactly(DATA, HISTORY, SNAPSHOTS, MANIFESTS, PARTITIONS, FILES, PROPERTIES, REFS); + .containsExactly(DATA, HISTORY, SNAPSHOTS, MANIFESTS, PARTITIONS, FILES, PROPERTIES, REFS, MATERIALIZED_VIEW_STORAGE); } finally { getQueryRunner().execute("DROP TABLE IF EXISTS test_select_snapshots"); @@ -468,7 +454,7 @@ public void testSelectSystemTable() } @Test - public void testInformationSchemaColumns() + public void testInformationSchemaTableAndColumns() { String schemaName = "test_i_s_columns_schema" + randomNameSuffix(); assertUpdate("CREATE SCHEMA " + schemaName); @@ -493,6 +479,7 @@ public void testInformationSchemaColumns() assertUpdate(session, "CREATE TABLE test_other_select_i_s_columns" + i + "(id varchar, age integer)"); // won't match the filter } + // Bulk columns retrieval assertInvocations( session, "SELECT * FROM information_schema.columns WHERE table_schema = CURRENT_SCHEMA AND table_name LIKE 'test_select_i_s_columns%'", @@ -501,6 +488,38 @@ public void testInformationSchemaColumns() .build(), ImmutableMultiset.of()); } + + // Tables listing + assertInvocations( + session, + "SELECT * FROM information_schema.tables WHERE table_schema = CURRENT_SCHEMA", + ImmutableMultiset.builder() + .add(GET_TABLES) + .build(), + ImmutableMultiset.of()); + + // Pointed columns lookup + assertInvocations( + session, + "SELECT * FROM information_schema.columns WHERE table_schema = CURRENT_SCHEMA AND table_name = 'test_select_i_s_columns0'", + ImmutableMultiset.builder() + .add(GET_TABLE) + .build(), + ImmutableMultiset.builder() + .add(new FileOperation(METADATA_JSON, INPUT_FILE_NEW_STREAM)) + .build()); + + // Pointed columns lookup via DESCRIBE (which does some additional things before delegating to information_schema.columns) + assertInvocations( + session, + "DESCRIBE test_select_i_s_columns0", + ImmutableMultiset.builder() + .add(GET_DATABASE) + .add(GET_TABLE) + .build(), + ImmutableMultiset.builder() + .add(new FileOperation(METADATA_JSON, INPUT_FILE_NEW_STREAM)) + .build()); } finally { for (int i = 0; i < tablesCreated; i++) { @@ -573,6 +592,16 @@ public void testSystemMetadataTableComments() } } + @Test + public void testShowTables() + { + assertGlueMetastoreApiInvocations("SHOW TABLES", + ImmutableMultiset.builder() + .add(GET_DATABASE) + .add(GET_TABLES) + .build()); + } + private void assertGlueMetastoreApiInvocations(@Language("SQL") String query, Multiset expectedInvocations) { assertGlueMetastoreApiInvocations(getSession(), query, expectedInvocations); diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/glue/TestIcebergGlueCatalogConnectorSmokeTest.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/glue/TestIcebergGlueCatalogConnectorSmokeTest.java index ae47487e45d6..f17b1a57f815 100644 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/glue/TestIcebergGlueCatalogConnectorSmokeTest.java +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/glue/TestIcebergGlueCatalogConnectorSmokeTest.java @@ -13,17 +13,6 @@ */ package io.trino.plugin.iceberg.catalog.glue; -import com.amazonaws.services.glue.AWSGlueAsync; -import com.amazonaws.services.glue.AWSGlueAsyncClientBuilder; -import com.amazonaws.services.glue.model.DeleteTableRequest; -import com.amazonaws.services.glue.model.EntityNotFoundException; -import com.amazonaws.services.glue.model.GetTableRequest; -import com.amazonaws.services.s3.AmazonS3; -import com.amazonaws.services.s3.AmazonS3ClientBuilder; -import com.amazonaws.services.s3.model.DeleteObjectsRequest; -import com.amazonaws.services.s3.model.ListObjectsV2Request; -import com.amazonaws.services.s3.model.ListObjectsV2Result; -import com.amazonaws.services.s3.model.S3ObjectSummary; import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; import io.trino.filesystem.Location; @@ -37,7 +26,6 @@ import io.trino.hdfs.HdfsEnvironment; import io.trino.hdfs.TrinoHdfsFileSystemStats; import io.trino.hdfs.authentication.NoHdfsAuthentication; -import io.trino.plugin.hive.aws.AwsApiCallStats; import io.trino.plugin.iceberg.BaseIcebergConnectorSmokeTest; import io.trino.plugin.iceberg.IcebergQueryRunner; import io.trino.plugin.iceberg.SchemaInitializer; @@ -45,12 +33,19 @@ import org.apache.iceberg.FileFormat; import org.testng.annotations.AfterClass; import org.testng.annotations.Test; +import software.amazon.awssdk.services.glue.GlueClient; +import software.amazon.awssdk.services.glue.model.EntityNotFoundException; +import software.amazon.awssdk.services.glue.model.Table; +import software.amazon.awssdk.services.s3.S3Client; +import software.amazon.awssdk.services.s3.model.DeleteObjectsRequest; +import software.amazon.awssdk.services.s3.model.ListObjectsRequest; +import software.amazon.awssdk.services.s3.model.ListObjectsV2Request; +import software.amazon.awssdk.services.s3.model.ObjectIdentifier; +import software.amazon.awssdk.services.s3.model.S3Object; import java.util.List; import static com.google.common.collect.ImmutableList.toImmutableList; -import static io.trino.plugin.hive.metastore.glue.AwsSdkUtil.getPaginatedResults; -import static io.trino.plugin.hive.metastore.glue.converter.GlueToTrinoConverter.getTableParameters; import static io.trino.plugin.iceberg.IcebergTestUtils.checkParquetFileSorting; import static io.trino.testing.TestingConnectorSession.SESSION; import static io.trino.testing.TestingNames.randomNameSuffix; @@ -69,7 +64,7 @@ public class TestIcebergGlueCatalogConnectorSmokeTest { private final String bucketName; private final String schemaName; - private final AWSGlueAsync glueClient; + private final GlueClient glueClient; private final TrinoFileSystemFactory fileSystemFactory; public TestIcebergGlueCatalogConnectorSmokeTest() @@ -77,7 +72,7 @@ public TestIcebergGlueCatalogConnectorSmokeTest() super(FileFormat.PARQUET); this.bucketName = requireNonNull(System.getenv("S3_BUCKET"), "Environment S3_BUCKET was not set"); this.schemaName = "test_iceberg_smoke_" + randomNameSuffix(); - glueClient = AWSGlueAsyncClientBuilder.defaultClient(); + glueClient = GlueClient.create(); HdfsConfigurationInitializer initializer = new HdfsConfigurationInitializer(new HdfsConfig(), ImmutableSet.of()); HdfsConfiguration hdfsConfiguration = new DynamicHdfsConfiguration(initializer, ImmutableSet.of()); @@ -143,53 +138,49 @@ public void testRenameSchema() .hasStackTraceContaining("renameNamespace is not supported for Iceberg Glue catalogs"); } + private Table getGlueTable(String tableName) + { + return glueClient.getTable(x -> x.databaseName(schemaName).name(tableName)).table(); + } + @Override protected void dropTableFromMetastore(String tableName) { - DeleteTableRequest deleteTableRequest = new DeleteTableRequest() - .withDatabaseName(schemaName) - .withName(tableName); - glueClient.deleteTable(deleteTableRequest); - GetTableRequest getTableRequest = new GetTableRequest() - .withDatabaseName(schemaName) - .withName(tableName); - assertThatThrownBy(() -> glueClient.getTable(getTableRequest)) + glueClient.deleteTable(x -> x.databaseName(schemaName).name(tableName)); + assertThatThrownBy(() -> getGlueTable(tableName)) .isInstanceOf(EntityNotFoundException.class); } @Override protected String getMetadataLocation(String tableName) { - GetTableRequest getTableRequest = new GetTableRequest() - .withDatabaseName(schemaName) - .withName(tableName); - return getTableParameters(glueClient.getTable(getTableRequest).getTable()) - .get("metadata_location"); + return getGlueTable(tableName).parameters().get("metadata_location"); } @Override protected void deleteDirectory(String location) { - AmazonS3 s3 = AmazonS3ClientBuilder.standard().build(); - - ListObjectsV2Request listObjectsRequest = new ListObjectsV2Request() - .withBucketName(bucketName) - .withPrefix(location); - List keysToDelete = getPaginatedResults( - s3::listObjectsV2, - listObjectsRequest, - ListObjectsV2Request::setContinuationToken, - ListObjectsV2Result::getNextContinuationToken, - new AwsApiCallStats()) - .map(ListObjectsV2Result::getObjectSummaries) - .flatMap(objectSummaries -> objectSummaries.stream().map(S3ObjectSummary::getKey)) - .map(DeleteObjectsRequest.KeyVersion::new) - .collect(toImmutableList()); - - if (!keysToDelete.isEmpty()) { - s3.deleteObjects(new DeleteObjectsRequest(bucketName).withKeys(keysToDelete)); + try (S3Client s3 = S3Client.create()) { + ListObjectsV2Request listObjectsRequest = ListObjectsV2Request.builder() + .bucket(bucketName) + .prefix(location) + .build(); + s3.listObjectsV2Paginator(listObjectsRequest).stream() + .forEach(listObjectsResponse -> { + List keys = listObjectsResponse.contents().stream().map(S3Object::key).collect(toImmutableList()); + if (!keys.isEmpty()) { + DeleteObjectsRequest deleteObjectsRequest = DeleteObjectsRequest.builder() + .bucket(bucketName) + .delete(builder -> builder.objects(keys.stream() + .map(key -> ObjectIdentifier.builder().key(key).build()) + .toList()).quiet(true)) + .build(); + s3.deleteObjects(deleteObjectsRequest); + } + }); + + assertThat(s3.listObjects(ListObjectsRequest.builder().bucket(bucketName).prefix(location).build()).contents()).isEmpty(); } - assertThat(s3.listObjects(bucketName, location).getObjectSummaries()).isEmpty(); } @Override @@ -208,13 +199,13 @@ protected String schemaPath() @Override protected boolean locationExists(String location) { - String prefix = "s3://" + bucketName + "/"; - AmazonS3 s3 = AmazonS3ClientBuilder.standard().build(); - ListObjectsV2Request request = new ListObjectsV2Request() - .withBucketName(bucketName) - .withPrefix(location.substring(prefix.length())) - .withMaxKeys(1); - return !s3.listObjectsV2(request) - .getObjectSummaries().isEmpty(); + try (S3Client s3 = S3Client.create()) { + ListObjectsV2Request request = ListObjectsV2Request.builder() + .bucket(bucketName) + .prefix(location) + .maxKeys(1) + .build(); + return !s3.listObjectsV2(request).contents().isEmpty(); + } } } diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/glue/TestIcebergGlueCatalogMaterializedView.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/glue/TestIcebergGlueCatalogMaterializedView.java index 95b1d315d7ef..741273f6393e 100644 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/glue/TestIcebergGlueCatalogMaterializedView.java +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/glue/TestIcebergGlueCatalogMaterializedView.java @@ -13,21 +13,16 @@ */ package io.trino.plugin.iceberg.catalog.glue; -import com.amazonaws.services.glue.AWSGlueAsync; -import com.amazonaws.services.glue.AWSGlueAsyncClientBuilder; -import com.amazonaws.services.glue.model.BatchDeleteTableRequest; -import com.amazonaws.services.glue.model.DeleteDatabaseRequest; -import com.amazonaws.services.glue.model.GetTablesRequest; -import com.amazonaws.services.glue.model.GetTablesResult; -import com.amazonaws.services.glue.model.Table; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; -import io.trino.plugin.hive.aws.AwsApiCallStats; import io.trino.plugin.iceberg.BaseIcebergMaterializedViewTest; import io.trino.plugin.iceberg.IcebergQueryRunner; import io.trino.plugin.iceberg.SchemaInitializer; import io.trino.testing.QueryRunner; import org.testng.annotations.AfterClass; +import software.amazon.awssdk.services.glue.GlueClient; +import software.amazon.awssdk.services.glue.model.GetTablesResponse; +import software.amazon.awssdk.services.glue.model.Table; import java.io.File; import java.nio.file.Files; @@ -35,8 +30,8 @@ import java.util.Set; import static com.google.common.collect.ImmutableSet.toImmutableSet; -import static io.trino.plugin.hive.metastore.glue.AwsSdkUtil.getPaginatedResults; import static io.trino.testing.TestingNames.randomNameSuffix; +import static org.apache.iceberg.BaseMetastoreTableOperations.METADATA_LOCATION_PROP; public class TestIcebergGlueCatalogMaterializedView extends BaseIcebergMaterializedViewTest @@ -71,30 +66,36 @@ protected String getSchemaDirectory() return new File(schemaDirectory, schemaName + ".db").getPath(); } + @Override + protected String getStorageMetadataLocation(String materializedViewName) + { + return GlueClient.create() + .getTable(x -> x + .databaseName(schemaName) + .name(materializedViewName)) + .table() + .parameters().get(METADATA_LOCATION_PROP); + } + @AfterClass(alwaysRun = true) public void cleanup() { cleanUpSchema(schemaName); - cleanUpSchema(storageSchemaName); } private static void cleanUpSchema(String schema) { - AWSGlueAsync glueClient = AWSGlueAsyncClientBuilder.defaultClient(); - Set tableNames = getPaginatedResults( - glueClient::getTables, - new GetTablesRequest().withDatabaseName(schema), - GetTablesRequest::setNextToken, - GetTablesResult::getNextToken, - new AwsApiCallStats()) - .map(GetTablesResult::getTableList) + GlueClient glueClient = GlueClient.create(); + Set tableNames = glueClient + .getTablesPaginator(x -> x.databaseName(schema)) + .stream() + .map(GetTablesResponse::tableList) .flatMap(Collection::stream) - .map(Table::getName) + .map(Table::name) .collect(toImmutableSet()); - glueClient.batchDeleteTable(new BatchDeleteTableRequest() - .withDatabaseName(schema) - .withTablesToDelete(tableNames)); - glueClient.deleteDatabase(new DeleteDatabaseRequest() - .withName(schema)); + glueClient.batchDeleteTable(x -> x + .databaseName(schema) + .tablesToDelete(tableNames)); + glueClient.deleteDatabase(x -> x.name(schema)); } } diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/glue/TestIcebergGlueCatalogSkipArchive.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/glue/TestIcebergGlueCatalogSkipArchive.java index b7014bf8e0d2..83d5b274d824 100644 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/glue/TestIcebergGlueCatalogSkipArchive.java +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/glue/TestIcebergGlueCatalogSkipArchive.java @@ -13,17 +13,7 @@ */ package io.trino.plugin.iceberg.catalog.glue; -import com.amazonaws.services.glue.AWSGlueAsync; -import com.amazonaws.services.glue.AWSGlueAsyncClientBuilder; -import com.amazonaws.services.glue.model.GetTableRequest; -import com.amazonaws.services.glue.model.GetTableVersionsRequest; -import com.amazonaws.services.glue.model.GetTableVersionsResult; -import com.amazonaws.services.glue.model.Table; -import com.amazonaws.services.glue.model.TableInput; -import com.amazonaws.services.glue.model.TableVersion; -import com.amazonaws.services.glue.model.UpdateTableRequest; import com.google.common.collect.ImmutableMap; -import io.trino.plugin.hive.aws.AwsApiCallStats; import io.trino.plugin.iceberg.IcebergQueryRunner; import io.trino.plugin.iceberg.SchemaInitializer; import io.trino.plugin.iceberg.fileio.ForwardingFileIo; @@ -35,6 +25,11 @@ import org.apache.iceberg.io.FileIO; import org.testng.annotations.AfterClass; import org.testng.annotations.Test; +import software.amazon.awssdk.services.glue.GlueClient; +import software.amazon.awssdk.services.glue.model.GetTableVersionsResponse; +import software.amazon.awssdk.services.glue.model.Table; +import software.amazon.awssdk.services.glue.model.TableInput; +import software.amazon.awssdk.services.glue.model.TableVersion; import java.io.File; import java.nio.file.Files; @@ -44,10 +39,7 @@ import java.util.Map; import java.util.Optional; -import static com.google.common.collect.ImmutableList.toImmutableList; import static com.google.common.collect.Iterables.getOnlyElement; -import static io.trino.plugin.hive.metastore.glue.AwsSdkUtil.getPaginatedResults; -import static io.trino.plugin.hive.metastore.glue.converter.GlueToTrinoConverter.getTableParameters; import static io.trino.plugin.iceberg.IcebergTestUtils.getFileSystemFactory; import static io.trino.plugin.iceberg.catalog.glue.GlueIcebergUtil.getTableInput; import static io.trino.testing.TestingConnectorSession.SESSION; @@ -65,13 +57,13 @@ public class TestIcebergGlueCatalogSkipArchive extends AbstractTestQueryFramework { private final String schemaName = "test_iceberg_skip_archive_" + randomNameSuffix(); - private AWSGlueAsync glueClient; + private GlueClient glueClient; @Override protected QueryRunner createQueryRunner() throws Exception { - glueClient = AWSGlueAsyncClientBuilder.defaultClient(); + glueClient = GlueClient.create(); File schemaDirectory = Files.createTempDirectory("test_iceberg").toFile(); schemaDirectory.deleteOnExit(); @@ -101,14 +93,14 @@ public void testSkipArchive() try (TestTable table = new TestTable(getQueryRunner()::execute, "test_skip_archive", "(col int)")) { List tableVersionsBeforeInsert = getTableVersions(schemaName, table.getName()); assertThat(tableVersionsBeforeInsert).hasSize(1); - String versionIdBeforeInsert = getOnlyElement(tableVersionsBeforeInsert).getVersionId(); + String versionIdBeforeInsert = getOnlyElement(tableVersionsBeforeInsert).versionId(); assertUpdate("INSERT INTO " + table.getName() + " VALUES 1", 1); // Verify count of table versions isn't increased, but version id is changed List tableVersionsAfterInsert = getTableVersions(schemaName, table.getName()); assertThat(tableVersionsAfterInsert).hasSize(1); - String versionIdAfterInsert = getOnlyElement(tableVersionsAfterInsert).getVersionId(); + String versionIdAfterInsert = getOnlyElement(tableVersionsAfterInsert).versionId(); assertThat(versionIdBeforeInsert).isNotEqualTo(versionIdAfterInsert); } } @@ -122,14 +114,14 @@ public void testNotRemoveExistingArchive() TableVersion initialVersion = getOnlyElement(tableVersionsBeforeInsert); // Add a new archive using Glue client - Table glueTable = glueClient.getTable(new GetTableRequest().withDatabaseName(schemaName).withName(table.getName())).getTable(); - Map tableParameters = new HashMap<>(getTableParameters(glueTable)); + Table glueTable = glueClient.getTable(builder -> builder.databaseName(schemaName).name(table.getName())).table(); + Map tableParameters = new HashMap<>(glueTable.parameters()); String metadataLocation = tableParameters.remove(METADATA_LOCATION_PROP); FileIO io = new ForwardingFileIo(getFileSystemFactory(getDistributedQueryRunner()).create(SESSION)); TableMetadata metadata = TableMetadataParser.read(io, io.newInputFile(metadataLocation)); boolean cacheTableMetadata = new IcebergGlueCatalogConfig().isCacheTableMetadata(); - TableInput tableInput = getTableInput(TESTING_TYPE_MANAGER, table.getName(), Optional.empty(), metadata, metadataLocation, tableParameters, cacheTableMetadata); - glueClient.updateTable(new UpdateTableRequest().withDatabaseName(schemaName).withTableInput(tableInput)); + TableInput tableInput = getTableInput(TESTING_TYPE_MANAGER, table.getName(), Optional.empty(), metadata, metadata.location(), metadataLocation, tableParameters, cacheTableMetadata); + glueClient.updateTable(builder -> builder.databaseName(schemaName).tableInput(tableInput)); assertThat(getTableVersions(schemaName, table.getName())).hasSize(2); assertUpdate("INSERT INTO " + table.getName() + " VALUES 1", 1); @@ -142,14 +134,13 @@ public void testNotRemoveExistingArchive() private List getTableVersions(String databaseName, String tableName) { - return getPaginatedResults( - glueClient::getTableVersions, - new GetTableVersionsRequest().withDatabaseName(databaseName).withTableName(tableName), - GetTableVersionsRequest::setNextToken, - GetTableVersionsResult::getNextToken, - new AwsApiCallStats()) - .map(GetTableVersionsResult::getTableVersions) + return glueClient + .getTableVersionsPaginator(x -> x + .databaseName(databaseName) + .tableName(tableName)) + .stream() + .map(GetTableVersionsResponse::tableVersions) .flatMap(Collection::stream) - .collect(toImmutableList()); + .toList(); } } diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/glue/TestIcebergGlueCreateTableFailure.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/glue/TestIcebergGlueCreateTableFailure.java deleted file mode 100644 index cda4705323cf..000000000000 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/glue/TestIcebergGlueCreateTableFailure.java +++ /dev/null @@ -1,190 +0,0 @@ -/* - * 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.iceberg.catalog.glue; - -import com.amazonaws.services.glue.AWSGlueAsync; -import com.amazonaws.services.glue.model.InvalidInputException; -import com.amazonaws.services.glue.model.OperationTimeoutException; -import com.google.common.collect.ImmutableMap; -import io.airlift.log.Logger; -import io.trino.Session; -import io.trino.filesystem.FileEntry; -import io.trino.filesystem.FileIterator; -import io.trino.filesystem.Location; -import io.trino.filesystem.TrinoFileSystem; -import io.trino.metadata.InternalFunctionBundle; -import io.trino.plugin.hive.metastore.Database; -import io.trino.plugin.hive.metastore.glue.GlueHiveMetastore; -import io.trino.plugin.iceberg.IcebergPlugin; -import io.trino.plugin.iceberg.TestingIcebergConnectorFactory; -import io.trino.spi.security.PrincipalType; -import io.trino.testing.AbstractTestQueryFramework; -import io.trino.testing.LocalQueryRunner; -import org.testng.annotations.AfterClass; -import org.testng.annotations.Test; - -import java.lang.reflect.InvocationTargetException; -import java.nio.file.Files; -import java.nio.file.Path; -import java.util.Optional; -import java.util.concurrent.atomic.AtomicReference; - -import static com.google.common.io.MoreFiles.deleteRecursively; -import static com.google.common.io.RecursiveDeleteOption.ALLOW_INSECURE; -import static com.google.common.reflect.Reflection.newProxy; -import static com.google.inject.util.Modules.EMPTY_MODULE; -import static io.trino.plugin.hive.HiveTestUtils.HDFS_FILE_SYSTEM_FACTORY; -import static io.trino.plugin.hive.metastore.glue.GlueHiveMetastore.createTestingGlueHiveMetastore; -import static io.trino.testing.TestingConnectorSession.SESSION; -import static io.trino.testing.TestingNames.randomNameSuffix; -import static io.trino.testing.TestingSession.testSessionBuilder; -import static org.assertj.core.api.Assertions.assertThat; -import static org.assertj.core.api.Assertions.assertThatThrownBy; - -/* - * The test currently uses AWS Default Credential Provider Chain, - * See https://docs.aws.amazon.com/sdk-for-java/v1/developer-guide/credentials.html#credentials-default - * on ways to set your AWS credentials which will be needed to run this test. - */ -@Test(singleThreaded = true) // testException is a shared mutable state -public class TestIcebergGlueCreateTableFailure - extends AbstractTestQueryFramework -{ - private static final Logger LOG = Logger.get(TestIcebergGlueCreateTableFailure.class); - - private static final String ICEBERG_CATALOG = "iceberg"; - - private final String schemaName = "test_iceberg_glue_" + randomNameSuffix(); - - private Path dataDirectory; - private TrinoFileSystem fileSystem; - private GlueHiveMetastore glueHiveMetastore; - private final AtomicReference testException = new AtomicReference<>(); - - @Override - protected LocalQueryRunner createQueryRunner() - throws Exception - { - Session session = testSessionBuilder() - .setCatalog(ICEBERG_CATALOG) - .setSchema(schemaName) - .build(); - LocalQueryRunner queryRunner = LocalQueryRunner.create(session); - - AWSGlueAsyncAdapterProvider awsGlueAsyncAdapterProvider = delegate -> newProxy(AWSGlueAsync.class, (proxy, method, methodArgs) -> { - Object result; - if (method.getName().equals("createTable")) { - throw testException.get(); - } - try { - result = method.invoke(delegate, methodArgs); - } - catch (InvocationTargetException e) { - throw e.getCause(); - } - return result; - }); - - InternalFunctionBundle.InternalFunctionBundleBuilder functions = InternalFunctionBundle.builder(); - new IcebergPlugin().getFunctions().forEach(functions::functions); - queryRunner.addFunctions(functions.build()); - - queryRunner.createCatalog( - ICEBERG_CATALOG, - new TestingIcebergConnectorFactory(Optional.of(new TestingIcebergGlueCatalogModule(awsGlueAsyncAdapterProvider)), Optional.empty(), EMPTY_MODULE), - ImmutableMap.of()); - - dataDirectory = Files.createTempDirectory("test_iceberg_create_table_failure"); - dataDirectory.toFile().deleteOnExit(); - - glueHiveMetastore = createTestingGlueHiveMetastore(dataDirectory); - fileSystem = HDFS_FILE_SYSTEM_FACTORY.create(SESSION); - - Database database = Database.builder() - .setDatabaseName(schemaName) - .setOwnerName(Optional.of("public")) - .setOwnerType(Optional.of(PrincipalType.ROLE)) - .setLocation(Optional.of(dataDirectory.toString())) - .build(); - glueHiveMetastore.createDatabase(database); - - return queryRunner; - } - - @AfterClass(alwaysRun = true) - public void cleanup() - { - try { - if (glueHiveMetastore != null) { - glueHiveMetastore.dropDatabase(schemaName, false); - } - if (dataDirectory != null) { - deleteRecursively(dataDirectory, ALLOW_INSECURE); - } - } - catch (Exception e) { - LOG.error(e, "Failed to clean up Glue database: %s", schemaName); - } - } - - @Test - public void testCreateTableFailureMetadataCleanedUp() - throws Exception - { - final String exceptionMessage = "Test-simulated metastore invalid input exception"; - testException.set(new InvalidInputException(exceptionMessage)); - testCreateTableFailure(exceptionMessage, false); - } - - @Test - public void testCreateTableFailureMetadataNotCleanedUp() - throws Exception - { - final String exceptionMessage = "Test-simulated metastore operation timeout exception"; - testException.set(new OperationTimeoutException(exceptionMessage)); - testCreateTableFailure(exceptionMessage, true); - } - - private void testCreateTableFailure(String expectedExceptionMessage, boolean shouldMetadataFileExist) - throws Exception - { - String tableName = "test_create_failure_" + randomNameSuffix(); - assertThatThrownBy(() -> getQueryRunner().execute("CREATE TABLE " + tableName + " (a_varchar) AS VALUES ('Trino')")) - .hasMessageContaining(expectedExceptionMessage); - - assertMetadataLocation(tableName, shouldMetadataFileExist); - } - - protected void assertMetadataLocation(String tableName, boolean shouldMetadataFileExist) - throws Exception - { - FileIterator fileIterator = fileSystem.listFiles(Location.of(dataDirectory.toString())); - String tableLocationPrefix = Path.of(dataDirectory.toString(), tableName).toString(); - boolean metadataFileFound = false; - while (fileIterator.hasNext()) { - FileEntry fileEntry = fileIterator.next(); - String location = fileEntry.location().toString(); - if (location.startsWith(tableLocationPrefix) && location.endsWith(".metadata.json")) { - metadataFileFound = true; - break; - } - } - if (shouldMetadataFileExist) { - assertThat(metadataFileFound).as("Metadata file should exist").isTrue(); - } - else { - assertThat(metadataFileFound).as("Metadata file should not exist").isFalse(); - } - } -} diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/glue/TestIcebergGlueTableOperationsInsertFailure.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/glue/TestIcebergGlueTableOperationsInsertFailure.java index 27e320f1bbb8..02f9c4ba4ee3 100644 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/glue/TestIcebergGlueTableOperationsInsertFailure.java +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/glue/TestIcebergGlueTableOperationsInsertFailure.java @@ -13,30 +13,34 @@ */ package io.trino.plugin.iceberg.catalog.glue; -import com.amazonaws.services.glue.AWSGlueAsync; import com.google.common.collect.ImmutableMap; +import com.google.inject.Binder; +import com.google.inject.multibindings.ProvidesIntoSet; +import io.airlift.configuration.AbstractConfigurationAwareModule; import io.airlift.log.Logger; import io.trino.Session; -import io.trino.metadata.InternalFunctionBundle; import io.trino.plugin.hive.metastore.Database; +import io.trino.plugin.hive.metastore.glue.ForGlueHiveMetastore; import io.trino.plugin.hive.metastore.glue.GlueHiveMetastore; -import io.trino.plugin.iceberg.IcebergPlugin; -import io.trino.plugin.iceberg.TestingIcebergConnectorFactory; +import io.trino.plugin.iceberg.TestingIcebergPlugin; +import io.trino.plugin.iceberg.catalog.IcebergCatalogModule; import io.trino.spi.security.PrincipalType; import io.trino.testing.AbstractTestQueryFramework; -import io.trino.testing.LocalQueryRunner; +import io.trino.testing.QueryRunner; +import io.trino.testing.StandaloneQueryRunner; import org.apache.iceberg.exceptions.CommitStateUnknownException; import org.testng.annotations.AfterClass; import org.testng.annotations.Test; +import software.amazon.awssdk.core.interceptor.Context; +import software.amazon.awssdk.core.interceptor.ExecutionAttributes; +import software.amazon.awssdk.core.interceptor.ExecutionInterceptor; +import software.amazon.awssdk.services.glue.model.UpdateTableRequest; -import java.lang.reflect.InvocationTargetException; import java.nio.file.Files; import java.nio.file.Path; import java.util.Optional; -import static com.google.common.reflect.Reflection.newProxy; -import static com.google.inject.util.Modules.EMPTY_MODULE; -import static io.trino.plugin.hive.metastore.glue.GlueHiveMetastore.createTestingGlueHiveMetastore; +import static io.trino.plugin.hive.metastore.glue.TestingGlueHiveMetastore.createTestingGlueHiveMetastore; import static io.trino.testing.TestingNames.randomNameSuffix; import static io.trino.testing.TestingSession.testSessionBuilder; import static java.lang.String.format; @@ -59,42 +63,25 @@ public class TestIcebergGlueTableOperationsInsertFailure private GlueHiveMetastore glueHiveMetastore; @Override - protected LocalQueryRunner createQueryRunner() + protected QueryRunner createQueryRunner() throws Exception { Session session = testSessionBuilder() .setCatalog(ICEBERG_CATALOG) .setSchema(schemaName) .build(); - LocalQueryRunner queryRunner = LocalQueryRunner.create(session); - - AWSGlueAsyncAdapterProvider awsGlueAsyncAdapterProvider = delegate -> newProxy(AWSGlueAsync.class, (proxy, method, methodArgs) -> { - Object result; - try { - result = method.invoke(delegate, methodArgs); - } - catch (InvocationTargetException e) { - throw e.getCause(); - } - if (method.getName().equals("updateTable")) { - throw new RuntimeException("Test-simulated Glue timeout exception"); - } - return result; - }); - - InternalFunctionBundle.InternalFunctionBundleBuilder functions = InternalFunctionBundle.builder(); - new IcebergPlugin().getFunctions().forEach(functions::functions); - queryRunner.addFunctions(functions.build()); - - queryRunner.createCatalog( - ICEBERG_CATALOG, - new TestingIcebergConnectorFactory(Optional.of(new TestingIcebergGlueCatalogModule(awsGlueAsyncAdapterProvider)), Optional.empty(), EMPTY_MODULE), - ImmutableMap.of()); + QueryRunner queryRunner = new StandaloneQueryRunner(session); Path dataDirectory = Files.createTempDirectory("iceberg_data"); dataDirectory.toFile().deleteOnExit(); - glueHiveMetastore = createTestingGlueHiveMetastore(dataDirectory); + queryRunner.installPlugin(new TestingIcebergPlugin(dataDirectory, Optional.of(new TestingGlueCatalogModule()))); + queryRunner.createCatalog(ICEBERG_CATALOG, "iceberg", ImmutableMap.builder() + .put("iceberg.catalog.type", "glue") + .put("fs.hadoop.enabled", "true") + .buildOrThrow()); + + glueHiveMetastore = createTestingGlueHiveMetastore(dataDirectory, this::closeAfterClass); Database database = Database.builder() .setDatabaseName(schemaName) @@ -132,4 +119,30 @@ public void testInsertFailureDoesNotCorruptTheTableMetadata() .hasMessageContaining("Test-simulated Glue timeout exception"); assertQuery("SELECT * FROM " + tableName, "VALUES 'Trino', 'rocks'"); } + + private static class TestingGlueCatalogModule + extends AbstractConfigurationAwareModule + { + @Override + protected void setup(Binder binder) + { + install(new IcebergCatalogModule()); + } + + @ProvidesIntoSet + @ForGlueHiveMetastore + public ExecutionInterceptor createExecutionInterceptor() + { + return new ExecutionInterceptor() + { + @Override + public void afterExecution(Context.AfterExecution context, ExecutionAttributes executionAttributes) + { + if (context.request() instanceof UpdateTableRequest) { + throw new RuntimeException("Test-simulated Glue timeout exception"); + } + } + }; + } + } } diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/glue/TestIcebergS3AndGlueMetastoreTest.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/glue/TestIcebergS3AndGlueMetastoreTest.java index cac690815439..73f8548f2b06 100644 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/glue/TestIcebergS3AndGlueMetastoreTest.java +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/glue/TestIcebergS3AndGlueMetastoreTest.java @@ -20,11 +20,11 @@ import io.trino.testing.QueryRunner; import org.testng.annotations.Test; -import java.nio.file.Path; +import java.net.URI; import java.util.Set; import java.util.stream.Collectors; -import static io.trino.plugin.hive.metastore.glue.GlueHiveMetastore.createTestingGlueHiveMetastore; +import static io.trino.plugin.hive.metastore.glue.TestingGlueHiveMetastore.createTestingGlueHiveMetastore; import static io.trino.testing.TestingNames.randomNameSuffix; import static java.util.Objects.requireNonNull; import static org.assertj.core.api.Assertions.assertThat; @@ -42,7 +42,7 @@ public TestIcebergS3AndGlueMetastoreTest() protected QueryRunner createQueryRunner() throws Exception { - metastore = createTestingGlueHiveMetastore(Path.of(schemaPath())); + metastore = createTestingGlueHiveMetastore(URI.create(schemaPath()), this::closeAfterClass); DistributedQueryRunner queryRunner = IcebergQueryRunner.builder() .setIcebergProperties(ImmutableMap.builder() .put("iceberg.catalog.type", "glue") diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/glue/TestSharedGlueMetastore.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/glue/TestSharedGlueMetastore.java index d8684dfb1f9c..9a6eef306d29 100644 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/glue/TestSharedGlueMetastore.java +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/glue/TestSharedGlueMetastore.java @@ -28,8 +28,9 @@ import org.testng.annotations.AfterClass; import java.nio.file.Path; +import java.util.Optional; -import static io.trino.plugin.hive.metastore.glue.GlueHiveMetastore.createTestingGlueHiveMetastore; +import static io.trino.plugin.hive.metastore.glue.TestingGlueHiveMetastore.createTestingGlueHiveMetastore; import static io.trino.plugin.iceberg.IcebergQueryRunner.ICEBERG_CATALOG; import static io.trino.plugin.tpch.TpchMetadata.TINY_SCHEMA_NAME; import static io.trino.testing.QueryAssertions.copyTpchTables; @@ -87,8 +88,8 @@ protected QueryRunner createQueryRunner() "hive.metastore.glue.default-warehouse-dir", dataDirectory.toString(), "iceberg.hive-catalog-name", "hive")); - this.glueMetastore = createTestingGlueHiveMetastore(dataDirectory); - queryRunner.installPlugin(new TestingHivePlugin(glueMetastore)); + this.glueMetastore = createTestingGlueHiveMetastore(dataDirectory, this::closeAfterClass); + queryRunner.installPlugin(new TestingHivePlugin(dataDirectory, Optional.of(glueMetastore))); queryRunner.createCatalog(HIVE_CATALOG, "hive"); queryRunner.createCatalog( "hive_with_redirections", diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/glue/TestTrinoGlueCatalog.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/glue/TestTrinoGlueCatalog.java index fd53d7a41bc2..36fd62048bb5 100644 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/glue/TestTrinoGlueCatalog.java +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/glue/TestTrinoGlueCatalog.java @@ -13,11 +13,6 @@ */ package io.trino.plugin.iceberg.catalog.glue; -import com.amazonaws.services.glue.AWSGlueAsync; -import com.amazonaws.services.glue.AWSGlueAsyncClientBuilder; -import com.amazonaws.services.glue.model.CreateDatabaseRequest; -import com.amazonaws.services.glue.model.DatabaseInput; -import com.amazonaws.services.glue.model.DeleteDatabaseRequest; import com.google.common.collect.ImmutableMap; import io.airlift.log.Logger; import io.trino.filesystem.TrinoFileSystemFactory; @@ -25,6 +20,7 @@ import io.trino.plugin.hive.NodeVersion; import io.trino.plugin.hive.metastore.glue.GlueMetastoreStats; import io.trino.plugin.iceberg.CommitTaskData; +import io.trino.plugin.iceberg.IcebergConfig; import io.trino.plugin.iceberg.IcebergMetadata; import io.trino.plugin.iceberg.TableStatisticsWriter; import io.trino.plugin.iceberg.catalog.BaseTrinoCatalogTest; @@ -36,6 +32,7 @@ import io.trino.spi.security.TrinoPrincipal; import io.trino.spi.type.TestingTypeManager; import org.testng.annotations.Test; +import software.amazon.awssdk.services.glue.GlueClient; import java.io.File; import java.io.IOException; @@ -43,6 +40,8 @@ import java.nio.file.Path; import java.util.Optional; +import static com.google.common.util.concurrent.MoreExecutors.directExecutor; +import static com.google.common.util.concurrent.MoreExecutors.newDirectExecutorService; import static io.airlift.json.JsonCodec.jsonCodec; import static io.trino.plugin.hive.HiveTestUtils.HDFS_FILE_SYSTEM_FACTORY; import static io.trino.sql.planner.TestingPlannerContext.PLANNER_CONTEXT; @@ -60,9 +59,14 @@ public class TestTrinoGlueCatalog @Override protected TrinoCatalog createTrinoCatalog(boolean useUniqueTableLocations) + { + return createTrinoCatalog(useUniqueTableLocations, false); + } + + private TrinoCatalog createTrinoCatalog(boolean useUniqueTableLocations, boolean useSystemSecurity) { TrinoFileSystemFactory fileSystemFactory = HDFS_FILE_SYSTEM_FACTORY; - AWSGlueAsync glueClient = AWSGlueAsyncClientBuilder.defaultClient(); + GlueClient glueClient = GlueClient.create(); IcebergGlueCatalogConfig catalogConfig = new IcebergGlueCatalogConfig(); return new TrinoGlueCatalog( new CatalogName("catalog_name"), @@ -78,8 +82,11 @@ protected TrinoCatalog createTrinoCatalog(boolean useUniqueTableLocations) "test", glueClient, new GlueMetastoreStats(), + useSystemSecurity, Optional.empty(), - useUniqueTableLocations); + useUniqueTableLocations, + new IcebergConfig().isHideMaterializedViewStorageTable(), + directExecutor()); } /** @@ -92,11 +99,11 @@ public void testNonLowercaseGlueDatabase() // Trino schema names are always lowercase (until https://github.com/trinodb/trino/issues/17) String trinoSchemaName = databaseName.toLowerCase(ENGLISH); - AWSGlueAsync glueClient = AWSGlueAsyncClientBuilder.defaultClient(); - glueClient.createDatabase(new CreateDatabaseRequest() - .withDatabaseInput(new DatabaseInput() + GlueClient glueClient = GlueClient.create(); + glueClient.createDatabase(database -> database + .databaseInput(input -> input // Currently this is actually stored in lowercase - .withName(databaseName))); + .name(databaseName))); try { TrinoCatalog catalog = createTrinoCatalog(false); assertThat(catalog.namespaceExists(SESSION, databaseName)).as("catalog.namespaceExists(databaseName)") @@ -114,10 +121,15 @@ public void testNonLowercaseGlueDatabase() CatalogHandle.fromId("iceberg:NORMAL:v12345"), jsonCodec(CommitTaskData.class), catalog, - connectorIdentity -> { + (connectorIdentity, fileIoProperties) -> { throw new UnsupportedOperationException(); }, - new TableStatisticsWriter(new NodeVersion("test-version"))); + new TableStatisticsWriter(new NodeVersion("test-version")), + Optional.empty(), + false, + ignore -> false, + newDirectExecutorService(), + directExecutor()); assertThat(icebergMetadata.schemaExists(SESSION, databaseName)).as("icebergMetadata.schemaExists(databaseName)") .isFalse(); assertThat(icebergMetadata.schemaExists(SESSION, trinoSchemaName)).as("icebergMetadata.schemaExists(trinoSchemaName)") @@ -127,8 +139,7 @@ public void testNonLowercaseGlueDatabase() .contains(trinoSchemaName); } finally { - glueClient.deleteDatabase(new DeleteDatabaseRequest() - .withName(databaseName)); + glueClient.deleteDatabase(delete -> delete.name(databaseName)); } } @@ -140,7 +151,7 @@ public void testDefaultLocation() tmpDirectory.toFile().deleteOnExit(); TrinoFileSystemFactory fileSystemFactory = HDFS_FILE_SYSTEM_FACTORY; - AWSGlueAsync glueClient = AWSGlueAsyncClientBuilder.defaultClient(); + GlueClient glueClient = GlueClient.create(); IcebergGlueCatalogConfig catalogConfig = new IcebergGlueCatalogConfig(); TrinoCatalog catalogWithDefaultLocation = new TrinoGlueCatalog( new CatalogName("catalog_name"), @@ -156,8 +167,11 @@ public void testDefaultLocation() "test", glueClient, new GlueMetastoreStats(), + false, Optional.of(tmpDirectory.toAbsolutePath().toString()), - false); + false, + new IcebergConfig().isHideMaterializedViewStorageTable(), + directExecutor()); String namespace = "test_default_location_" + randomNameSuffix(); String table = "tableName"; diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/glue/TestingGlueIcebergTableOperationsProvider.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/glue/TestingGlueIcebergTableOperationsProvider.java deleted file mode 100644 index 84c9b6390dd3..000000000000 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/glue/TestingGlueIcebergTableOperationsProvider.java +++ /dev/null @@ -1,87 +0,0 @@ -/* - * 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.iceberg.catalog.glue; - -import com.amazonaws.auth.AWSCredentialsProvider; -import com.amazonaws.services.glue.AWSGlueAsync; -import com.google.common.collect.ImmutableSet; -import com.google.inject.Inject; -import io.trino.filesystem.TrinoFileSystemFactory; -import io.trino.plugin.hive.metastore.glue.GlueHiveMetastoreConfig; -import io.trino.plugin.hive.metastore.glue.GlueMetastoreStats; -import io.trino.plugin.iceberg.catalog.IcebergTableOperations; -import io.trino.plugin.iceberg.catalog.IcebergTableOperationsProvider; -import io.trino.plugin.iceberg.catalog.TrinoCatalog; -import io.trino.plugin.iceberg.fileio.ForwardingFileIo; -import io.trino.spi.connector.ConnectorSession; -import io.trino.spi.type.TypeManager; - -import java.util.Optional; - -import static io.trino.plugin.hive.metastore.glue.GlueClientUtil.createAsyncGlueClient; -import static java.util.Objects.requireNonNull; - -public class TestingGlueIcebergTableOperationsProvider - implements IcebergTableOperationsProvider -{ - private final TypeManager typeManager; - private final boolean cacheTableMetadata; - private final TrinoFileSystemFactory fileSystemFactory; - private final AWSGlueAsync glueClient; - private final GlueMetastoreStats stats; - - @Inject - public TestingGlueIcebergTableOperationsProvider( - TypeManager typeManager, - IcebergGlueCatalogConfig catalogConfig, - TrinoFileSystemFactory fileSystemFactory, - GlueMetastoreStats stats, - GlueHiveMetastoreConfig glueConfig, - AWSCredentialsProvider credentialsProvider, - AWSGlueAsyncAdapterProvider awsGlueAsyncAdapterProvider) - { - this.typeManager = requireNonNull(typeManager, "typeManager is null"); - this.cacheTableMetadata = catalogConfig.isCacheTableMetadata(); - this.fileSystemFactory = requireNonNull(fileSystemFactory, "fileSystemFactory is null"); - this.stats = requireNonNull(stats, "stats is null"); - requireNonNull(glueConfig, "glueConfig is null"); - requireNonNull(credentialsProvider, "credentialsProvider is null"); - requireNonNull(awsGlueAsyncAdapterProvider, "awsGlueAsyncAdapterProvider is null"); - this.glueClient = awsGlueAsyncAdapterProvider.createAWSGlueAsyncAdapter( - createAsyncGlueClient(glueConfig, credentialsProvider, ImmutableSet.of(), stats.newRequestMetricsCollector())); - } - - @Override - public IcebergTableOperations createTableOperations( - TrinoCatalog catalog, - ConnectorSession session, - String database, - String table, - Optional owner, - Optional location) - { - return new GlueIcebergTableOperations( - typeManager, - cacheTableMetadata, - glueClient, - stats, - ((TrinoGlueCatalog) catalog)::getTable, - new ForwardingFileIo(fileSystemFactory.create(session)), - session, - database, - table, - owner, - location); - } -} diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/glue/TestingIcebergGlueCatalogModule.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/glue/TestingIcebergGlueCatalogModule.java deleted file mode 100644 index 46493fea3191..000000000000 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/glue/TestingIcebergGlueCatalogModule.java +++ /dev/null @@ -1,76 +0,0 @@ -/* - * 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.iceberg.catalog.glue; - -import com.amazonaws.auth.AWSCredentialsProvider; -import com.amazonaws.handlers.RequestHandler2; -import com.amazonaws.services.glue.model.Table; -import com.google.inject.Binder; -import com.google.inject.Key; -import com.google.inject.Scopes; -import com.google.inject.TypeLiteral; -import io.airlift.configuration.AbstractConfigurationAwareModule; -import io.trino.plugin.hive.HideDeltaLakeTables; -import io.trino.plugin.hive.metastore.glue.ForGlueHiveMetastore; -import io.trino.plugin.hive.metastore.glue.GlueCredentialsProvider; -import io.trino.plugin.hive.metastore.glue.GlueHiveMetastoreConfig; -import io.trino.plugin.hive.metastore.glue.GlueMetastoreModule; -import io.trino.plugin.hive.metastore.glue.GlueMetastoreStats; -import io.trino.plugin.iceberg.catalog.IcebergTableOperationsProvider; -import io.trino.plugin.iceberg.catalog.TrinoCatalogFactory; - -import java.util.function.Predicate; - -import static com.google.inject.multibindings.Multibinder.newSetBinder; -import static com.google.inject.multibindings.OptionalBinder.newOptionalBinder; -import static io.airlift.configuration.ConditionalModule.conditionalModule; -import static io.airlift.configuration.ConfigBinder.configBinder; -import static java.util.Objects.requireNonNull; -import static org.weakref.jmx.guice.ExportBinder.newExporter; - -public class TestingIcebergGlueCatalogModule - extends AbstractConfigurationAwareModule -{ - private final AWSGlueAsyncAdapterProvider awsGlueAsyncAdapterProvider; - - public TestingIcebergGlueCatalogModule(AWSGlueAsyncAdapterProvider awsGlueAsyncAdapterProvider) - { - this.awsGlueAsyncAdapterProvider = requireNonNull(awsGlueAsyncAdapterProvider, "awsGlueAsyncAdapterProvider is null"); - } - - @Override - protected void setup(Binder binder) - { - configBinder(binder).bindConfig(GlueHiveMetastoreConfig.class); - configBinder(binder).bindConfig(IcebergGlueCatalogConfig.class); - binder.bind(GlueMetastoreStats.class).in(Scopes.SINGLETON); - newExporter(binder).export(GlueMetastoreStats.class).withGeneratedName(); - binder.bind(AWSCredentialsProvider.class).toProvider(GlueCredentialsProvider.class).in(Scopes.SINGLETON); - binder.bind(IcebergTableOperationsProvider.class).to(TestingGlueIcebergTableOperationsProvider.class).in(Scopes.SINGLETON); - binder.bind(TrinoCatalogFactory.class).to(TrinoGlueCatalogFactory.class).in(Scopes.SINGLETON); - newExporter(binder).export(TrinoCatalogFactory.class).withGeneratedName(); - binder.bind(AWSGlueAsyncAdapterProvider.class).toInstance(awsGlueAsyncAdapterProvider); - - install(conditionalModule( - IcebergGlueCatalogConfig.class, - IcebergGlueCatalogConfig::isSkipArchive, - internalBinder -> newSetBinder(internalBinder, RequestHandler2.class, ForGlueHiveMetastore.class).addBinding().toInstance(new SkipArchiveRequestHandler()))); - - // Required to inject HiveMetastoreFactory for migrate procedure - binder.bind(Key.get(boolean.class, HideDeltaLakeTables.class)).toInstance(false); - newOptionalBinder(binder, Key.get(new TypeLiteral>() {}, ForGlueHiveMetastore.class)) - .setBinding().toInstance(table -> true); - install(new GlueMetastoreModule()); - } -} diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/hms/TestIcebergHiveMetastoreTableOperationsReleaseLockFailure.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/hms/TestIcebergHiveMetastoreTableOperationsReleaseLockFailure.java deleted file mode 100644 index 73c018d781eb..000000000000 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/hms/TestIcebergHiveMetastoreTableOperationsReleaseLockFailure.java +++ /dev/null @@ -1,140 +0,0 @@ -/* - * 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.iceberg.catalog.hms; - -import com.google.common.collect.ImmutableMap; -import io.trino.Session; -import io.trino.hive.thrift.metastore.Table; -import io.trino.metadata.InternalFunctionBundle; -import io.trino.plugin.hive.metastore.AcidTransactionOwner; -import io.trino.plugin.hive.metastore.Database; -import io.trino.plugin.hive.metastore.HiveMetastore; -import io.trino.plugin.hive.metastore.thrift.BridgingHiveMetastore; -import io.trino.plugin.hive.metastore.thrift.InMemoryThriftMetastore; -import io.trino.plugin.hive.metastore.thrift.ThriftMetastore; -import io.trino.plugin.hive.metastore.thrift.ThriftMetastoreConfig; -import io.trino.plugin.hive.metastore.thrift.ThriftMetastoreFactory; -import io.trino.plugin.iceberg.IcebergPlugin; -import io.trino.plugin.iceberg.TestingIcebergConnectorFactory; -import io.trino.spi.security.ConnectorIdentity; -import io.trino.spi.security.PrincipalType; -import io.trino.testing.AbstractTestQueryFramework; -import io.trino.testing.LocalQueryRunner; -import org.testng.annotations.Test; - -import java.io.File; -import java.nio.file.Files; -import java.util.Optional; - -import static com.google.inject.util.Modules.EMPTY_MODULE; -import static io.trino.testing.TestingSession.testSessionBuilder; -import static java.lang.String.format; - -public class TestIcebergHiveMetastoreTableOperationsReleaseLockFailure - extends AbstractTestQueryFramework -{ - private static final String ICEBERG_CATALOG = "iceberg"; - private static final String SCHEMA_NAME = "test_schema"; - private File baseDir; - - @Override - protected LocalQueryRunner createQueryRunner() - throws Exception - { - Session session = testSessionBuilder() - .setCatalog(ICEBERG_CATALOG) - .setSchema(SCHEMA_NAME) - .build(); - - baseDir = Files.createTempDirectory(null).toFile(); - baseDir.deleteOnExit(); - - LocalQueryRunner queryRunner = LocalQueryRunner.create(session); - - InternalFunctionBundle.InternalFunctionBundleBuilder functions = InternalFunctionBundle.builder(); - new IcebergPlugin().getFunctions().forEach(functions::functions); - queryRunner.addFunctions(functions.build()); - - ThriftMetastore thriftMetastore = createMetastoreWithReleaseLockFailure(); - HiveMetastore hiveMetastore = new BridgingHiveMetastore(thriftMetastore); - TestingIcebergHiveMetastoreCatalogModule testModule = new TestingIcebergHiveMetastoreCatalogModule(hiveMetastore, buildThriftMetastoreFactory(thriftMetastore)); - - queryRunner.createCatalog( - ICEBERG_CATALOG, - new TestingIcebergConnectorFactory(Optional.of(testModule), Optional.empty(), EMPTY_MODULE), - ImmutableMap.of()); - - Database database = Database.builder() - .setDatabaseName(SCHEMA_NAME) - .setOwnerName(Optional.of("public")) - .setOwnerType(Optional.of(PrincipalType.ROLE)) - .build(); - hiveMetastore.createDatabase(database); - - return queryRunner; - } - - @Test - public void testReleaseLockFailureDoesNotCorruptTheTable() - { - String tableName = "test_release_lock_failure"; - query(format("CREATE TABLE %s (a_varchar) AS VALUES ('Trino')", tableName)); - query(format("INSERT INTO %s VALUES 'rocks'", tableName)); - assertQuery("SELECT * FROM " + tableName, "VALUES 'Trino', 'rocks'"); - } - - private InMemoryThriftMetastore createMetastoreWithReleaseLockFailure() - { - return new InMemoryThriftMetastore(new File(baseDir + "/metastore"), new ThriftMetastoreConfig()) { - @Override - public long acquireTableExclusiveLock(AcidTransactionOwner transactionOwner, String queryId, String dbName, String tableName) - { - // returning dummy lock - return 100; - } - - @Override - public void releaseTableLock(long lockId) - { - throw new RuntimeException("Release table lock has failed!"); - } - - @Override - public synchronized void createTable(Table table) - { - // InMemoryThriftMetastore throws an exception if the table has any privileges set - table.setPrivileges(null); - super.createTable(table); - } - }; - } - - private static ThriftMetastoreFactory buildThriftMetastoreFactory(ThriftMetastore thriftMetastore) - { - return new ThriftMetastoreFactory() - { - @Override - public boolean isImpersonationEnabled() - { - return false; - } - - @Override - public ThriftMetastore createMetastore(Optional identity) - { - return thriftMetastore; - } - }; - } -} diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/hms/TestTrinoHiveCatalogWithHiveMetastore.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/hms/TestTrinoHiveCatalogWithHiveMetastore.java index 620e53812137..4dee4b886f2f 100644 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/hms/TestTrinoHiveCatalogWithHiveMetastore.java +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/hms/TestTrinoHiveCatalogWithHiveMetastore.java @@ -26,6 +26,7 @@ import io.trino.hdfs.s3.HiveS3Config; import io.trino.hdfs.s3.TrinoS3ConfigurationInitializer; import io.trino.plugin.base.CatalogName; +import io.trino.plugin.base.util.AutoCloseableCloser; import io.trino.plugin.hive.TrinoViewHiveMetastore; import io.trino.plugin.hive.containers.HiveMinioDataLake; import io.trino.plugin.hive.metastore.cache.CachingHiveMetastore; @@ -46,9 +47,10 @@ import java.util.Set; import static com.google.common.base.Verify.verify; +import static com.google.common.util.concurrent.MoreExecutors.directExecutor; import static io.trino.plugin.hive.TestingThriftHiveMetastoreBuilder.testingThriftHiveMetastoreBuilder; import static io.trino.plugin.hive.containers.HiveHadoop.HIVE3_IMAGE; -import static io.trino.plugin.hive.metastore.cache.CachingHiveMetastore.memoizeMetastore; +import static io.trino.plugin.hive.metastore.cache.CachingHiveMetastore.createPerTransactionCache; import static io.trino.testing.TestingNames.randomNameSuffix; import static io.trino.testing.containers.Minio.MINIO_ACCESS_KEY; import static io.trino.testing.containers.Minio.MINIO_SECRET_KEY; @@ -59,6 +61,7 @@ public class TestTrinoHiveCatalogWithHiveMetastore { private static final String bucketName = "test-hive-catalog-with-hms-" + randomNameSuffix(); + private final AutoCloseableCloser closer = AutoCloseableCloser.create(); // Use MinIO for storage, since HDFS is hard to get working in a unit test private HiveMinioDataLake dataLake; @@ -77,6 +80,7 @@ public void tearDown() dataLake.stop(); dataLake = null; } + closer.close(); } @Override @@ -99,10 +103,10 @@ protected TrinoCatalog createTrinoCatalog(boolean useUniqueTableLocations) ThriftMetastore thriftMetastore = testingThriftHiveMetastoreBuilder() .thriftMetastoreConfig(new ThriftMetastoreConfig() // Read timed out sometimes happens with the default timeout - .setMetastoreTimeout(new Duration(1, MINUTES))) + .setReadTimeout(new Duration(1, MINUTES))) .metastoreClient(dataLake.getHiveHadoop().getHiveMetastoreEndpoint()) - .build(); - CachingHiveMetastore metastore = memoizeMetastore(new BridgingHiveMetastore(thriftMetastore), 1000); + .build(closer::register); + CachingHiveMetastore metastore = createPerTransactionCache(new BridgingHiveMetastore(thriftMetastore), 1000); return new TrinoHiveCatalog( new CatalogName("catalog"), metastore, @@ -123,10 +127,17 @@ public ThriftMetastore createMetastore(Optional identity) { return thriftMetastore; } - }), + }, new IcebergHiveCatalogConfig()), useUniqueTableLocations, false, - false); + false, + isHideMaterializedViewStorageTable(), + directExecutor()); + } + + protected boolean isHideMaterializedViewStorageTable() + { + return true; } @Override diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/nessie/TestTrinoNessieCatalog.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/nessie/TestTrinoNessieCatalog.java index 45030e2f09e5..086330c9b26f 100644 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/nessie/TestTrinoNessieCatalog.java +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/nessie/TestTrinoNessieCatalog.java @@ -31,8 +31,8 @@ import io.trino.spi.security.TrinoPrincipal; import io.trino.spi.type.TestingTypeManager; import org.apache.iceberg.nessie.NessieIcebergClient; +import org.projectnessie.client.NessieClientBuilder; import org.projectnessie.client.api.NessieApiV1; -import org.projectnessie.client.http.HttpClientBuilder; import org.testng.annotations.AfterClass; import org.testng.annotations.BeforeClass; import org.testng.annotations.Test; @@ -42,7 +42,10 @@ import java.net.URI; import java.nio.file.Path; import java.util.Map; +import java.util.Optional; +import static com.google.common.util.concurrent.MoreExecutors.directExecutor; +import static com.google.common.util.concurrent.MoreExecutors.newDirectExecutorService; import static io.airlift.json.JsonCodec.jsonCodec; import static io.trino.plugin.hive.HiveTestUtils.HDFS_ENVIRONMENT; import static io.trino.plugin.hive.HiveTestUtils.HDFS_FILE_SYSTEM_STATS; @@ -88,7 +91,7 @@ protected TrinoCatalog createTrinoCatalog(boolean useUniqueTableLocations) TrinoFileSystemFactory fileSystemFactory = new HdfsFileSystemFactory(HDFS_ENVIRONMENT, HDFS_FILE_SYSTEM_STATS); IcebergNessieCatalogConfig icebergNessieCatalogConfig = new IcebergNessieCatalogConfig() .setServerUri(URI.create(nessieContainer.getRestApiUri())); - NessieApiV1 nessieApi = HttpClientBuilder.builder() + NessieApiV1 nessieApi = NessieClientBuilder.createClientBuilder("test", null) .withUri(nessieContainer.getRestApiUri()) .build(NessieApiV1.class); NessieIcebergClient nessieClient = new NessieIcebergClient(nessieApi, icebergNessieCatalogConfig.getDefaultReferenceName(), null, ImmutableMap.of()); @@ -112,7 +115,7 @@ public void testDefaultLocation() IcebergNessieCatalogConfig icebergNessieCatalogConfig = new IcebergNessieCatalogConfig() .setDefaultWarehouseDir(tmpDirectory.toAbsolutePath().toString()) .setServerUri(URI.create(nessieContainer.getRestApiUri())); - NessieApiV1 nessieApi = HttpClientBuilder.builder() + NessieApiV1 nessieApi = NessieClientBuilder.createClientBuilder("http", null) .withUri(nessieContainer.getRestApiUri()) .build(NessieApiV1.class); NessieIcebergClient nessieClient = new NessieIcebergClient(nessieApi, icebergNessieCatalogConfig.getDefaultReferenceName(), null, ImmutableMap.of()); @@ -176,10 +179,15 @@ public void testNonLowercaseNamespace() CatalogHandle.fromId("iceberg:NORMAL:v12345"), jsonCodec(CommitTaskData.class), catalog, - connectorIdentity -> { + (connectorIdentity, fileIoProperties) -> { throw new UnsupportedOperationException(); }, - new TableStatisticsWriter(new NodeVersion("test-version"))); + new TableStatisticsWriter(new NodeVersion("test-version")), + Optional.empty(), + false, + ignore -> false, + newDirectExecutorService(), + directExecutor()); assertThat(icebergMetadata.schemaExists(SESSION, namespace)).as("icebergMetadata.schemaExists(namespace)") .isTrue(); assertThat(icebergMetadata.schemaExists(SESSION, schema)).as("icebergMetadata.schemaExists(schema)") diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/rest/RestCatalogTestUtils.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/rest/RestCatalogTestUtils.java index cc0961b86804..23760a69cf5e 100644 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/rest/RestCatalogTestUtils.java +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/rest/RestCatalogTestUtils.java @@ -24,25 +24,25 @@ import io.trino.hdfs.authentication.NoHdfsAuthentication; import io.trino.spi.connector.ConnectorSession; import io.trino.testing.TestingConnectorSession; -import org.apache.hadoop.fs.Path; import org.apache.iceberg.CatalogProperties; import org.apache.iceberg.catalog.Catalog; import org.apache.iceberg.jdbc.JdbcCatalog; -import org.assertj.core.util.Files; import java.io.File; +import java.nio.file.Path; public final class RestCatalogTestUtils { private RestCatalogTestUtils() {} - public static Catalog backendCatalog(File warehouseLocation) + public static Catalog backendCatalog(Path warehouseLocation) { ImmutableMap.Builder properties = ImmutableMap.builder(); - properties.put(CatalogProperties.URI, "jdbc:h2:file:" + Files.newTemporaryFile().getAbsolutePath()); + properties.put(CatalogProperties.URI, "jdbc:h2:file:" + createTempFile(null, null).toAbsolutePath()); properties.put(JdbcCatalog.PROPERTY_PREFIX + "username", "user"); properties.put(JdbcCatalog.PROPERTY_PREFIX + "password", "password"); - properties.put(CatalogProperties.WAREHOUSE_LOCATION, warehouseLocation.toPath().resolve("iceberg_data").toFile().getAbsolutePath()); + properties.put(JdbcCatalog.PROPERTY_PREFIX + "schema-version", "V1"); + properties.put(CatalogProperties.WAREHOUSE_LOCATION, warehouseLocation.resolve("iceberg_data").toFile().getAbsolutePath()); ConnectorSession connectorSession = TestingConnectorSession.builder().build(); HdfsConfig hdfsConfig = new HdfsConfig(); @@ -51,9 +51,19 @@ public static Catalog backendCatalog(File warehouseLocation) HdfsContext context = new HdfsContext(connectorSession); JdbcCatalog catalog = new JdbcCatalog(); - catalog.setConf(hdfsEnvironment.getConfiguration(context, new Path(warehouseLocation.getAbsolutePath()))); + catalog.setConf(hdfsEnvironment.getConfiguration(context, new org.apache.hadoop.fs.Path(warehouseLocation.toAbsolutePath().toString()))); catalog.initialize("backend_jdbc", properties.buildOrThrow()); return catalog; } + + private static Path createTempFile(String prefix, String suffix) + { + try { + return File.createTempFile(prefix == null ? "data" : prefix, suffix == null ? ".tmp" : suffix).toPath(); + } + catch (Exception e) { + throw new RuntimeException("Failed to create temp file", e); + } + } } diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/rest/TestIcebergRestCatalogCaseInsensitiveMapping.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/rest/TestIcebergRestCatalogCaseInsensitiveMapping.java new file mode 100644 index 000000000000..81196dd58590 --- /dev/null +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/rest/TestIcebergRestCatalogCaseInsensitiveMapping.java @@ -0,0 +1,273 @@ +/* + * 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.iceberg.catalog.rest; + +import io.airlift.http.server.testing.TestingHttpServer; +import io.trino.Session; +import io.trino.plugin.iceberg.IcebergQueryRunner; +import io.trino.testing.AbstractTestQueryFramework; +import io.trino.testing.QueryRunner; +import org.apache.iceberg.Schema; +import org.apache.iceberg.catalog.Namespace; +import org.apache.iceberg.catalog.TableIdentifier; +import org.apache.iceberg.jdbc.JdbcCatalog; +import org.apache.iceberg.rest.DelegatingRestSessionCatalog; +import org.apache.iceberg.types.Types; +import org.apache.iceberg.view.ViewBuilder; +import org.testng.annotations.BeforeMethod; +import org.testng.annotations.Test; + +import java.io.IOException; +import java.io.UncheckedIOException; +import java.net.URI; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.util.Map; +import java.util.Optional; +import java.util.Set; + +import static com.google.common.io.MoreFiles.deleteRecursively; +import static com.google.common.io.RecursiveDeleteOption.ALLOW_INSECURE; +import static io.trino.plugin.iceberg.IcebergSchemaProperties.LOCATION_PROPERTY; +import static io.trino.plugin.iceberg.catalog.rest.RestCatalogTestUtils.backendCatalog; +import static io.trino.testing.TestingNames.randomNameSuffix; +import static java.nio.file.Files.createDirectories; +import static java.util.Locale.ENGLISH; +import static org.apache.iceberg.types.Types.NestedField.required; +import static org.assertj.core.api.Assertions.assertThat; + +public final class TestIcebergRestCatalogCaseInsensitiveMapping + extends AbstractTestQueryFramework +{ + private static final String SCHEMA = "LeVeL1_" + randomNameSuffix(); + private static final String LOWERCASE_SCHEMA = SCHEMA.toLowerCase(ENGLISH); + private static final Namespace NAMESPACE = Namespace.of(SCHEMA); + + private JdbcCatalog backend; + + @Override + protected QueryRunner createQueryRunner() + throws Exception + { + Path warehouseLocation = Files.createTempDirectory(null); + closeAfterClass(() -> deleteRecursively(warehouseLocation, ALLOW_INSECURE)); + + backend = closeAfterClass((JdbcCatalog) backendCatalog(warehouseLocation)); + + DelegatingRestSessionCatalog delegatingCatalog = DelegatingRestSessionCatalog.builder() + .delegate(backend) + .build(); + + TestingHttpServer testServer = delegatingCatalog.testServer(); + testServer.start(); + closeAfterClass(testServer::stop); + + return IcebergQueryRunner.builder(LOWERCASE_SCHEMA) + .setBaseDataDir(Optional.of(warehouseLocation)) + .addIcebergProperty("iceberg.catalog.type", "rest") + .addIcebergProperty("iceberg.rest-catalog.uri", testServer.getBaseUrl().toString()) + .addIcebergProperty("iceberg.rest-catalog.case-insensitive-name-matching", "true") + .addIcebergProperty("iceberg.register-table-procedure.enabled", "true") + .build(); + } + + @BeforeMethod + public void setup() + { + if (backend.namespaceExists(NAMESPACE)) { + return; + } + backend.createNamespace(NAMESPACE); + assertThat(computeActual("SHOW SCHEMAS").getOnlyColumnAsSet()) + .containsExactlyInAnyOrder( + "information_schema", + "tpch", + LOWERCASE_SCHEMA); + + assertThat(computeActual("SHOW SCHEMAS LIKE 'level%'").getOnlyColumnAsSet()) + .containsExactlyInAnyOrder( + LOWERCASE_SCHEMA); + + assertQuery("SELECT * FROM information_schema.schemata", + """ + VALUES + ('iceberg', 'information_schema'), + ('iceberg', '%s'), + ('iceberg', 'tpch') + """.formatted(LOWERCASE_SCHEMA)); + } + + @Test + public void testCaseInsensitiveMatchingForTable() + { + Map namespaceMetadata = backend.loadNamespaceMetadata(NAMESPACE); + String namespaceLocation = namespaceMetadata.get(LOCATION_PROPERTY); + createDir(namespaceLocation); + + // Create and query a mixed case letter table from Trino + String tableName1 = "MiXed_CaSe_TaBlE1_" + randomNameSuffix(); + String lowercaseTableName1 = tableName1.toLowerCase(ENGLISH); + String table1Location = namespaceLocation + "/" + lowercaseTableName1; + assertUpdate("CREATE TABLE " + tableName1 + " WITH (location = '" + table1Location + "') AS SELECT BIGINT '42' a, DOUBLE '-38.5' b", 1); + assertQuery("SELECT * FROM " + tableName1, "VALUES (42, -38.5)"); + + // Create a mixed case letter table directly using rest catalog and query from Trino + String tableName2 = "mIxEd_cAsE_tAbLe2_" + randomNameSuffix(); + String lowercaseTableName2 = tableName2.toLowerCase(ENGLISH); + String table2Location = namespaceLocation + "/" + lowercaseTableName2; + createDir(table2Location); + createDir(table2Location + "/data"); + createDir(table2Location + "/metadata"); + backend + .buildTable(TableIdentifier.of(NAMESPACE, tableName2), new Schema(required(1, "x", Types.LongType.get()))) + .withLocation(table2Location) + .createTransaction() + .commitTransaction(); + assertUpdate("INSERT INTO " + tableName2 + " VALUES (78)", 1); + assertQuery("SELECT * FROM " + tableName2, "VALUES (78)"); + + // Test register/unregister table. Re-register for further testing. + assertThat(backend.dropTable(TableIdentifier.of(NAMESPACE, lowercaseTableName1), false)).isTrue(); + assertQueryFails("SELECT * FROM " + tableName1, ".*'iceberg.%s.%s' does not exist".formatted(LOWERCASE_SCHEMA, lowercaseTableName1)); + assertUpdate("CALL system.register_table (CURRENT_SCHEMA, '" + tableName1 + "', '" + table1Location + "')"); + assertQuery("SELECT * FROM " + tableName1, "VALUES (42, -38.5)"); + assertUpdate("CALL system.unregister_table (CURRENT_SCHEMA, '" + tableName1 + "')"); + assertQueryFails("SELECT * FROM " + tableName1, ".*'iceberg.%s.%s' does not exist".formatted(LOWERCASE_SCHEMA, lowercaseTableName1)); + assertUpdate("CALL system.register_table (CURRENT_SCHEMA, '" + tableName1 + "', '" + table1Location + "')"); + + // Query information_schema and list objects + assertThat(computeActual("SHOW TABLES IN " + SCHEMA).getOnlyColumnAsSet()).contains(lowercaseTableName1, lowercaseTableName2); + assertThat(computeActual("SHOW TABLES IN " + SCHEMA + " LIKE 'mixed_case_table%'").getOnlyColumnAsSet()).isEqualTo(Set.of(lowercaseTableName1, lowercaseTableName2)); + assertQuery("SELECT * FROM information_schema.tables WHERE table_schema != 'information_schema' AND table_type = 'BASE TABLE'", + """ + VALUES + ('iceberg', '%1$s', '%2$s', 'BASE TABLE'), + ('iceberg', '%1$s', '%3$s', 'BASE TABLE') + """.formatted(LOWERCASE_SCHEMA, lowercaseTableName1, lowercaseTableName2)); + + // Add table comment + assertUpdate("COMMENT ON TABLE " + tableName1 + " IS 'test comment' "); + assertThat(getTableComment(lowercaseTableName1)).isEqualTo("test comment"); + + // Add table column comment + assertUpdate("COMMENT ON COLUMN " + tableName1 + ".a IS 'test column comment'"); + assertThat(getColumnComment(lowercaseTableName1, "a")).isEqualTo("test column comment"); + + // Rename table + String renamedTableName1 = tableName1 + "_renamed"; + assertUpdate("ALTER TABLE " + lowercaseTableName1 + " RENAME TO " + renamedTableName1); + assertQueryFails("SELECT * FROM " + tableName1, ".*'iceberg.%s.%s' does not exist".formatted(LOWERCASE_SCHEMA, lowercaseTableName1)); + assertQuery("SELECT * FROM " + renamedTableName1, "VALUES (42, -38.5)"); + + // Drop tables + assertUpdate("DROP TABLE " + renamedTableName1); + assertUpdate("DROP TABLE " + tableName2); + + // Query dropped tablesd + assertQueryFails("SELECT * FROM " + renamedTableName1, ".*'iceberg.%s.%s' does not exist".formatted(LOWERCASE_SCHEMA, renamedTableName1.toLowerCase(ENGLISH))); + assertQueryFails("SELECT * FROM " + tableName2, ".*'iceberg.%s.%s' does not exist".formatted(LOWERCASE_SCHEMA, lowercaseTableName2)); + } + + @Test + public void testCaseInsensitiveMatchingForView() + { + Map namespaceMetadata = backend.loadNamespaceMetadata(NAMESPACE); + String namespaceLocation = namespaceMetadata.get(LOCATION_PROPERTY); + createDir(namespaceLocation); + + // Create and query a mixed case letter view from Trino + String viewName1 = "MiXed_CaSe_vIeW1_" + randomNameSuffix(); + String lowercaseViewName1 = viewName1.toLowerCase(ENGLISH); + assertUpdate("CREATE VIEW " + viewName1 + " AS SELECT BIGINT '25' a, DOUBLE '99.4' b"); + assertQuery("SELECT * FROM " + viewName1, "VALUES (25, 99.4)"); + + // Create a mixed case letter view directly using rest catalog and query from Trino + String viewName2 = "mIxEd_cAsE_ViEw2_" + randomNameSuffix(); + String lowercaseViewName2 = viewName2.toLowerCase(ENGLISH); + String view2Location = namespaceLocation + "/" + lowercaseViewName2; + createDir(view2Location); + createDir(view2Location + "/data"); + createDir(view2Location + "/metadata"); + ViewBuilder viewBuilder = backend.buildView(TableIdentifier.of(NAMESPACE, viewName2)); + viewBuilder + .withQuery("trino", "SELECT BIGINT '34' y") + .withSchema(new Schema(required(1, "y", Types.LongType.get()))) + .withDefaultNamespace(NAMESPACE) + .withLocation(view2Location) + .createOrReplace(); + assertQuery("SELECT * FROM " + viewName2, "VALUES (34)"); + + // Query information_schema and list objects + assertThat(computeActual("SHOW TABLES IN " + SCHEMA).getOnlyColumnAsSet()).contains(lowercaseViewName1, lowercaseViewName2); + assertThat(computeActual("SHOW TABLES IN " + SCHEMA + " LIKE 'mixed_case_view%'").getOnlyColumnAsSet()).contains(lowercaseViewName1, lowercaseViewName2); + assertQuery("SELECT * FROM information_schema.tables WHERE table_schema != 'information_schema' AND table_type = 'VIEW'", + """ + VALUES + ('iceberg', '%1$s', '%2$s', 'VIEW'), + ('iceberg', '%1$s', '%3$s', 'VIEW') + """.formatted(LOWERCASE_SCHEMA, lowercaseViewName1, lowercaseViewName2)); + + // Add view comment + assertUpdate("COMMENT ON VIEW " + viewName1 + " IS 'test comment' "); + assertThat(getTableComment(lowercaseViewName1)).isEqualTo("test comment"); + + // Add view column comment + assertUpdate("COMMENT ON COLUMN " + viewName1 + ".a IS 'test column comment'"); + assertThat(getColumnComment(lowercaseViewName1, "a")).isEqualTo("test column comment"); + + // Rename view + String renamedViewName1 = viewName1 + "_renamed"; + assertUpdate("ALTER VIEW " + lowercaseViewName1 + " RENAME TO " + renamedViewName1); + assertQueryFails("SELECT * FROM " + viewName1, ".*'iceberg.%s.%s' does not exist".formatted(LOWERCASE_SCHEMA, lowercaseViewName1)); + assertQuery("SELECT * FROM " + renamedViewName1, "VALUES (25, 99.4)"); + + // Drop views + assertUpdate("DROP VIEW " + renamedViewName1); + assertUpdate("DROP VIEW " + viewName2); + + // Query dropped views + assertQueryFails("SELECT * FROM " + renamedViewName1, ".*'iceberg.%s.%s' does not exist".formatted(LOWERCASE_SCHEMA, renamedViewName1.toLowerCase(ENGLISH))); + assertQueryFails("SELECT * FROM " + viewName2, ".*'iceberg.%s.%s' does not exist".formatted(LOWERCASE_SCHEMA, lowercaseViewName2)); + } + + private String getColumnComment(String tableName, String columnName) + { + return (String) computeScalar("SELECT comment FROM information_schema.columns " + + "WHERE table_schema = '" + LOWERCASE_SCHEMA + "' AND table_name = '" + tableName + "' AND column_name = '" + columnName + "'"); + } + + private static void createDir(String absoluteDirPath) + { + Path path = Paths.get(URI.create(absoluteDirPath).getPath()); + try { + createDirectories(path); + } + catch (IOException e) { + throw new UncheckedIOException("Cannot create %s directory".formatted(absoluteDirPath), e); + } + } + + protected String getTableComment(String tableName) + { + Session session = getSession(); + return getTableComment(session.getCatalog().orElseThrow(), session.getSchema().orElseThrow(), tableName); + } + + protected String getTableComment(String catalogName, String schemaName, String tableName) + { + return (String) computeScalar("SELECT comment FROM system.metadata.table_comments " + + "WHERE catalog_name = '" + catalogName + "' AND schema_name = '" + schemaName + "' AND table_name = '" + tableName + "'"); + } +} diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/rest/TestIcebergRestCatalogConfig.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/rest/TestIcebergRestCatalogConfig.java index 5e411999f304..1a757c764bfd 100644 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/rest/TestIcebergRestCatalogConfig.java +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/rest/TestIcebergRestCatalogConfig.java @@ -14,6 +14,7 @@ package io.trino.plugin.iceberg.catalog.rest; import com.google.common.collect.ImmutableMap; +import io.airlift.units.Duration; import org.testng.annotations.Test; import java.util.Map; @@ -21,6 +22,7 @@ import static io.airlift.configuration.testing.ConfigAssertions.assertFullMapping; import static io.airlift.configuration.testing.ConfigAssertions.assertRecordedDefaults; import static io.airlift.configuration.testing.ConfigAssertions.recordDefaults; +import static java.util.concurrent.TimeUnit.MINUTES; public class TestIcebergRestCatalogConfig { @@ -29,9 +31,16 @@ public void testDefaults() { assertRecordedDefaults(recordDefaults(IcebergRestCatalogConfig.class) .setBaseUri(null) + .setPrefix(null) .setWarehouse(null) + .setNestedNamespaceEnabled(false) .setSessionType(IcebergRestCatalogConfig.SessionType.NONE) - .setSecurity(IcebergRestCatalogConfig.Security.NONE)); + .setSecurity(IcebergRestCatalogConfig.Security.NONE) + .setVendedCredentialsEnabled(false) + .setViewEndpointsEnabled(true) + .setSigV4Enabled(false) + .setCaseInsensitiveNameMatching(false) + .setCaseInsensitiveNameMatchingCacheTtl(new Duration(1, MINUTES))); } @Test @@ -39,16 +48,30 @@ public void testExplicitPropertyMappings() { Map properties = ImmutableMap.builder() .put("iceberg.rest-catalog.uri", "http://localhost:1234") + .put("iceberg.rest-catalog.prefix", "dev") .put("iceberg.rest-catalog.warehouse", "test_warehouse_identifier") + .put("iceberg.rest-catalog.nested-namespace-enabled", "true") .put("iceberg.rest-catalog.security", "OAUTH2") .put("iceberg.rest-catalog.session", "USER") + .put("iceberg.rest-catalog.vended-credentials-enabled", "true") + .put("iceberg.rest-catalog.view-endpoints-enabled", "false") + .put("iceberg.rest-catalog.sigv4-enabled", "true") + .put("iceberg.rest-catalog.case-insensitive-name-matching", "true") + .put("iceberg.rest-catalog.case-insensitive-name-matching.cache-ttl", "3m") .buildOrThrow(); IcebergRestCatalogConfig expected = new IcebergRestCatalogConfig() .setBaseUri("http://localhost:1234") + .setPrefix("dev") .setWarehouse("test_warehouse_identifier") + .setNestedNamespaceEnabled(true) .setSessionType(IcebergRestCatalogConfig.SessionType.USER) - .setSecurity(IcebergRestCatalogConfig.Security.OAUTH2); + .setSecurity(IcebergRestCatalogConfig.Security.OAUTH2) + .setVendedCredentialsEnabled(true) + .setViewEndpointsEnabled(false) + .setSigV4Enabled(true) + .setCaseInsensitiveNameMatching(true) + .setCaseInsensitiveNameMatchingCacheTtl(new Duration(3, MINUTES)); assertFullMapping(properties, expected); } diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/rest/TestIcebergRestCatalogSigV4Config.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/rest/TestIcebergRestCatalogSigV4Config.java new file mode 100644 index 000000000000..2ea5254c5147 --- /dev/null +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/rest/TestIcebergRestCatalogSigV4Config.java @@ -0,0 +1,46 @@ +/* + * 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.iceberg.catalog.rest; + +import com.google.common.collect.ImmutableMap; +import org.junit.jupiter.api.Test; + +import java.util.Map; + +import static io.airlift.configuration.testing.ConfigAssertions.assertFullMapping; +import static io.airlift.configuration.testing.ConfigAssertions.assertRecordedDefaults; +import static io.airlift.configuration.testing.ConfigAssertions.recordDefaults; + +final class TestIcebergRestCatalogSigV4Config +{ + @Test + void testDefaults() + { + assertRecordedDefaults(recordDefaults(IcebergRestCatalogSigV4Config.class) + .setSigningName("execute-api")); + } + + @Test + void testExplicitPropertyMappings() + { + Map properties = ImmutableMap.builder() + .put("iceberg.rest-catalog.signing-name", "glue") + .buildOrThrow(); + + IcebergRestCatalogSigV4Config expected = new IcebergRestCatalogSigV4Config() + .setSigningName("glue"); + + assertFullMapping(properties, expected); + } +} diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/rest/TestIcebergS3TablesConnectorSmokeTest.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/rest/TestIcebergS3TablesConnectorSmokeTest.java new file mode 100644 index 000000000000..676b9476314f --- /dev/null +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/rest/TestIcebergS3TablesConnectorSmokeTest.java @@ -0,0 +1,299 @@ +/* + * 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.iceberg.catalog.rest; + +import io.trino.filesystem.Location; +import io.trino.plugin.iceberg.BaseIcebergConnectorSmokeTest; +import io.trino.plugin.iceberg.IcebergConfig; +import io.trino.plugin.iceberg.IcebergConnector; +import io.trino.plugin.iceberg.IcebergQueryRunner; +import io.trino.plugin.iceberg.catalog.TrinoCatalog; +import io.trino.plugin.iceberg.catalog.TrinoCatalogFactory; +import io.trino.spi.connector.SchemaTableName; +import io.trino.testing.DistributedQueryRunner; +import io.trino.testing.QueryRunner; +import io.trino.testing.TestingConnectorBehavior; +import org.apache.iceberg.BaseTable; +import org.testng.annotations.Test; + +import static io.trino.testing.TestingNames.randomNameSuffix; +import static java.lang.String.format; +import static java.util.Objects.requireNonNull; +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +public final class TestIcebergS3TablesConnectorSmokeTest + extends BaseIcebergConnectorSmokeTest +{ + public static final String S3_TABLES_BUCKET = requireEnv("S3_TABLES_BUCKET"); + public static final String AWS_ACCESS_KEY_ID = System.getenv("AWS_ACCESS_KEY_ID"); + public static final String AWS_SECRET_ACCESS_KEY = System.getenv("AWS_SECRET_ACCESS_KEY"); + public static final String AWS_REGION = requireEnv("AWS_REGION"); + public static final String AWS_IAM_ROLE = requireEnv("AWS_IAM_ROLE"); + + private static String requireEnv(String key) + { + return requireNonNull(System.getenv(key), "Environment variable '%s' is not set".formatted(key)); + } + + public TestIcebergS3TablesConnectorSmokeTest() + { + super(new IcebergConfig().getFileFormat().toIceberg()); + } + + @Override + protected boolean hasBehavior(TestingConnectorBehavior connectorBehavior) + { + return switch (connectorBehavior) { + case SUPPORTS_CREATE_MATERIALIZED_VIEW, + SUPPORTS_RENAME_MATERIALIZED_VIEW, + SUPPORTS_RENAME_SCHEMA, + SUPPORTS_RENAME_TABLE -> false; + default -> super.hasBehavior(connectorBehavior); + }; + } + + @Override + protected QueryRunner createQueryRunner() + throws Exception + { + return IcebergQueryRunner.builder("tpch") + .addIcebergProperty("iceberg.file-format", format.name()) + .addIcebergProperty("iceberg.register-table-procedure.enabled", "true") + .addIcebergProperty("iceberg.catalog.type", "rest") + .addIcebergProperty("iceberg.rest-catalog.uri", "https://glue.%s.amazonaws.com/iceberg".formatted(AWS_REGION)) + .addIcebergProperty("iceberg.rest-catalog.warehouse", "s3tablescatalog/" + S3_TABLES_BUCKET) + .addIcebergProperty("iceberg.rest-catalog.view-endpoints-enabled", "false") + .addIcebergProperty("iceberg.rest-catalog.sigv4-enabled", "true") + .addIcebergProperty("iceberg.rest-catalog.signing-name", "glue") + .addIcebergProperty("iceberg.writer-sort-buffer-size", "1MB") + .addIcebergProperty("iceberg.allowed-extra-properties", "write.metadata.delete-after-commit.enabled,write.metadata.previous-versions-max") + .addIcebergProperty("fs.native-s3.enabled", "true") + .addIcebergProperty("s3.iam-role", AWS_IAM_ROLE) + .addIcebergProperty("s3.region", AWS_REGION) + .addIcebergProperty("s3.aws-access-key", AWS_ACCESS_KEY_ID) + .addIcebergProperty("s3.aws-secret-key", AWS_SECRET_ACCESS_KEY) + .setInitialTables(REQUIRED_TPCH_TABLES) + .build(); + } + + @Override + protected String getMetadataLocation(String tableName) + { + TrinoCatalogFactory catalogFactory = ((IcebergConnector) ((DistributedQueryRunner) getQueryRunner()).getCoordinator().getConnector("iceberg")).getInjector().getInstance(TrinoCatalogFactory.class); + TrinoCatalog trinoCatalog = catalogFactory.create(getSession().getIdentity().toConnectorIdentity()); + BaseTable table = (BaseTable) trinoCatalog.loadTable(getSession().toConnectorSession(), new SchemaTableName(getSession().getSchema().orElseThrow(), tableName)); + return table.operations().current().metadataFileLocation(); + } + + @Override + protected String schemaPath() + { + return "dummy"; + } + + @Override + protected boolean locationExists(String location) + { + throw new UnsupportedOperationException(); + } + + @Override + protected boolean isFileSorted(Location path, String sortColumnName) + { + throw new UnsupportedOperationException(); + } + + @Override + protected void deleteDirectory(String location) + { + throw new UnsupportedOperationException(); + } + + @Override + protected void dropTableFromMetastore(String tableName) + { + throw new UnsupportedOperationException(); + } + + @Test + @Override // Override because the location pattern differs + public void testShowCreateTable() + { + assertThat((String) computeScalar("SHOW CREATE TABLE region")) + .matches("CREATE TABLE iceberg.tpch.region \\(\n" + + " regionkey bigint,\n" + + " name varchar,\n" + + " comment varchar\n" + + "\\)\n" + + "WITH \\(\n" + + " format = 'PARQUET',\n" + + " format_version = 2,\n" + + " location = 's3://.*--table-s3'\n" + + " max_commit_retry = 4\n" + + "\\)"); + } + + @Test + @Override + public void testRenameSchema() + { + assertThatThrownBy(super::testRenameSchema) + .hasMessageContaining("renameNamespace is not supported for Iceberg REST catalog"); + } + + @Test + @Override + public void testMaterializedView() + { + assertThatThrownBy(super::testMaterializedView) + .hasMessageContaining("createMaterializedView is not supported for Iceberg REST catalog"); + } + + @Test + @Override // Override because S3 Tables does not support specifying the location + public void testCreateTableWithTrailingSpaceInLocation() + { + String tableName = "test_create_table_with_trailing_space_" + randomNameSuffix(); + String tableLocationWithTrailingSpace = schemaPath() + tableName + " "; + + assertQueryFails( + format("CREATE TABLE %s WITH (location = '%s') AS SELECT 1 AS a, 'INDIA' AS b, true AS c", tableName, tableLocationWithTrailingSpace), + "Failed to create transaction"); + } + + @Test + @Override + public void testRenameTable() + { + assertThatThrownBy(super::testRenameTable) + .hasStackTraceContaining("Unable to process: RenameTable endpoint is not supported for Glue Catalog"); + } + + @Test + @Override + public void testView() + { + assertThatThrownBy(super::testView) + .hasMessageContaining("Server does not support endpoint: POST /v1/{prefix}/namespaces/{namespace}/views"); + } + + @Test + @Override + public void testCommentViewColumn() + { + assertThatThrownBy(super::testCommentViewColumn) + .hasMessageContaining("Server does not support endpoint: POST /v1/{prefix}/namespaces/{namespace}/views"); + } + + @Test + @Override + public void testCommentView() + { + assertThatThrownBy(super::testCommentView) + .hasMessageContaining("Server does not support endpoint: POST /v1/{prefix}/namespaces/{namespace}/views"); + } + + @Test + @Override // The locationExists helper method is unsupported + public void testCreateTableWithNonExistingSchemaVerifyLocation() {} + + @Test + @Override // The TrinoFileSystem.deleteFile is unsupported + public void testDropTableWithMissingMetadataFile() {} + + @Test + @Override // The TrinoFileSystem.deleteFile is unsupported + public void testDropTableWithMissingManifestListFile() {} + + @Test + //@Override // The TrinoFileSystem.listFiles is unsupported + public void testMetadataDeleteAfterCommitEnabled() {} + + @Test + @Override // The TrinoFileSystem.deleteFile is unsupported + public void testDropTableWithMissingSnapshotFile() {} + + @Test + @Override // The TrinoFileSystem.listFiles is unsupported + public void testDropTableWithMissingDataFile() {} + + @Test + @Override // The TrinoFileSystem.deleteDirectory is unsupported + public void testDropTableWithNonExistentTableLocation() {} + + @Test + @Override // BaseIcebergConnectorSmokeTest.isFileSorted method is unsupported + public void testSortedNationTable() {} + + @Test + @Override // The TrinoFileSystem.deleteFile is unsupported + public void testFileSortingWithLargerTable() {} + + @Test + @Override // The procedure is unsupported in S3 Tables + public void testRegisterTableWithTableLocation() {} + + @Test + @Override // The procedure is unsupported in S3 Tables + public void testRegisterTableWithComments() {} + + @Test + @Override // The procedure is unsupported in S3 Tables + public void testRegisterTableWithShowCreateTable() {} + + @Test + @Override // The procedure is unsupported in S3 Tables + public void testRegisterTableWithReInsert() {} + + @Test + @Override // The procedure is unsupported in S3 Tables + public void testRegisterTableWithDroppedTable() {} + + @Test + @Override // The procedure is unsupported in S3 Tables + public void testRegisterTableWithDifferentTableName() {} + + @Test + @Override // The procedure is unsupported in S3 Tables + public void testRegisterTableWithMetadataFile() {} + + @Test + @Override // The procedure is unsupported in S3 Tables + public void testRegisterTableWithTrailingSpaceInLocation() {} + + @Test + @Override // The procedure is unsupported in S3 Tables + public void testUnregisterTable() {} + + @Test + @Override // The procedure is unsupported in S3 Tables + public void testUnregisterBrokenTable() {} + + @Test + @Override // The procedure is unsupported in S3 Tables + public void testUnregisterTableNotExistingSchema() {} + + @Test + @Override // The procedure is unsupported in S3 Tables + public void testUnregisterTableNotExistingTable() {} + + @Test + @Override // The procedure is unsupported in S3 Tables + public void testRepeatUnregisterTable() {} + + @Test + @Override // The procedure is unsupported in S3 Tables + public void testUnregisterTableAccessControl() {} +} diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/rest/TestIcebergTrinoRestCatalogConnectorSmokeTest.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/rest/TestIcebergTrinoRestCatalogConnectorSmokeTest.java index 860bfe7f1470..ce5e80c4c8f5 100644 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/rest/TestIcebergTrinoRestCatalogConnectorSmokeTest.java +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/rest/TestIcebergTrinoRestCatalogConnectorSmokeTest.java @@ -19,17 +19,22 @@ import io.trino.plugin.iceberg.BaseIcebergConnectorSmokeTest; import io.trino.plugin.iceberg.IcebergConfig; import io.trino.plugin.iceberg.IcebergQueryRunner; +import io.trino.testing.QueryFailedException; import io.trino.testing.QueryRunner; import io.trino.testing.TestingConnectorBehavior; import org.apache.iceberg.BaseTable; -import org.apache.iceberg.catalog.Catalog; +import org.apache.iceberg.Schema; +import org.apache.iceberg.catalog.Namespace; import org.apache.iceberg.catalog.TableIdentifier; import org.apache.iceberg.jdbc.JdbcCatalog; import org.apache.iceberg.rest.DelegatingRestSessionCatalog; +import org.apache.iceberg.types.Types; import org.assertj.core.util.Files; import org.testng.annotations.AfterClass; +import org.testng.annotations.Test; -import java.io.File; +import java.io.IOException; +import java.io.UncheckedIOException; import java.nio.file.Path; import java.util.Optional; @@ -38,15 +43,18 @@ import static io.trino.plugin.iceberg.IcebergTestUtils.checkOrcFileSorting; import static io.trino.plugin.iceberg.IcebergTestUtils.checkParquetFileSorting; import static io.trino.plugin.iceberg.catalog.rest.RestCatalogTestUtils.backendCatalog; +import static io.trino.testing.TestingNames.randomNameSuffix; import static java.lang.String.format; import static org.apache.iceberg.FileFormat.PARQUET; +import static org.apache.iceberg.types.Types.NestedField.required; +import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatThrownBy; public class TestIcebergTrinoRestCatalogConnectorSmokeTest extends BaseIcebergConnectorSmokeTest { - private File warehouseLocation; - private Catalog backend; + private Path warehouseLocation; + private JdbcCatalog backend; public TestIcebergTrinoRestCatalogConnectorSmokeTest() { @@ -58,9 +66,9 @@ public TestIcebergTrinoRestCatalogConnectorSmokeTest() protected boolean hasBehavior(TestingConnectorBehavior connectorBehavior) { return switch (connectorBehavior) { - case SUPPORTS_RENAME_SCHEMA -> false; - case SUPPORTS_CREATE_VIEW, SUPPORTS_COMMENT_ON_VIEW, SUPPORTS_COMMENT_ON_VIEW_COLUMN -> false; - case SUPPORTS_CREATE_MATERIALIZED_VIEW, SUPPORTS_RENAME_MATERIALIZED_VIEW -> false; + case SUPPORTS_CREATE_MATERIALIZED_VIEW, + SUPPORTS_RENAME_MATERIALIZED_VIEW, + SUPPORTS_RENAME_SCHEMA -> false; default -> super.hasBehavior(connectorBehavior); }; } @@ -69,8 +77,8 @@ protected boolean hasBehavior(TestingConnectorBehavior connectorBehavior) protected QueryRunner createQueryRunner() throws Exception { - warehouseLocation = Files.newTemporaryFolder(); - closeAfterClass(() -> deleteRecursively(warehouseLocation.toPath(), ALLOW_INSECURE)); + warehouseLocation = Files.newTemporaryFolder().toPath(); + closeAfterClass(() -> deleteRecursively(warehouseLocation, ALLOW_INSECURE)); backend = closeAfterClass((JdbcCatalog) backendCatalog(warehouseLocation)); @@ -83,7 +91,7 @@ protected QueryRunner createQueryRunner() closeAfterClass(testServer::stop); return IcebergQueryRunner.builder() - .setBaseDataDir(Optional.of(warehouseLocation.toPath())) + .setBaseDataDir(Optional.of(warehouseLocation)) .setIcebergProperties( ImmutableMap.builder() .put("iceberg.file-format", format.name()) @@ -91,6 +99,7 @@ protected QueryRunner createQueryRunner() .put("iceberg.rest-catalog.uri", testServer.getBaseUrl().toString()) .put("iceberg.register-table-procedure.enabled", "true") .put("iceberg.writer-sort-buffer-size", "1MB") + .put("iceberg.allowed-extra-properties", "write.metadata.delete-after-commit.enabled,write.metadata.previous-versions-max") .buildOrThrow()) .setInitialTables(REQUIRED_TPCH_TABLES) .build(); @@ -102,13 +111,69 @@ public void teardown() backend = null; // closed by closeAfterClass } - @Override - public void testView() + @Test + public void testDropSchemaCascadeWithViews() { - assertThatThrownBy(super::testView) - .hasMessageContaining("createView is not supported for Iceberg REST catalog"); + String schemaName = "test_drop_schema_cascade" + randomNameSuffix(); + + assertUpdate("CREATE SCHEMA " + schemaName); + + TableIdentifier sparkViewIdentifier = TableIdentifier.of(schemaName, "test_spark_views" + randomNameSuffix()); + backend.buildView(sparkViewIdentifier) + .withDefaultNamespace(Namespace.of(schemaName)) + .withDefaultCatalog("iceberg") + .withQuery("spark", "SELECT 1 x") + .withSchema(new Schema(required(1, "x", Types.LongType.get()))) + .create(); + + TableIdentifier trinoViewIdentifier = TableIdentifier.of(schemaName, "test_trino_views" + randomNameSuffix()); + backend.buildView(trinoViewIdentifier) + .withDefaultNamespace(Namespace.of(schemaName)) + .withDefaultCatalog("iceberg") + .withQuery("trino", "SELECT 1 x") + .withSchema(new Schema(required(1, "x", Types.LongType.get()))) + .create(); + + assertThat(backend.viewExists(sparkViewIdentifier)).isTrue(); + assertThat(backend.viewExists(trinoViewIdentifier)).isTrue(); + + assertUpdate("DROP SCHEMA " + schemaName + " CASCADE"); + + assertThat(backend.viewExists(sparkViewIdentifier)).isFalse(); + assertThat(backend.viewExists(trinoViewIdentifier)).isFalse(); + } + + @Test + public void testUnsupportedViewDialect() + { + String viewName = "test_unsupported_dialect" + randomNameSuffix(); + TableIdentifier identifier = TableIdentifier.of("tpch", viewName); + + backend.buildView(identifier) + .withDefaultNamespace(Namespace.of("tpch")) + .withDefaultCatalog("iceberg") + .withQuery("spark", "SELECT 1 x") + .withSchema(new Schema(required(1, "x", Types.LongType.get()))) + .create(); + + assertThat(computeActual("SHOW TABLES FROM iceberg.tpch").getOnlyColumnAsSet()) + .contains(viewName); + + assertThat(computeActual("SELECT table_name FROM information_schema.views WHERE table_schema = 'tpch'").getOnlyColumnAsSet()) + .doesNotContain(viewName); + + assertThat(computeActual("SELECT table_name FROM information_schema.columns WHERE table_schema = 'tpch'").getOnlyColumnAsSet()) + .doesNotContain(viewName); + + assertQueryReturnsEmptyResult("SELECT * FROM information_schema.columns WHERE table_schema = 'tpch' AND table_name = '" + viewName + "'"); + + // no isView function to check view existence in Iceberg REST catalog in 423 + //assertQueryFails("SELECT * FROM " + viewName, "Cannot read unsupported dialect 'spark' for view '.*'"); + + backend.dropView(identifier); } + @Test @Override public void testMaterializedView() { @@ -116,6 +181,7 @@ public void testMaterializedView() .hasMessageContaining("createMaterializedView is not supported for Iceberg REST catalog"); } + @Test @Override public void testRenameSchema() { @@ -126,7 +192,7 @@ public void testRenameSchema() @Override protected void dropTableFromMetastore(String tableName) { - // used when registering a table, which is not supported by the REST catalog + backend.dropTable(toIdentifier(tableName), false); } @Override @@ -148,83 +214,7 @@ protected boolean locationExists(String location) return java.nio.file.Files.exists(Path.of(location)); } - @Override - public void testRegisterTableWithTableLocation() - { - assertThatThrownBy(super::testRegisterTableWithTableLocation) - .hasMessageContaining("registerTable is not supported for Iceberg REST catalog"); - } - - @Override - public void testRegisterTableWithComments() - { - assertThatThrownBy(super::testRegisterTableWithComments) - .hasMessageContaining("registerTable is not supported for Iceberg REST catalog"); - } - - @Override - public void testRegisterTableWithShowCreateTable() - { - assertThatThrownBy(super::testRegisterTableWithShowCreateTable) - .hasMessageContaining("registerTable is not supported for Iceberg REST catalog"); - } - - @Override - public void testRegisterTableWithReInsert() - { - assertThatThrownBy(super::testRegisterTableWithReInsert) - .hasMessageContaining("registerTable is not supported for Iceberg REST catalog"); - } - - @Override - public void testRegisterTableWithDifferentTableName() - { - assertThatThrownBy(super::testRegisterTableWithDifferentTableName) - .hasMessageContaining("registerTable is not supported for Iceberg REST catalog"); - } - - @Override - public void testRegisterTableWithMetadataFile() - { - assertThatThrownBy(super::testRegisterTableWithMetadataFile) - .hasMessageContaining("registerTable is not supported for Iceberg REST catalog"); - } - - @Override - public void testRegisterTableWithTrailingSpaceInLocation() - { - assertThatThrownBy(super::testRegisterTableWithTrailingSpaceInLocation) - .hasMessageContaining("registerTable is not supported for Iceberg REST catalog"); - } - - @Override - public void testUnregisterTable() - { - assertThatThrownBy(super::testUnregisterTable) - .hasMessageContaining("unregisterTable is not supported for Iceberg REST catalogs"); - } - - @Override - public void testUnregisterBrokenTable() - { - assertThatThrownBy(super::testUnregisterBrokenTable) - .hasMessageContaining("unregisterTable is not supported for Iceberg REST catalogs"); - } - - @Override - public void testUnregisterTableNotExistingTable() - { - assertThatThrownBy(super::testUnregisterTableNotExistingTable) - .hasMessageContaining("unregisterTable is not supported for Iceberg REST catalogs"); - } - - @Override - public void testRepeatUnregisterTable() - { - assertThatThrownBy(super::testRepeatUnregisterTable) - .hasMessageContaining("unregisterTable is not supported for Iceberg REST catalogs"); - } - + @Test @Override public void testDropTableWithMissingMetadataFile() { @@ -232,27 +222,27 @@ public void testDropTableWithMissingMetadataFile() .hasMessageMatching("Failed to load table: (.*)"); } + @Test @Override public void testDropTableWithMissingSnapshotFile() { assertThatThrownBy(super::testDropTableWithMissingSnapshotFile) + .isInstanceOf(QueryFailedException.class) + .cause() + .hasMessageContaining("Failed to drop table") + .cause() .hasMessageMatching("Server error: NotFoundException: Failed to open input stream for file: (.*)"); } + @Test @Override public void testDropTableWithMissingManifestListFile() { assertThatThrownBy(super::testDropTableWithMissingManifestListFile) - .hasMessageContaining("Table location should not exist expected [false] but found [true]"); - } - - @Override - public void testDropTableWithMissingDataFile() - { - assertThatThrownBy(super::testDropTableWithMissingDataFile) - .hasMessageContaining("Table location should not exist expected [false] but found [true]"); + .hasMessageContaining("Table location should not exist"); } + @Test @Override public void testDropTableWithNonExistentTableLocation() { @@ -272,7 +262,12 @@ protected boolean isFileSorted(Location path, String sortColumnName) @Override protected void deleteDirectory(String location) { - // used when unregistering a table, which is not supported by the REST catalog + try { + deleteRecursively(Path.of(location), ALLOW_INSECURE); + } + catch (IOException e) { + throw new UncheckedIOException(e); + } } private TableIdentifier toIdentifier(String tableName) diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/rest/TestOAuth2SecurityConfig.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/rest/TestOAuth2SecurityConfig.java index 39c54c7defaa..de72a8e4aef9 100644 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/rest/TestOAuth2SecurityConfig.java +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/rest/TestOAuth2SecurityConfig.java @@ -16,12 +16,13 @@ import com.google.common.collect.ImmutableMap; import org.testng.annotations.Test; +import java.net.URI; import java.util.Map; import static io.airlift.configuration.testing.ConfigAssertions.assertFullMapping; import static io.airlift.configuration.testing.ConfigAssertions.assertRecordedDefaults; import static io.airlift.configuration.testing.ConfigAssertions.recordDefaults; -import static org.testng.Assert.assertTrue; +import static org.assertj.core.api.Assertions.assertThat; public class TestOAuth2SecurityConfig { @@ -30,7 +31,9 @@ public void testDefaults() { assertRecordedDefaults(recordDefaults(OAuth2SecurityConfig.class) .setCredential(null) - .setToken(null)); + .setToken(null) + .setScope(null) + .setServerUri(null)); } @Test @@ -39,12 +42,17 @@ public void testExplicitPropertyMappings() Map properties = ImmutableMap.builder() .put("iceberg.rest-catalog.oauth2.token", "token") .put("iceberg.rest-catalog.oauth2.credential", "credential") + .put("iceberg.rest-catalog.oauth2.scope", "scope") + .put("iceberg.rest-catalog.oauth2.server-uri", "http://localhost:8080/realms/iceberg/protocol/openid-connect/token") .buildOrThrow(); OAuth2SecurityConfig expected = new OAuth2SecurityConfig() .setCredential("credential") - .setToken("token"); - assertTrue(expected.credentialOrTokenPresent()); + .setToken("token") + .setScope("scope") + .setServerUri(URI.create("http://localhost:8080/realms/iceberg/protocol/openid-connect/token")); + assertThat(expected.credentialOrTokenPresent()).isTrue(); + assertThat(expected.scopePresentOnlyWithCredential()).isFalse(); assertFullMapping(properties, expected); } } diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/rest/TestSigV4AwsCredentialProvider.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/rest/TestSigV4AwsCredentialProvider.java new file mode 100644 index 000000000000..72a540ad4733 --- /dev/null +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/rest/TestSigV4AwsCredentialProvider.java @@ -0,0 +1,35 @@ +/* + * 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.iceberg.catalog.rest; + +import org.junit.jupiter.api.Test; +import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider; + +import static io.trino.spi.testing.InterfaceTestUtils.assertAllMethodsOverridden; +import static io.trino.spi.testing.InterfaceTestUtils.assertProperForwardingMethodsAreCalled; + +class TestSigV4AwsCredentialProvider +{ + @Test + void testEverythingImplemented() + { + assertAllMethodsOverridden(AwsCredentialsProvider.class, SigV4AwsCredentialProvider.class); + } + + @Test + void testProperForwardingMethodsAreCalled() + { + assertProperForwardingMethodsAreCalled(AwsCredentialsProvider.class, SigV4AwsCredentialProvider::new); + } +} diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/rest/TestTrinoRestCatalog.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/rest/TestTrinoRestCatalog.java index 405bb5c90a44..2d76afccccbc 100644 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/rest/TestTrinoRestCatalog.java +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/rest/TestTrinoRestCatalog.java @@ -14,30 +14,41 @@ package io.trino.plugin.iceberg.catalog.rest; import com.google.common.collect.ImmutableMap; +import com.google.common.io.Files; +import io.trino.cache.EvictableCacheBuilder; import io.trino.plugin.base.CatalogName; import io.trino.plugin.hive.NodeVersion; +import io.trino.plugin.hive.metastore.TableInfo; import io.trino.plugin.iceberg.CommitTaskData; import io.trino.plugin.iceberg.IcebergMetadata; import io.trino.plugin.iceberg.TableStatisticsWriter; import io.trino.plugin.iceberg.catalog.BaseTrinoCatalogTest; import io.trino.plugin.iceberg.catalog.TrinoCatalog; +import io.trino.spi.TrinoException; import io.trino.spi.connector.CatalogHandle; import io.trino.spi.connector.ConnectorMetadata; import io.trino.spi.security.PrincipalType; import io.trino.spi.security.TrinoPrincipal; +import io.trino.spi.type.TestingTypeManager; import org.apache.iceberg.rest.DelegatingRestSessionCatalog; import org.apache.iceberg.rest.RESTSessionCatalog; -import org.assertj.core.util.Files; +import org.testng.annotations.Test; -import java.io.File; +import java.nio.file.Path; +import java.util.Map; +import java.util.Optional; +import static com.google.common.util.concurrent.MoreExecutors.directExecutor; +import static com.google.common.util.concurrent.MoreExecutors.newDirectExecutorService; import static io.airlift.json.JsonCodec.jsonCodec; +import static io.trino.plugin.hive.metastore.TableInfo.ExtendedRelationType.OTHER_VIEW; import static io.trino.plugin.iceberg.catalog.rest.IcebergRestCatalogConfig.SessionType.NONE; import static io.trino.plugin.iceberg.catalog.rest.RestCatalogTestUtils.backendCatalog; import static io.trino.sql.planner.TestingPlannerContext.PLANNER_CONTEXT; import static io.trino.testing.TestingConnectorSession.SESSION; import static io.trino.testing.TestingNames.randomNameSuffix; import static java.util.Locale.ENGLISH; +import static java.util.concurrent.TimeUnit.MILLISECONDS; import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatThrownBy; @@ -47,8 +58,13 @@ public class TestTrinoRestCatalog @Override protected TrinoCatalog createTrinoCatalog(boolean useUniqueTableLocations) { - File warehouseLocation = Files.newTemporaryFolder(); - warehouseLocation.deleteOnExit(); + return createTrinoRestCatalog(useUniqueTableLocations, ImmutableMap.of()); + } + + private static TrinoRestCatalog createTrinoRestCatalog(boolean useUniqueTableLocations, Map properties) + { + Path warehouseLocation = Files.createTempDir().toPath(); + warehouseLocation.toFile().deleteOnExit(); String catalogName = "iceberg_rest"; RESTSessionCatalog restSessionCatalog = DelegatingRestSessionCatalog @@ -56,18 +72,23 @@ protected TrinoCatalog createTrinoCatalog(boolean useUniqueTableLocations) .delegate(backendCatalog(warehouseLocation)) .build(); - restSessionCatalog.initialize(catalogName, ImmutableMap.of()); - - return new TrinoRestCatalog(restSessionCatalog, new CatalogName(catalogName), NONE, "test", useUniqueTableLocations); - } + restSessionCatalog.initialize(catalogName, properties); - @Override - public void testView() - { - assertThatThrownBy(super::testView) - .hasMessageContaining("createView is not supported for Iceberg REST catalog"); + return new TrinoRestCatalog( + restSessionCatalog, + new CatalogName(catalogName), + NONE, + ImmutableMap.of(), + false, + "test", + new TestingTypeManager(), + useUniqueTableLocations, + false, + EvictableCacheBuilder.newBuilder().expireAfterWrite(1000, MILLISECONDS).shareNothingWhenDisabled().build(), + EvictableCacheBuilder.newBuilder().expireAfterWrite(1000, MILLISECONDS).shareNothingWhenDisabled().build()); } + @Test @Override public void testNonLowercaseNamespace() { @@ -93,10 +114,15 @@ public void testNonLowercaseNamespace() CatalogHandle.fromId("iceberg:NORMAL:v12345"), jsonCodec(CommitTaskData.class), catalog, - connectorIdentity -> { + (connectorIdentity, fileIoProperties) -> { throw new UnsupportedOperationException(); }, - new TableStatisticsWriter(new NodeVersion("test-version"))); + new TableStatisticsWriter(new NodeVersion("test-version")), + Optional.empty(), + false, + ignore -> false, + newDirectExecutorService(), + directExecutor()); assertThat(icebergMetadata.schemaExists(SESSION, namespace)).as("icebergMetadata.schemaExists(namespace)") .isTrue(); assertThat(icebergMetadata.schemaExists(SESSION, schema)).as("icebergMetadata.schemaExists(schema)") @@ -109,4 +135,31 @@ public void testNonLowercaseNamespace() catalog.dropNamespace(SESSION, namespace); } } + + @Test + public void testPrefix() + throws Exception + { + TrinoCatalog catalog = createTrinoRestCatalog(false, ImmutableMap.of("prefix", "dev")); + + String namespace = "testPrefixNamespace" + randomNameSuffix(); + + assertThatThrownBy(() -> + catalog.createNamespace( + SESSION, + namespace, + defaultNamespaceProperties(namespace), + new TrinoPrincipal(PrincipalType.USER, SESSION.getUser()))) + .isInstanceOf(TrinoException.class) + .hasMessageContaining("Failed to create namespace") + .cause() + .as("should fail as the prefix dev is not implemented for the current endpoint") + .hasMessageContaining("Malformed request: No route for request: POST v1/dev/namespaces"); + } + + @Override + protected TableInfo.ExtendedRelationType getViewType() + { + return OTHER_VIEW; + } } diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/optimizer/TestConnectorPushdownRulesWithIceberg.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/optimizer/TestConnectorPushdownRulesWithIceberg.java index d5f0c9994c5d..e0e000a0c711 100644 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/optimizer/TestConnectorPushdownRulesWithIceberg.java +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/optimizer/TestConnectorPushdownRulesWithIceberg.java @@ -61,7 +61,6 @@ import static com.google.common.io.MoreFiles.deleteRecursively; import static com.google.common.io.RecursiveDeleteOption.ALLOW_INSECURE; -import static com.google.inject.util.Modules.EMPTY_MODULE; import static io.trino.plugin.hive.metastore.file.TestingFileHiveMetastore.createTestingFileHiveMetastore; import static io.trino.plugin.iceberg.ColumnIdentity.TypeCategory.STRUCT; import static io.trino.plugin.iceberg.ColumnIdentity.primitiveColumnIdentity; @@ -115,7 +114,6 @@ protected Optional createLocalQueryRunner() metastore.createDatabase(database); - HiveMetastore metastore = createTestingFileHiveMetastore(baseDir); LocalQueryRunner queryRunner = LocalQueryRunner.create(ICEBERG_SESSION); InternalFunctionBundle.InternalFunctionBundleBuilder functions = InternalFunctionBundle.builder(); @@ -124,7 +122,7 @@ protected Optional createLocalQueryRunner() queryRunner.createCatalog( TEST_CATALOG_NAME, - new TestingIcebergConnectorFactory(Optional.of(new TestingIcebergFileMetastoreCatalogModule(metastore)), Optional.empty(), EMPTY_MODULE), + new TestingIcebergConnectorFactory(baseDir.toPath(), Optional.of(new TestingIcebergFileMetastoreCatalogModule(metastore))), ImmutableMap.of()); catalogHandle = queryRunner.getCatalogHandle(TEST_CATALOG_NAME); return Optional.of(queryRunner); @@ -151,6 +149,7 @@ public void testProjectionPushdown() baseType, ImmutableList.of(1), BIGINT, + true, Optional.empty()); IcebergTableHandle icebergTable = new IcebergTableHandle( @@ -169,7 +168,10 @@ public void testProjectionPushdown() Optional.empty(), "", ImmutableMap.of(), + Optional.empty(), false, + Optional.empty(), + ImmutableSet.of(), Optional.empty()); TableHandle table = new TableHandle(catalogHandle, icebergTable, new HiveTransactionHandle(false)); @@ -251,11 +253,14 @@ public void testPredicatePushdown() Optional.empty(), "", ImmutableMap.of(), + Optional.empty(), false, + Optional.empty(), + ImmutableSet.of(), Optional.empty()); TableHandle table = new TableHandle(catalogHandle, icebergTable, new HiveTransactionHandle(false)); - IcebergColumnHandle column = new IcebergColumnHandle(primitiveColumnIdentity(1, "a"), INTEGER, ImmutableList.of(), INTEGER, Optional.empty()); + IcebergColumnHandle column = new IcebergColumnHandle(primitiveColumnIdentity(1, "a"), INTEGER, ImmutableList.of(), INTEGER, false, Optional.empty()); tester().assertThat(pushPredicateIntoTableScan) .on(p -> @@ -300,12 +305,15 @@ public void testColumnPruningProjectionPushdown() Optional.empty(), "", ImmutableMap.of(), + Optional.empty(), false, + Optional.empty(), + ImmutableSet.of(), Optional.empty()); TableHandle table = new TableHandle(catalogHandle, icebergTable, new HiveTransactionHandle(false)); - IcebergColumnHandle columnA = new IcebergColumnHandle(primitiveColumnIdentity(0, "a"), INTEGER, ImmutableList.of(), INTEGER, Optional.empty()); - IcebergColumnHandle columnB = new IcebergColumnHandle(primitiveColumnIdentity(1, "b"), INTEGER, ImmutableList.of(), INTEGER, Optional.empty()); + IcebergColumnHandle columnA = new IcebergColumnHandle(primitiveColumnIdentity(0, "a"), INTEGER, ImmutableList.of(), INTEGER, true, Optional.empty()); + IcebergColumnHandle columnB = new IcebergColumnHandle(primitiveColumnIdentity(1, "b"), INTEGER, ImmutableList.of(), INTEGER, true, Optional.empty()); tester().assertThat(pruneTableScanColumns) .on(p -> { @@ -360,16 +368,20 @@ public void testPushdownWithDuplicateExpressions() Optional.empty(), "", ImmutableMap.of(), + Optional.empty(), false, + Optional.empty(), + ImmutableSet.of(), Optional.empty()); TableHandle table = new TableHandle(catalogHandle, icebergTable, new HiveTransactionHandle(false)); - IcebergColumnHandle bigintColumn = new IcebergColumnHandle(primitiveColumnIdentity(1, "just_bigint"), BIGINT, ImmutableList.of(), BIGINT, Optional.empty()); + IcebergColumnHandle bigintColumn = new IcebergColumnHandle(primitiveColumnIdentity(1, "just_bigint"), BIGINT, ImmutableList.of(), BIGINT, false, Optional.empty()); IcebergColumnHandle partialColumn = new IcebergColumnHandle( new ColumnIdentity(3, "struct_of_bigint", STRUCT, ImmutableList.of(primitiveColumnIdentity(1, "a"), primitiveColumnIdentity(2, "b"))), ROW_TYPE, ImmutableList.of(1), BIGINT, + false, Optional.empty()); // Test projection pushdown with duplicate column references diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/procedure/TestIcebergOptimizeManifestsProcedure.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/procedure/TestIcebergOptimizeManifestsProcedure.java new file mode 100644 index 000000000000..cff248bb4910 --- /dev/null +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/procedure/TestIcebergOptimizeManifestsProcedure.java @@ -0,0 +1,140 @@ +/* + * 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.iceberg.procedure; + +import com.google.common.collect.ImmutableList; +import io.trino.plugin.iceberg.IcebergQueryRunner; +import io.trino.testing.AbstractTestQueryFramework; +import io.trino.testing.QueryRunner; +import io.trino.testing.sql.TestTable; +import org.intellij.lang.annotations.Language; +import org.testng.annotations.Test; + +import java.util.List; +import java.util.Set; + +import static com.google.common.collect.ImmutableSet.toImmutableSet; +import static org.assertj.core.api.Assertions.assertThat; + +final class TestIcebergOptimizeManifestsProcedure + extends AbstractTestQueryFramework +{ + @Override + protected QueryRunner createQueryRunner() + throws Exception + { + return IcebergQueryRunner.builder().build(); + } + + protected TestTable newTrinoTable(String namePrefix, @Language("SQL") String tableDefinition) + { + return newTrinoTable(namePrefix, tableDefinition, ImmutableList.of()); + } + + protected TestTable newTrinoTable(String namePrefix, @Language("SQL") String tableDefinition, List rowsToInsert) + { + return new TestTable(getQueryRunner()::execute, namePrefix, tableDefinition, rowsToInsert); + } + + @Test + public void testOptimizeManifests() + { + try (TestTable table = newTrinoTable("test_optimize_manifests", "(x int)")) { + assertUpdate("INSERT INTO " + table.getName() + " VALUES 1", 1); + assertUpdate("INSERT INTO " + table.getName() + " VALUES 2", 1); + + Set manifestFiles = manifestFiles(table.getName()); + assertThat(manifestFiles).hasSize(2); + + assertUpdate("ALTER TABLE " + table.getName() + " EXECUTE optimize_manifests"); + assertThat(manifestFiles(table.getName())) + .hasSize(1) + .doesNotContainAnyElementsOf(manifestFiles); + + assertThat(query("SELECT * FROM " + table.getName())) + .matches("VALUES 1, 2"); + } + } + + @Test + public void testPartitionTable() + { + try (TestTable table = newTrinoTable("test_partition", "(id int, part int) WITH (partitioning = ARRAY['part'])")) { + assertUpdate("INSERT INTO " + table.getName() + " VALUES (1, 10)", 1); + assertUpdate("INSERT INTO " + table.getName() + " VALUES (2, 10)", 1); + assertUpdate("INSERT INTO " + table.getName() + " VALUES (3, 20)", 1); + assertUpdate("INSERT INTO " + table.getName() + " VALUES (4, 20)", 1); + + Set manifestFiles = manifestFiles(table.getName()); + assertThat(manifestFiles).hasSize(4); + + assertUpdate("ALTER TABLE " + table.getName() + " EXECUTE optimize_manifests"); + assertThat(manifestFiles(table.getName())) + .hasSize(1) + .doesNotContainAnyElementsOf(manifestFiles); + + assertThat(query("SELECT * FROM " + table.getName())) + .matches("VALUES (1, 10), (2, 10), (3, 20), (4, 20)"); + } + } + + @Test + public void testEmptyManifest() + { + try (TestTable table = newTrinoTable("test_no_rewrite", "(x int)")) { + Set manifestFiles = manifestFiles(table.getName()); + assertThat(manifestFiles).isEmpty(); + + assertUpdate("ALTER TABLE " + table.getName() + " EXECUTE optimize_manifests"); + assertThat(manifestFiles(table.getName())).isEmpty(); + + assertQueryReturnsEmptyResult("SELECT * FROM " + table.getName()); + } + } + + @Test + public void testNotRewriteSingleManifest() + { + try (TestTable table = newTrinoTable("test_no_rewrite", "(x int)")) { + assertUpdate("INSERT INTO " + table.getName() + " VALUES 1", 1); + + Set manifestFiles = manifestFiles(table.getName()); + assertThat(manifestFiles).hasSize(1); + + assertUpdate("ALTER TABLE " + table.getName() + " EXECUTE optimize_manifests"); + assertThat(manifestFiles(table.getName())) + .hasSize(1) + .isEqualTo(manifestFiles); + + assertThat(query("SELECT * FROM " + table.getName())) + .matches("VALUES 1"); + } + } + + @Test + public void testUnsupportedWhere() + { + try (TestTable table = newTrinoTable("test_unsupported_where", "WITH (partitioning = ARRAY['part']) AS SELECT 1 id, 1 part")) { + assertQueryFails("ALTER TABLE " + table.getName() + " EXECUTE optimize_manifests WHERE id = 1", ".* WHERE not supported for procedure OPTIMIZE_MANIFESTS"); + assertQueryFails("ALTER TABLE " + table.getName() + " EXECUTE optimize_manifests WHERE part = 10", ".* WHERE not supported for procedure OPTIMIZE_MANIFESTS"); + } + } + + private Set manifestFiles(String tableName) + { + return computeActual("SELECT path FROM \"" + tableName + "$manifests\"").getOnlyColumnAsSet().stream() + .map(path -> (String) path) + .collect(toImmutableSet()); + } +} diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/util/EqualityDeleteUtils.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/util/EqualityDeleteUtils.java new file mode 100644 index 000000000000..c7ddba271d63 --- /dev/null +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/util/EqualityDeleteUtils.java @@ -0,0 +1,94 @@ +/* + * 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.iceberg.util; + +import io.trino.filesystem.TrinoFileSystemFactory; +import io.trino.plugin.iceberg.PartitionData; +import io.trino.plugin.iceberg.fileio.ForwardingFileIo; +import org.apache.iceberg.PartitionSpec; +import org.apache.iceberg.Schema; +import org.apache.iceberg.Table; +import org.apache.iceberg.data.GenericRecord; +import org.apache.iceberg.data.Record; +import org.apache.iceberg.data.parquet.GenericParquetWriter; +import org.apache.iceberg.deletes.EqualityDeleteWriter; +import org.apache.iceberg.io.FileIO; +import org.apache.iceberg.parquet.Parquet; + +import java.io.Closeable; +import java.io.IOException; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.UUID; + +import static com.google.common.collect.ImmutableList.toImmutableList; +import static io.trino.testing.TestingConnectorSession.SESSION; + +public final class EqualityDeleteUtils +{ + private EqualityDeleteUtils() {} + + public static void writeEqualityDeleteForTable( + Table icebergTable, + TrinoFileSystemFactory fileSystemFactory, + Optional partitionSpec, + Optional partitionData, + Map overwriteValues, + Optional> deleteFileColumns) + throws IOException + { + List deleteColumns = deleteFileColumns.orElse(new ArrayList<>(overwriteValues.keySet())); + Schema deleteRowSchema = icebergTable.schema().select(deleteColumns); + List equalityDeleteFieldIds = deleteColumns.stream() + .map(name -> deleteRowSchema.findField(name).fieldId()) + .collect(toImmutableList()); + writeEqualityDeleteForTableWithSchema(icebergTable, fileSystemFactory, partitionSpec, partitionData, deleteRowSchema, equalityDeleteFieldIds, overwriteValues); + } + + public static void writeEqualityDeleteForTableWithSchema( + Table icebergTable, + TrinoFileSystemFactory fileSystemFactory, + Optional partitionSpec, + Optional partitionData, + Schema deleteRowSchema, + List equalityDeleteFieldIds, + Map overwriteValues) + throws IOException + { + String deleteFileName = "local:///delete_file_" + UUID.randomUUID() + ".parquet"; + FileIO fileIo = new ForwardingFileIo(fileSystemFactory.create(SESSION)); + + Parquet.DeleteWriteBuilder writerBuilder = Parquet.writeDeletes(fileIo.newOutputFile(deleteFileName)) + .forTable(icebergTable) + .rowSchema(deleteRowSchema) + .createWriterFunc(GenericParquetWriter::buildWriter) + .equalityFieldIds(equalityDeleteFieldIds) + .overwrite(); + if (partitionSpec.isPresent() && partitionData.isPresent()) { + writerBuilder = writerBuilder + .withSpec(partitionSpec.get()) + .withPartition(partitionData.get()); + } + EqualityDeleteWriter writer = writerBuilder.buildEqualityWriter(); + + Record dataDelete = GenericRecord.create(deleteRowSchema); + try (Closeable ignored = writer) { + writer.write(dataDelete.copy(overwriteValues)); + } + + icebergTable.newRowDelta().addDeletes(writer.toDeleteFile()).commit(); + } +} diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/util/FileOperationUtils.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/util/FileOperationUtils.java new file mode 100644 index 000000000000..d93e19527807 --- /dev/null +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/util/FileOperationUtils.java @@ -0,0 +1,107 @@ +/* + * 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.iceberg.util; + +import com.google.common.collect.HashMultiset; +import com.google.common.collect.Multiset; +import io.opentelemetry.sdk.trace.data.SpanData; + +import java.util.List; +import java.util.function.Predicate; + +import static io.trino.filesystem.tracing.FileSystemAttributes.FILE_LOCATION; +import static io.trino.plugin.iceberg.util.FileOperationUtils.FileType.DATA; +import static io.trino.plugin.iceberg.util.FileOperationUtils.FileType.METASTORE; +import static io.trino.plugin.iceberg.util.FileOperationUtils.FileType.fromFilePath; +import static java.util.Objects.requireNonNull; +import static java.util.stream.Collectors.toCollection; + +public final class FileOperationUtils +{ + private FileOperationUtils() {} + + public static Multiset getOperations(List spans) + { + return spans.stream() + .filter(span -> span.getName().startsWith("InputFile.") || span.getName().startsWith("OutputFile.")) + .map(span -> new FileOperation(fromFilePath(requireNonNull(span.getAttributes().get(FILE_LOCATION))), span.getName())) + .collect(toCollection(HashMultiset::create)); + } + + public record FileOperation(FileType fileType, String operationType) + { + public FileOperation + { + requireNonNull(fileType, "fileType is null"); + requireNonNull(operationType, "operationType is null"); + } + } + + public enum Scope + implements Predicate + { + METADATA_FILES { + @Override + public boolean test(FileOperation fileOperation) + { + return fileOperation.fileType() != DATA && fileOperation.fileType() != METASTORE; + } + }, + ALL_FILES { + @Override + public boolean test(FileOperation fileOperation) + { + return fileOperation.fileType() != METASTORE; + } + }, + } + + public enum FileType + { + METADATA_JSON, + SNAPSHOT, + MANIFEST, + STATS, + DATA, + DELETE, + METASTORE, + /**/; + + public static FileType fromFilePath(String path) + { + if (path.endsWith("metadata.json")) { + return METADATA_JSON; + } + if (path.contains("/snap-")) { + return SNAPSHOT; + } + if (path.endsWith("-m0.avro")) { + return MANIFEST; + } + if (path.endsWith(".stats")) { + return STATS; + } + if (path.contains("/data/") && (path.endsWith(".orc") || path.endsWith(".parquet"))) { + return DATA; + } + if (path.contains("delete_file") && (path.endsWith(".orc") || path.endsWith(".parquet"))) { + return DELETE; + } + if (path.endsWith(".trinoSchema") || path.contains("/.trinoPermissions/")) { + return METASTORE; + } + throw new IllegalArgumentException("File not recognized: " + path); + } + } +} diff --git a/plugin/trino-iceberg/src/test/java/org/apache/iceberg/jdbc/TestingTrinoIcebergJdbcUtil.java b/plugin/trino-iceberg/src/test/java/org/apache/iceberg/jdbc/TestingTrinoIcebergJdbcUtil.java index f15d86b24b66..4967d9708188 100644 --- a/plugin/trino-iceberg/src/test/java/org/apache/iceberg/jdbc/TestingTrinoIcebergJdbcUtil.java +++ b/plugin/trino-iceberg/src/test/java/org/apache/iceberg/jdbc/TestingTrinoIcebergJdbcUtil.java @@ -16,8 +16,8 @@ public final class TestingTrinoIcebergJdbcUtil { - public static final String CREATE_CATALOG_TABLE = JdbcUtil.CREATE_CATALOG_TABLE; - public static final String CREATE_NAMESPACE_PROPERTIES_TABLE = JdbcUtil.CREATE_NAMESPACE_PROPERTIES_TABLE; + public static final String CREATE_CATALOG_TABLE = JdbcUtil.V0_CREATE_CATALOG_SQL; + public static final String CREATE_NAMESPACE_PROPERTIES_TABLE = JdbcUtil.CREATE_NAMESPACE_PROPERTIES_TABLE_SQL; private TestingTrinoIcebergJdbcUtil() {} } diff --git a/plugin/trino-iceberg/src/test/java/org/apache/iceberg/rest/RestCatalogServlet.java b/plugin/trino-iceberg/src/test/java/org/apache/iceberg/rest/RestCatalogServlet.java index 0ccb550e88b2..3fc7e41bb67a 100644 --- a/plugin/trino-iceberg/src/test/java/org/apache/iceberg/rest/RestCatalogServlet.java +++ b/plugin/trino-iceberg/src/test/java/org/apache/iceberg/rest/RestCatalogServlet.java @@ -22,7 +22,7 @@ import org.apache.iceberg.exceptions.RESTException; import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap; import org.apache.iceberg.relocated.com.google.common.io.CharStreams; -import org.apache.iceberg.rest.RESTCatalogAdapter.HTTPMethod; +import org.apache.iceberg.rest.HTTPRequest.HTTPMethod; import org.apache.iceberg.rest.RESTCatalogAdapter.Route; import org.apache.iceberg.rest.responses.ErrorResponse; import org.apache.iceberg.util.Pair; @@ -98,15 +98,19 @@ protected void execute(ServletRequestContext context, HttpServletResponse respon return; } + HTTPRequest request = restCatalogAdapter.buildRequest( + context.method(), + context.path(), + context.queryParams(), + context.headers(), + context.body()); + try { Object responseBody = restCatalogAdapter.execute( - context.method(), - context.path(), - context.queryParams(), - context.body(), + request, context.route().responseClass(), - context.headers(), - handle(response)); + handle(response), + x -> {}); if (responseBody != null) { RESTObjectMapper.mapper().writeValue(response.getWriter(), responseBody); diff --git a/plugin/trino-pinot/pom.xml b/plugin/trino-pinot/pom.xml index bc33a4d1b8c6..fb7f6fea290e 100755 --- a/plugin/trino-pinot/pom.xml +++ b/plugin/trino-pinot/pom.xml @@ -495,19 +495,24 @@ org.apache.avro avro runtime + + + org.apache.commons + commons-compress + + org.apache.commons commons-lang3 - 3.11 + 3.17.0 runtime org.apache.httpcomponents httpcore - 4.4.13 runtime diff --git a/plugin/trino-raptor-legacy/src/main/java/io/trino/plugin/raptor/legacy/RaptorMergeSink.java b/plugin/trino-raptor-legacy/src/main/java/io/trino/plugin/raptor/legacy/RaptorMergeSink.java index 5fa9fc37c051..0d7fcc59d425 100644 --- a/plugin/trino-raptor-legacy/src/main/java/io/trino/plugin/raptor/legacy/RaptorMergeSink.java +++ b/plugin/trino-raptor-legacy/src/main/java/io/trino/plugin/raptor/legacy/RaptorMergeSink.java @@ -23,7 +23,6 @@ import io.trino.plugin.raptor.legacy.storage.StorageManager; import io.trino.spi.Page; import io.trino.spi.block.Block; -import io.trino.spi.block.ColumnarRow; import io.trino.spi.connector.ConnectorMergeSink; import io.trino.spi.connector.ConnectorPageSink; import io.trino.spi.connector.MergePage; @@ -43,7 +42,7 @@ import static com.google.common.base.Verify.verify; import static com.google.common.collect.ImmutableList.toImmutableList; import static io.airlift.json.JsonCodec.jsonCodec; -import static io.trino.spi.block.ColumnarRow.toColumnarRow; +import static io.trino.spi.block.RowBlock.getRowFieldsFromBlock; import static io.trino.spi.connector.MergePage.createDeleteAndInsertPages; import static io.trino.spi.type.BigintType.BIGINT; import static io.trino.spi.type.IntegerType.INTEGER; @@ -81,12 +80,12 @@ public void storeMergedRows(Page page) mergePage.getInsertionsPage().ifPresent(pageSink::appendPage); mergePage.getDeletionsPage().ifPresent(deletions -> { - ColumnarRow rowIdRow = toColumnarRow(deletions.getBlock(deletions.getChannelCount() - 1)); - Block shardBucketBlock = rowIdRow.getField(0); - Block shardUuidBlock = rowIdRow.getField(1); - Block shardRowIdBlock = rowIdRow.getField(2); + List fields = getRowFieldsFromBlock(deletions.getBlock(deletions.getChannelCount() - 1)); + Block shardBucketBlock = fields.get(0); + Block shardUuidBlock = fields.get(1); + Block shardRowIdBlock = fields.get(2); - for (int position = 0; position < rowIdRow.getPositionCount(); position++) { + for (int position = 0; position < shardBucketBlock.getPositionCount(); position++) { OptionalInt bucketNumber = shardBucketBlock.isNull(position) ? OptionalInt.empty() : OptionalInt.of(INTEGER.getInt(shardBucketBlock, position)); diff --git a/pom.xml b/pom.xml index 2b12d5a52bd5..956cefcd5b3f 100644 --- a/pom.xml +++ b/pom.xml @@ -153,36 +153,37 @@ 4.13.0 235 12.0.1 - 1.11.1 + 1.12.0 ${dep.airlift.version} 2.1.1 1.12.505 - 2.20.93 + 2.30.23 0.11.5 21.9.0.0 1.21 - 200 + 202 2.2.8 2.21.1 - 1.18.3 + 1.20.6 1.0.8 7.3.1 3.3.2 - 4.14.0 + 4.17.0 8.4.5 - 1.3.0 + 5.3.4 + 1.8.1 3.23.2 4.5.0 - 4.1.93.Final + 4.1.118.Final 5.13.0 - 3.3.0 + 3.11.0 9.21.0 - 1.13.1 + 1.15.1 1.37 - 81 + 110 **/Test*.java diff --git a/testing/trino-product-tests/pom.xml b/testing/trino-product-tests/pom.xml index 816a244fcbad..def12d4884d8 100644 --- a/testing/trino-product-tests/pom.xml +++ b/testing/trino-product-tests/pom.xml @@ -23,11 +23,6 @@ aws-java-sdk-core - - com.amazonaws - aws-java-sdk-glue - - com.amazonaws aws-java-sdk-s3 @@ -249,6 +244,11 @@ testng + + software.amazon.awssdk + glue + + com.clickhouse clickhouse-jdbc diff --git a/testing/trino-product-tests/src/main/java/io/trino/tests/product/deltalake/TestDatabricksWithGlueMetastoreCleanUp.java b/testing/trino-product-tests/src/main/java/io/trino/tests/product/deltalake/TestDatabricksWithGlueMetastoreCleanUp.java index f728dcdb8c45..568cea2800a5 100644 --- a/testing/trino-product-tests/src/main/java/io/trino/tests/product/deltalake/TestDatabricksWithGlueMetastoreCleanUp.java +++ b/testing/trino-product-tests/src/main/java/io/trino/tests/product/deltalake/TestDatabricksWithGlueMetastoreCleanUp.java @@ -13,16 +13,14 @@ */ package io.trino.tests.product.deltalake; -import com.amazonaws.services.glue.AWSGlueAsync; -import com.amazonaws.services.glue.AWSGlueAsyncClientBuilder; -import com.amazonaws.services.glue.model.GetTableRequest; -import com.amazonaws.services.glue.model.Table; import com.google.common.collect.ImmutableSet; import io.airlift.log.Logger; import io.trino.tempto.ProductTest; import io.trino.tempto.query.QueryResult; import io.trino.testng.services.Flaky; import org.testng.annotations.Test; +import software.amazon.awssdk.services.glue.GlueClient; +import software.amazon.awssdk.services.glue.model.Table; import java.time.Instant; import java.time.temporal.ChronoUnit; @@ -31,7 +29,7 @@ import java.util.stream.Collectors; import static com.google.common.collect.ImmutableSet.toImmutableSet; -import static io.trino.plugin.hive.metastore.glue.converter.GlueToTrinoConverter.getTableType; +import static io.trino.plugin.hive.metastore.glue.GlueConverter.getTableType; import static io.trino.tests.product.TestGroups.DELTA_LAKE_DATABRICKS; import static io.trino.tests.product.TestGroups.PROFILE_SPECIFIC_TESTS; import static io.trino.tests.product.deltalake.util.DeltaLakeTestUtils.DATABRICKS_COMMUNICATION_FAILURE_ISSUE; @@ -54,7 +52,7 @@ public class TestDatabricksWithGlueMetastoreCleanUp @Flaky(issue = DATABRICKS_COMMUNICATION_FAILURE_ISSUE, match = DATABRICKS_COMMUNICATION_FAILURE_MATCH) public void testCleanUpOldTablesUsingDelta() { - AWSGlueAsync glueClient = AWSGlueAsyncClientBuilder.standard().build(); + GlueClient glueClient = GlueClient.create(); long startTime = currentTimeMillis(); List schemas = onTrino().executeQuery("SELECT DISTINCT(table_schema) FROM information_schema.tables") .rows().stream() @@ -67,7 +65,7 @@ public void testCleanUpOldTablesUsingDelta() schemas.forEach(schema -> cleanSchema(schema, startTime, glueClient)); } - private void cleanSchema(String schema, long startTime, AWSGlueAsync glueClient) + private void cleanSchema(String schema, long startTime, GlueClient glueClient) { Set allTestTableNames = findAllTablesInSchema(schema).stream() .filter(name -> name.toLowerCase(ENGLISH).startsWith("test")) @@ -76,8 +74,8 @@ private void cleanSchema(String schema, long startTime, AWSGlueAsync glueClient) int droppedTablesCount = 0; for (String tableName : allTestTableNames) { try { - Table table = glueClient.getTable(new GetTableRequest().withDatabaseName(schema).withName(tableName)).getTable(); - Instant createTime = table.getCreateTime().toInstant(); + Table table = glueClient.getTable(builder -> builder.databaseName(schema).name(tableName)).table(); + Instant createTime = table.createTime(); if (createTime.isBefore(SCHEMA_CLEANUP_THRESHOLD)) { if (getTableType(table).contains("VIEW")) { onTrino().executeQuery(format("DROP VIEW IF EXISTS %s.%s", schema, tableName)); diff --git a/testing/trino-product-tests/src/main/java/io/trino/tests/product/deltalake/TestDeltaLakeDatabricksCleanUpGlueMetastore.java b/testing/trino-product-tests/src/main/java/io/trino/tests/product/deltalake/TestDeltaLakeDatabricksCleanUpGlueMetastore.java index b08200cd6ae1..36a47598e848 100644 --- a/testing/trino-product-tests/src/main/java/io/trino/tests/product/deltalake/TestDeltaLakeDatabricksCleanUpGlueMetastore.java +++ b/testing/trino-product-tests/src/main/java/io/trino/tests/product/deltalake/TestDeltaLakeDatabricksCleanUpGlueMetastore.java @@ -13,22 +13,16 @@ */ package io.trino.tests.product.deltalake; -import com.amazonaws.services.glue.AWSGlueAsync; -import com.amazonaws.services.glue.AWSGlueAsyncClientBuilder; -import com.amazonaws.services.glue.model.Database; -import com.amazonaws.services.glue.model.DeleteDatabaseRequest; -import com.amazonaws.services.glue.model.EntityNotFoundException; -import com.amazonaws.services.glue.model.GetDatabasesRequest; -import com.amazonaws.services.glue.model.GetDatabasesResult; import io.airlift.log.Logger; -import io.trino.plugin.hive.aws.AwsApiCallStats; import io.trino.tempto.ProductTest; import org.testng.annotations.Test; +import software.amazon.awssdk.services.glue.GlueClient; +import software.amazon.awssdk.services.glue.model.Database; +import software.amazon.awssdk.services.glue.model.EntityNotFoundException; +import software.amazon.awssdk.services.glue.model.GetDatabasesResponse; import java.util.List; -import static com.google.common.collect.ImmutableList.toImmutableList; -import static io.trino.plugin.hive.metastore.glue.AwsSdkUtil.getPaginatedResults; import static io.trino.tests.product.TestGroups.DELTA_LAKE_DATABRICKS; import static io.trino.tests.product.TestGroups.PROFILE_SPECIFIC_TESTS; import static java.lang.System.currentTimeMillis; @@ -43,27 +37,21 @@ public class TestDeltaLakeDatabricksCleanUpGlueMetastore @Test(groups = {DELTA_LAKE_DATABRICKS, PROFILE_SPECIFIC_TESTS}) public void testCleanupOrphanedDatabases() { - AWSGlueAsync glueClient = AWSGlueAsyncClientBuilder.defaultClient(); + GlueClient glueClient = GlueClient.create(); long creationTimeMillisThreshold = currentTimeMillis() - DAYS.toMillis(1); - List orphanedDatabases = getPaginatedResults( - glueClient::getDatabases, - new GetDatabasesRequest(), - GetDatabasesRequest::setNextToken, - GetDatabasesResult::getNextToken, - new AwsApiCallStats()) - .map(GetDatabasesResult::getDatabaseList) + List orphanedDatabases = glueClient.getDatabasesPaginator(x -> {}).stream() + .map(GetDatabasesResponse::databaseList) .flatMap(List::stream) .filter(database -> isOrphanedTestDatabase(database, creationTimeMillisThreshold)) - .map(Database::getName) - .collect(toImmutableList()); + .map(Database::name) + .toList(); if (!orphanedDatabases.isEmpty()) { log.info("Found %s %s* databases that look orphaned, removing", orphanedDatabases.size(), TEST_DATABASE_NAME_PREFIX); orphanedDatabases.forEach(database -> { try { log.info("Deleting %s database", database); - glueClient.deleteDatabase(new DeleteDatabaseRequest() - .withName(database)); + glueClient.deleteDatabase(builder -> builder.name(database)); } catch (EntityNotFoundException e) { log.info("Database [%s] not found, could be removed by other cleanup process", database); @@ -77,7 +65,7 @@ public void testCleanupOrphanedDatabases() private static boolean isOrphanedTestDatabase(Database database, long creationTimeMillisThreshold) { - return database.getName().startsWith(TEST_DATABASE_NAME_PREFIX) && - database.getCreateTime().getTime() <= creationTimeMillisThreshold; + return database.name().startsWith(TEST_DATABASE_NAME_PREFIX) && + database.createTime().toEpochMilli() <= creationTimeMillisThreshold; } } diff --git a/testing/trino-product-tests/src/main/java/io/trino/tests/product/deltalake/util/DeltaLakeTestUtils.java b/testing/trino-product-tests/src/main/java/io/trino/tests/product/deltalake/util/DeltaLakeTestUtils.java index 4d84b362acf7..e6bfa6dd68ba 100644 --- a/testing/trino-product-tests/src/main/java/io/trino/tests/product/deltalake/util/DeltaLakeTestUtils.java +++ b/testing/trino-product-tests/src/main/java/io/trino/tests/product/deltalake/util/DeltaLakeTestUtils.java @@ -13,7 +13,6 @@ */ package io.trino.tests.product.deltalake.util; -import com.amazonaws.services.glue.model.ConcurrentModificationException; import com.amazonaws.services.s3.AmazonS3; import com.amazonaws.services.s3.model.DeleteObjectsRequest; import com.amazonaws.services.s3.model.ObjectListing; @@ -24,6 +23,7 @@ import io.airlift.log.Logger; import io.trino.tempto.query.QueryResult; import org.intellij.lang.annotations.Language; +import software.amazon.awssdk.services.glue.model.ConcurrentModificationException; import java.time.temporal.ChronoUnit; import java.util.List; diff --git a/testing/trino-product-tests/src/main/java/io/trino/tests/product/hive/TestHiveMetastoreClientFactory.java b/testing/trino-product-tests/src/main/java/io/trino/tests/product/hive/TestHiveMetastoreClientFactory.java index c66395b2ed1f..17097d85d269 100644 --- a/testing/trino-product-tests/src/main/java/io/trino/tests/product/hive/TestHiveMetastoreClientFactory.java +++ b/testing/trino-product-tests/src/main/java/io/trino/tests/product/hive/TestHiveMetastoreClientFactory.java @@ -13,7 +13,6 @@ */ package io.trino.tests.product.hive; -import com.google.common.net.HostAndPort; import com.google.inject.Inject; import com.google.inject.name.Named; import io.airlift.units.Duration; @@ -34,8 +33,10 @@ public final class TestHiveMetastoreClientFactory Optional.empty(), Optional.empty(), new Duration(10, SECONDS), + new Duration(10, SECONDS), new NoHiveMetastoreAuthentication(), - "localhost"); + "localhost", + Optional.empty()); @Inject @Named("databases.hive.metastore.host") @@ -50,7 +51,7 @@ public ThriftMetastoreClient createMetastoreClient() { URI metastore = URI.create("thrift://" + metastoreHost + ":" + metastorePort); return thriftMetastoreClientFactory.create( - HostAndPort.fromParts(metastore.getHost(), metastore.getPort()), + metastore, Optional.empty()); } } diff --git a/testing/trino-product-tests/src/main/java/io/trino/tests/product/iceberg/TestIcebergHiveMetadataListing.java b/testing/trino-product-tests/src/main/java/io/trino/tests/product/iceberg/TestIcebergHiveMetadataListing.java index 7564761a75a0..f2e4c4e6c87e 100644 --- a/testing/trino-product-tests/src/main/java/io/trino/tests/product/iceberg/TestIcebergHiveMetadataListing.java +++ b/testing/trino-product-tests/src/main/java/io/trino/tests/product/iceberg/TestIcebergHiveMetadataListing.java @@ -22,7 +22,6 @@ import java.util.List; -import static com.google.common.collect.Iterators.getOnlyElement; import static io.trino.tempto.assertions.QueryAssert.Row.row; import static io.trino.tests.product.TestGroups.HMS_ONLY; import static io.trino.tests.product.TestGroups.ICEBERG; @@ -34,7 +33,6 @@ public class TestIcebergHiveMetadataListing extends ProductTest { - private String storageTable; private List preexistingTables; private List preexistingColumns; @@ -55,11 +53,6 @@ public void setUp() onTrino().executeQuery("CREATE TABLE iceberg.default.iceberg_table1 (_string VARCHAR, _integer INTEGER)"); onTrino().executeQuery("CREATE MATERIALIZED VIEW iceberg.default.iceberg_materialized_view AS " + "SELECT * FROM iceberg.default.iceberg_table1"); - storageTable = getOnlyElement(onTrino().executeQuery("SHOW TABLES FROM iceberg.default") - .column(1).stream() - .map(String.class::cast) - .filter(name -> name.startsWith("st_")) - .iterator()); onTrino().executeQuery("CREATE TABLE hive.default.hive_table (_double DOUBLE)"); onTrino().executeQuery("CREATE VIEW hive.default.hive_view AS SELECT * FROM hive.default.hive_table"); @@ -84,7 +77,6 @@ public void testTableListing() .addAll(preexistingTables) .add(row("iceberg_table1")) .add(row("iceberg_materialized_view")) - .add(row(storageTable)) .add(row("iceberg_view")) .add(row("hive_table")) // Iceberg connector supports Trino views created via Hive connector @@ -104,8 +96,6 @@ public void testColumnListing() .add(row("iceberg_table1", "_integer")) .add(row("iceberg_materialized_view", "_string")) .add(row("iceberg_materialized_view", "_integer")) - .add(row(storageTable, "_string")) - .add(row(storageTable, "_integer")) .add(row("iceberg_view", "_string")) .add(row("iceberg_view", "_integer")) .add(row("hive_view", "_double")) diff --git a/testing/trino-product-tests/src/main/java/io/trino/tests/product/iceberg/TestIcebergSparkCompatibility.java b/testing/trino-product-tests/src/main/java/io/trino/tests/product/iceberg/TestIcebergSparkCompatibility.java index 9d51741b127e..a29baad6f7ea 100644 --- a/testing/trino-product-tests/src/main/java/io/trino/tests/product/iceberg/TestIcebergSparkCompatibility.java +++ b/testing/trino-product-tests/src/main/java/io/trino/tests/product/iceberg/TestIcebergSparkCompatibility.java @@ -1899,6 +1899,26 @@ public void testOptimizeOnV2IcebergTable() .containsOnly(row(1, 2), row(2, 2), row(3, 2), row(11, 12), row(12, 12), row(13, 12)); } + @Test(groups = {ICEBERG, PROFILE_SPECIFIC_TESTS}) + public void testOptimizeManifests() + { + String tableName = "test_optimize_manifests_" + randomNameSuffix(); + String sparkTableName = sparkTableName(tableName); + String trinoTableName = trinoTableName(tableName); + + onSpark().executeQuery("CREATE TABLE " + sparkTableName + "(a INT) USING ICEBERG"); + onSpark().executeQuery("INSERT INTO " + sparkTableName + " VALUES (1)"); + onSpark().executeQuery("INSERT INTO " + sparkTableName + " VALUES (2)"); + + onTrino().executeQuery("ALTER TABLE " + trinoTableName + " EXECUTE optimize_manifests"); + assertThat(onTrino().executeQuery("SELECT * FROM " + trinoTableName)) + .containsOnly(row(1), row(2)); + assertThat(onSpark().executeQuery("SELECT * FROM " + sparkTableName)) + .containsOnly(row(1), row(2)); + + onSpark().executeQuery("DROP TABLE " + sparkTableName); + } + @Test(groups = {ICEBERG, PROFILE_SPECIFIC_TESTS}) public void testAlterTableExecuteProceduresOnEmptyTable() { diff --git a/testing/trino-testing-containers/src/main/java/io/trino/testing/containers/IcebergRestCatalogBackendContainer.java b/testing/trino-testing-containers/src/main/java/io/trino/testing/containers/IcebergRestCatalogBackendContainer.java new file mode 100644 index 000000000000..a66e676a33df --- /dev/null +++ b/testing/trino-testing-containers/src/main/java/io/trino/testing/containers/IcebergRestCatalogBackendContainer.java @@ -0,0 +1,57 @@ +/* + * 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.testing.containers; + +import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableSet; +import org.testcontainers.containers.Network; + +import java.util.Optional; + +import static io.trino.testing.containers.Minio.MINIO_REGION; + +public class IcebergRestCatalogBackendContainer + extends BaseTestContainer +{ + public IcebergRestCatalogBackendContainer( + Optional network, + String warehouseLocation, + String minioAccessKey, + String minioSecretKey, + String minioSessionToken) + { + super( + "tabulario/iceberg-rest:0.12.0", + "iceberg-rest", + ImmutableSet.of(8181), + ImmutableMap.of(), + ImmutableMap.of( + "CATALOG_INCLUDE__CREDENTIALS", "true", + "CATALOG_WAREHOUSE", warehouseLocation, + "CATALOG_IO__IMPL", "org.apache.iceberg.aws.s3.S3FileIO", + "AWS_REGION", MINIO_REGION, + "CATALOG_S3_ACCESS__KEY__ID", minioAccessKey, + "CATALOG_S3_SECRET__ACCESS__KEY", minioSecretKey, + "CATALOG_S3_SESSION__TOKEN", minioSessionToken, + "CATALOG_S3_ENDPOINT", "http://minio:4566", + "CATALOG_S3_PATH__STYLE__ACCESS", "true"), + network, + 5); + } + + public String getRestCatalogEndpoint() + { + return getMappedHostAndPortForExposedPort(8181).toString(); + } +} diff --git a/testing/trino-testing-containers/src/main/java/io/trino/testing/containers/Minio.java b/testing/trino-testing-containers/src/main/java/io/trino/testing/containers/Minio.java index 624d7c5f9ac4..76c2b2d18262 100644 --- a/testing/trino-testing-containers/src/main/java/io/trino/testing/containers/Minio.java +++ b/testing/trino-testing-containers/src/main/java/io/trino/testing/containers/Minio.java @@ -41,7 +41,7 @@ public class Minio { private static final Logger log = Logger.get(Minio.class); - public static final String DEFAULT_IMAGE = "minio/minio:RELEASE.2023-05-18T00-05-36Z"; + public static final String DEFAULT_IMAGE = "minio/minio:RELEASE.2024-12-18T13-15-44Z"; public static final String DEFAULT_HOST_NAME = "minio"; public static final int MINIO_API_PORT = 4566; diff --git a/testing/trino-testing/src/main/java/io/trino/testing/AbstractTestQueryFramework.java b/testing/trino-testing/src/main/java/io/trino/testing/AbstractTestQueryFramework.java index 1bd83cf7564b..4aec2519b733 100644 --- a/testing/trino-testing/src/main/java/io/trino/testing/AbstractTestQueryFramework.java +++ b/testing/trino-testing/src/main/java/io/trino/testing/AbstractTestQueryFramework.java @@ -18,6 +18,7 @@ import com.google.errorprone.annotations.CanIgnoreReturnValue; import io.airlift.units.Duration; import io.trino.Session; +import io.trino.cost.StatsAndCosts; import io.trino.execution.QueryInfo; import io.trino.execution.QueryManager; import io.trino.execution.QueryState; @@ -29,6 +30,8 @@ import io.trino.execution.warnings.WarningCollector; import io.trino.memory.LocalMemoryManager; import io.trino.memory.MemoryPool; +import io.trino.metadata.FunctionManager; +import io.trino.metadata.Metadata; import io.trino.metadata.QualifiedObjectName; import io.trino.metadata.TableHandle; import io.trino.operator.OperatorStats; @@ -78,6 +81,7 @@ import static io.trino.sql.ParsingUtil.createParsingOptions; import static io.trino.sql.SqlFormatter.formatSql; import static io.trino.sql.planner.OptimizerConfig.JoinReorderingStrategy; +import static io.trino.sql.planner.planprinter.PlanPrinter.textLogicalPlan; import static io.trino.testing.assertions.Assert.assertEventually; import static io.trino.transaction.TransactionBuilder.transaction; import static java.lang.String.format; @@ -475,6 +479,13 @@ private void executeExclusively(Session session, @Language("SQL") String sql, Te }); } + protected String formatPlan(Session session, Plan plan) + { + Metadata metadata = getDistributedQueryRunner().getCoordinator().getMetadata(); + FunctionManager functionManager = getDistributedQueryRunner().getCoordinator().getFunctionManager(); + return inTransaction(session, transactionSession -> textLogicalPlan(plan.getRoot(), plan.getTypes(), metadata, functionManager, StatsAndCosts.empty(), transactionSession, 0, false)); + } + protected void assertAccessDenied(@Language("SQL") String sql, @Language("RegExp") String exceptionsMessageRegExp, TestingPrivilege... deniedPrivileges) { assertAccessDenied(getSession(), sql, exceptionsMessageRegExp, deniedPrivileges); diff --git a/testing/trino-testing/src/main/java/io/trino/testing/BaseConnectorTest.java b/testing/trino-testing/src/main/java/io/trino/testing/BaseConnectorTest.java index d566126a54aa..49d7bb00ee1f 100644 --- a/testing/trino-testing/src/main/java/io/trino/testing/BaseConnectorTest.java +++ b/testing/trino-testing/src/main/java/io/trino/testing/BaseConnectorTest.java @@ -4652,7 +4652,7 @@ public void testUpdate() } // Repeat test with invocationCount for better test coverage, since the tested aspect is inherently non-deterministic. - @Test(timeOut = 60_000, invocationCount = 4) + @Test(timeOut = 60_000, invocationCount = 1) public void testUpdateRowConcurrently() throws Exception { @@ -4661,7 +4661,7 @@ public void testUpdateRowConcurrently() return; } - int threads = 4; + int threads = 2; CyclicBarrier barrier = new CyclicBarrier(threads); ExecutorService executor = newFixedThreadPool(threads); try (TestTable table = new TestTable( @@ -6244,7 +6244,7 @@ public void testProjectionPushdown() .isFullyPushedDown(); // With Projection Pushdown disabled - Session sessionWithoutPushdown = sessionWithProjectionPushdownDisabled(); + Session sessionWithoutPushdown = sessionWithProjectionPushdownDisabled(getSession()); assertThat(query(sessionWithoutPushdown, selectQuery)) .matches(expectedResult) .isNotFullyPushedDown(ProjectNode.class); @@ -6368,10 +6368,11 @@ public void testProjectionPushdownReadsLessData() "AS SELECT val AS id, CAST(ROW(val + 1, val + 2) AS ROW(leaf1 BIGINT, leaf2 BIGINT)) AS root FROM UNNEST(SEQUENCE(1, 10)) AS t(val)")) { MaterializedResult expectedResult = computeActual("SELECT val + 2 FROM UNNEST(SEQUENCE(1, 10)) AS t(val)"); String selectQuery = "SELECT root.leaf2 FROM " + testTable.getName(); - Session sessionWithoutPushdown = sessionWithProjectionPushdownDisabled(); + Session sessionWithoutSmallFileThreshold = withoutSmallFileThreshold(getSession()); + Session sessionWithoutPushdown = sessionWithProjectionPushdownDisabled(sessionWithoutSmallFileThreshold); assertQueryStats( - getSession(), + sessionWithoutSmallFileThreshold, selectQuery, statsWithPushdown -> { DataSize physicalInputDataSizeWithPushdown = statsWithPushdown.getPhysicalInputDataSize(); @@ -6405,12 +6406,13 @@ public void testProjectionPushdownPhysicalInputSize() "test_projection_pushdown_physical_input_size_", "AS SELECT val AS id, CAST(ROW(val + 1, val + 2) AS ROW(leaf1 BIGINT, leaf2 BIGINT)) AS root FROM UNNEST(SEQUENCE(1, 10)) AS t(val)")) { // Verify that the physical input size is smaller when reading the root.leaf1 field compared to reading the root field + Session sessionWithoutSmallFileThreshold = withoutSmallFileThreshold(getSession()); assertQueryStats( - getSession(), + sessionWithoutSmallFileThreshold, "SELECT root FROM " + testTable.getName(), statsWithSelectRootField -> { assertQueryStats( - getSession(), + sessionWithoutSmallFileThreshold, "SELECT root.leaf1 FROM " + testTable.getName(), statsWithSelectLeafField -> { if (supportsPhysicalPushdown()) { @@ -6427,11 +6429,11 @@ public void testProjectionPushdownPhysicalInputSize() // Verify that the physical input size is the same when reading the root field compared to reading both the root and root.leaf1 fields assertQueryStats( - getSession(), + sessionWithoutSmallFileThreshold, "SELECT root FROM " + testTable.getName(), statsWithSelectRootField -> { assertQueryStats( - getSession(), + sessionWithoutSmallFileThreshold, "SELECT root, root.leaf1 FROM " + testTable.getName(), statsWithSelectRootAndLeafField -> { assertThat(statsWithSelectRootAndLeafField.getPhysicalInputDataSize()).isEqualTo(statsWithSelectRootField.getPhysicalInputDataSize()); @@ -6486,13 +6488,18 @@ protected boolean supportsPhysicalPushdown() return true; } - protected Session sessionWithProjectionPushdownDisabled() + protected Session sessionWithProjectionPushdownDisabled(Session session) { - return Session.builder(getSession()) + return Session.builder(session) .setCatalogSessionProperty(getSession().getCatalog().orElseThrow(), "projection_pushdown_enabled", "false") .build(); } + protected Session withoutSmallFileThreshold(Session session) + { + throw new UnsupportedOperationException(); + } + protected static final class DataMappingTestSetup { private final String trinoTypeName; diff --git a/testing/trino-testing/src/main/java/io/trino/testing/DistributedQueryRunner.java b/testing/trino-testing/src/main/java/io/trino/testing/DistributedQueryRunner.java index 28c7553364f4..52203e4ba91b 100644 --- a/testing/trino-testing/src/main/java/io/trino/testing/DistributedQueryRunner.java +++ b/testing/trino-testing/src/main/java/io/trino/testing/DistributedQueryRunner.java @@ -78,6 +78,8 @@ import static com.google.common.base.MoreObjects.firstNonNull; import static com.google.common.base.Preconditions.checkArgument; +import static com.google.common.base.Preconditions.checkState; +import static com.google.common.base.Throwables.throwIfInstanceOf; import static com.google.common.base.Throwables.throwIfUnchecked; import static com.google.inject.util.Modules.EMPTY_MODULE; import static io.airlift.log.Level.DEBUG; @@ -200,6 +202,22 @@ private DistributedQueryRunner( waitForAllNodesGloballyVisible(); } + public void registerResource(AutoCloseable resource) + { + requireNonNull(resource, "resource is null"); + checkState(!closed, "already closed"); + closer.register(() -> { + try { + resource.close(); + } + catch (Exception e) { + throwIfUnchecked(e); + throwIfInstanceOf(e, IOException.class); + throw new RuntimeException(e); + } + }); + } + private TestingTrinoServer createServer( boolean coordinator, Map extraCoordinatorProperties, diff --git a/testing/trino-testing/src/main/java/io/trino/testing/MultisetAssertions.java b/testing/trino-testing/src/main/java/io/trino/testing/MultisetAssertions.java new file mode 100644 index 000000000000..9505068907bb --- /dev/null +++ b/testing/trino-testing/src/main/java/io/trino/testing/MultisetAssertions.java @@ -0,0 +1,54 @@ +/* + * 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.testing; + +import com.google.common.collect.Multiset; +import com.google.common.collect.Sets; + +import java.util.List; +import java.util.stream.Stream; + +import static com.google.common.collect.ImmutableList.toImmutableList; +import static java.lang.String.format; +import static java.lang.String.join; +import static org.assertj.core.api.Fail.fail; + +public final class MultisetAssertions +{ + private MultisetAssertions() {} + + public static void assertMultisetsEqual(Multiset actual, Multiset expected) + { + if (expected.equals(actual)) { + return; + } + + List mismatchReport = Sets.union(expected.elementSet(), actual.elementSet()).stream() + .filter(key -> expected.count(key) != actual.count(key)) + .flatMap(key -> { + int expectedCount = expected.count(key); + int actualCount = actual.count(key); + if (actualCount < expectedCount) { + return Stream.of(format("%s more occurrences of %s", expectedCount - actualCount, key)); + } + if (actualCount > expectedCount) { + return Stream.of(format("%s fewer occurrences of %s", actualCount - expectedCount, key)); + } + return Stream.of(); + }) + .collect(toImmutableList()); + + fail("Expected: \n\t\t" + join(",\n\t\t", mismatchReport)); + } +} diff --git a/testing/trino-testing/src/test/java/io/trino/testing/TestH2QueryRunner.java b/testing/trino-testing/src/test/java/io/trino/testing/TestH2QueryRunner.java index 2ed5f3aa1f55..e02809b195f7 100644 --- a/testing/trino-testing/src/test/java/io/trino/testing/TestH2QueryRunner.java +++ b/testing/trino-testing/src/test/java/io/trino/testing/TestH2QueryRunner.java @@ -44,7 +44,7 @@ public void close() h2QueryRunner = null; } - @Test + @Test(enabled = false) public void testDateToTimestampCoercion() { // allow running tests with a connector that supports TIMESTAMP but not DATE diff --git a/testing/trino-tests/pom.xml b/testing/trino-tests/pom.xml index 69c984232580..2111543bd93f 100644 --- a/testing/trino-tests/pom.xml +++ b/testing/trino-tests/pom.xml @@ -195,6 +195,18 @@ test + + io.trino + trino-filesystem + test + + + + io.trino + trino-filesystem-s3 + test + + io.trino trino-hdfs diff --git a/testing/trino-tests/src/test/java/io/trino/sql/planner/BaseHiveCostBasedPlanTest.java b/testing/trino-tests/src/test/java/io/trino/sql/planner/BaseHiveCostBasedPlanTest.java index c627ac63fe83..c154599ac8d4 100644 --- a/testing/trino-tests/src/test/java/io/trino/sql/planner/BaseHiveCostBasedPlanTest.java +++ b/testing/trino-tests/src/test/java/io/trino/sql/planner/BaseHiveCostBasedPlanTest.java @@ -24,6 +24,7 @@ import java.io.File; import java.net.URL; +import java.nio.file.Path; import java.nio.file.Paths; import java.util.Arrays; import java.util.Optional; @@ -59,7 +60,7 @@ protected ConnectorFactory createConnectorFactory() RecordingHiveMetastore metastore = new RecordingHiveMetastore( new UnimplementedHiveMetastore(), new HiveMetastoreRecording(recordingConfig, createJsonCodec())); - return new TestingHiveConnectorFactory(metastore); + return new TestingHiveConnectorFactory(Path.of(metadataDir), Optional.of(metastore)); } private static String getSchema(String metadataDir) diff --git a/testing/trino-tests/src/test/java/io/trino/sql/planner/BaseIcebergCostBasedPlanTest.java b/testing/trino-tests/src/test/java/io/trino/sql/planner/BaseIcebergCostBasedPlanTest.java index a76924bccaef..529cad734be8 100644 --- a/testing/trino-tests/src/test/java/io/trino/sql/planner/BaseIcebergCostBasedPlanTest.java +++ b/testing/trino-tests/src/test/java/io/trino/sql/planner/BaseIcebergCostBasedPlanTest.java @@ -17,13 +17,11 @@ import com.google.common.collect.ImmutableMap; import com.google.errorprone.annotations.concurrent.GuardedBy; import io.airlift.log.Logger; -import io.trino.hdfs.DynamicHdfsConfiguration; -import io.trino.hdfs.HdfsConfig; -import io.trino.hdfs.HdfsConfigurationInitializer; -import io.trino.hdfs.HdfsEnvironment; -import io.trino.hdfs.authentication.NoHdfsAuthentication; -import io.trino.hdfs.s3.HiveS3Config; -import io.trino.hdfs.s3.TrinoS3ConfigurationInitializer; +import io.airlift.units.DataSize; +import io.opentelemetry.api.OpenTelemetry; +import io.trino.filesystem.s3.S3FileSystemConfig; +import io.trino.filesystem.s3.S3FileSystemFactory; +import io.trino.filesystem.s3.S3FileSystemStats; import io.trino.plugin.hive.NodeVersion; import io.trino.plugin.hive.metastore.Database; import io.trino.plugin.hive.metastore.Table; @@ -47,7 +45,6 @@ import java.util.List; import java.util.Map; import java.util.Optional; -import java.util.Set; import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.base.Verify.verify; @@ -118,23 +115,19 @@ protected ConnectorFactory createConnectorFactory() throw new UncheckedIOException(e); } - HdfsConfig hdfsConfig = new HdfsConfig(); - HiveS3Config s3Config = new HiveS3Config() - .setS3Endpoint(minio.getMinioAddress()) - .setS3AwsAccessKey(MINIO_ACCESS_KEY) - .setS3AwsSecretKey(MINIO_SECRET_KEY) - .setS3PathStyleAccess(true); - HdfsEnvironment hdfsEnvironment = new HdfsEnvironment( - new DynamicHdfsConfiguration( - new HdfsConfigurationInitializer(hdfsConfig, Set.of(new TrinoS3ConfigurationInitializer(s3Config))), - Set.of()), - hdfsConfig, - new NoHdfsAuthentication()); + S3FileSystemFactory s3FileSystemFactory = + new S3FileSystemFactory( + OpenTelemetry.noop(), + new S3FileSystemConfig() + .setAwsAccessKey(MINIO_ACCESS_KEY) + .setAwsSecretKey(MINIO_SECRET_KEY) + .setEndpoint(minio.getMinioAddress()) + .setStreamingPartSize(DataSize.valueOf("5.5MB")), new S3FileSystemStats()); fileMetastore = new FileHiveMetastore( // Must match the version picked by the LocalQueryRunner new NodeVersion(""), - hdfsEnvironment, + s3FileSystemFactory, false, new FileHiveMetastoreConfig() .setCatalogDirectory(temporaryMetastoreDirectory.toString())); diff --git a/testing/trino-tests/src/test/java/io/trino/sql/planner/HiveMetadataRecorder.java b/testing/trino-tests/src/test/java/io/trino/sql/planner/HiveMetadataRecorder.java index 5791e87af4dd..5ca9fb83f410 100644 --- a/testing/trino-tests/src/test/java/io/trino/sql/planner/HiveMetadataRecorder.java +++ b/testing/trino-tests/src/test/java/io/trino/sql/planner/HiveMetadataRecorder.java @@ -77,7 +77,7 @@ private static DistributedQueryRunner createQueryRunner(String schema, String co String recordingPath = getResourcePath(format("%s/%s.json.gz", recordingDir, schema)); Path.of(recordingPath).toFile().getParentFile().mkdirs(); - queryRunner.installPlugin(new TestingHivePlugin()); + queryRunner.installPlugin(new TestingHivePlugin(Path.of(recordingPath))); queryRunner.createCatalog("hive", "hive", configBuilder .putAll(loadPropertiesFrom(configPath)) .put("hive.metastore-recording-path", recordingPath)