Skip to content

Commit

Permalink
POST the metadata results to worker
Browse files Browse the repository at this point in the history
  • Loading branch information
Nikhil Collooru authored and highker committed Oct 1, 2020
1 parent c2b2b0d commit a786699
Show file tree
Hide file tree
Showing 19 changed files with 693 additions and 16 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -126,6 +126,9 @@ public void configure(Binder binder)
binder.bind(HiveWriterStats.class).in(Scopes.SINGLETON);
newExporter(binder).export(HiveWriterStats.class).as(generatedNameOf(HiveWriterStats.class, connectorId));

binder.bind(HiveFileRenamer.class).in(Scopes.SINGLETON);
newExporter(binder).export(HiveFileRenamer.class).as(generatedNameOf(HiveFileRenamer.class, connectorId));

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);
Expand Down
Original file line number Diff line number Diff line change
@@ -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 com.facebook.presto.hive;

import com.facebook.presto.spi.ConnectorMetadataUpdateHandle;
import com.facebook.presto.spi.QueryId;
import com.facebook.presto.spi.SchemaTableName;
import com.google.common.collect.ImmutableList;
import org.weakref.jmx.Managed;

import javax.annotation.PreDestroy;

import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicLong;

import static com.google.common.base.Verify.verify;

public class HiveFileRenamer
{
private final Map<QueryId, Map<HiveMetadataUpdateKey, AtomicLong>> queryPartitionFileCounterMap = new ConcurrentHashMap<>();
private final Map<QueryId, Map<HiveMetadataUpdateHandle, String>> queryHiveMetadataResultMap = new ConcurrentHashMap<>();

public List<ConnectorMetadataUpdateHandle> getMetadataUpdateResults(List<ConnectorMetadataUpdateHandle> metadataUpdateRequests, QueryId queryId)
{
ImmutableList.Builder<ConnectorMetadataUpdateHandle> metadataUpdateResults = ImmutableList.builder();

for (ConnectorMetadataUpdateHandle connectorMetadataUpdateHandle : metadataUpdateRequests) {
HiveMetadataUpdateHandle request = (HiveMetadataUpdateHandle) connectorMetadataUpdateHandle;
String fileName = getFileName(request, queryId);
metadataUpdateResults.add(new HiveMetadataUpdateHandle(request.getRequestId(), request.getSchemaTableName(), request.getPartitionName(), Optional.of(fileName)));
}
return metadataUpdateResults.build();
}

public void cleanup(QueryId queryId)
{
queryPartitionFileCounterMap.remove(queryId);
queryHiveMetadataResultMap.remove(queryId);
}

private String getFileName(HiveMetadataUpdateHandle request, QueryId queryId)
{
if (!queryPartitionFileCounterMap.containsKey(queryId) || !queryHiveMetadataResultMap.containsKey(queryId)) {
queryPartitionFileCounterMap.putIfAbsent(queryId, new ConcurrentHashMap<>());
queryHiveMetadataResultMap.putIfAbsent(queryId, new ConcurrentHashMap<>());
}

// To keep track of the file counter per query per partition
Map<HiveMetadataUpdateKey, AtomicLong> partitionFileCounterMap = queryPartitionFileCounterMap.get(queryId);

// To keep track of the file name result per query per request
// This is to make sure that request - fileName mapping is 1:1
Map<HiveMetadataUpdateHandle, String> hiveMetadataResultMap = queryHiveMetadataResultMap.get(queryId);

// If we have seen this request before then directly return the result.
if (hiveMetadataResultMap.containsKey(request)) {
// We come here if for some reason the worker did not receive the fileName and it retried the request.
return hiveMetadataResultMap.get(request);
}

HiveMetadataUpdateKey key = new HiveMetadataUpdateKey(request);
// File names start from 0
partitionFileCounterMap.putIfAbsent(key, new AtomicLong(0));

AtomicLong fileCount = partitionFileCounterMap.get(key);
String fileName = Long.valueOf(fileCount.getAndIncrement()).toString();

// Store the request - fileName mapping
hiveMetadataResultMap.put(request, fileName);

return fileName;
}

@PreDestroy
public void stop()
{
// Mappings should be deleted when query finishes. So verify that map is empty before its closed.
verify(queryPartitionFileCounterMap.isEmpty(), "Query partition file counter map has %s entries left behind", queryPartitionFileCounterMap.size());
verify(queryHiveMetadataResultMap.isEmpty(), "Query hive metadata result map has %s entries left behind", queryHiveMetadataResultMap.size());
}

@Managed
public int getQueryPartitionFileCounterMapSize()
{
return queryPartitionFileCounterMap.size();
}

@Managed
public int getHiveMetadataUpdateResultMapSize()
{
return queryHiveMetadataResultMap.size();
}

private static class HiveMetadataUpdateKey
{
private final SchemaTableName schemaTableName;
private final Optional<String> partitionName;

private HiveMetadataUpdateKey(HiveMetadataUpdateHandle hiveMetadataUpdateHandle)
{
this.schemaTableName = hiveMetadataUpdateHandle.getSchemaTableName();
this.partitionName = hiveMetadataUpdateHandle.getPartitionName();
}

@Override
public boolean equals(Object obj)
{
if (this == obj) {
return true;
}
if (obj == null || getClass() != obj.getClass()) {
return false;
}
HiveMetadataUpdateKey o = (HiveMetadataUpdateKey) obj;
return schemaTableName.equals(o.schemaTableName) &&
partitionName.equals(o.partitionName);
}

@Override
public int hashCode()
{
return Objects.hash(schemaTableName, partitionName);
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,7 @@
import com.facebook.presto.spi.ColumnHandle;
import com.facebook.presto.spi.ColumnMetadata;
import com.facebook.presto.spi.ConnectorInsertTableHandle;
import com.facebook.presto.spi.ConnectorMetadataUpdateHandle;
import com.facebook.presto.spi.ConnectorNewTableLayout;
import com.facebook.presto.spi.ConnectorOutputTableHandle;
import com.facebook.presto.spi.ConnectorSession;
Expand All @@ -61,6 +62,7 @@
import com.facebook.presto.spi.DiscretePredicates;
import com.facebook.presto.spi.InMemoryRecordSet;
import com.facebook.presto.spi.PrestoException;
import com.facebook.presto.spi.QueryId;
import com.facebook.presto.spi.RecordCursor;
import com.facebook.presto.spi.SchemaTableName;
import com.facebook.presto.spi.SchemaTablePrefix;
Expand Down Expand Up @@ -356,6 +358,7 @@ public class HiveMetadata
private final PartitionObjectBuilder partitionObjectBuilder;
private final HiveEncryptionInformationProvider encryptionInformationProvider;
private final HivePartitionStats hivePartitionStats;
private final HiveFileRenamer hiveFileRenamer;

public HiveMetadata(
SemiTransactionalHiveMetastore metastore,
Expand All @@ -380,7 +383,8 @@ public HiveMetadata(
ZeroRowFileCreator zeroRowFileCreator,
PartitionObjectBuilder partitionObjectBuilder,
HiveEncryptionInformationProvider encryptionInformationProvider,
HivePartitionStats hivePartitionStats)
HivePartitionStats hivePartitionStats,
HiveFileRenamer hiveFileRenamer)
{
this.allowCorruptWritesForTesting = allowCorruptWritesForTesting;

Expand All @@ -406,6 +410,7 @@ public HiveMetadata(
this.partitionObjectBuilder = requireNonNull(partitionObjectBuilder, "partitionObjectBuilder is null");
this.encryptionInformationProvider = requireNonNull(encryptionInformationProvider, "encryptionInformationProvider is null");
this.hivePartitionStats = requireNonNull(hivePartitionStats, "hivePartitionStats is null");
this.hiveFileRenamer = requireNonNull(hiveFileRenamer, "hiveFileRenamer is null");
}

public SemiTransactionalHiveMetastore getMetastore()
Expand Down Expand Up @@ -2879,6 +2884,18 @@ public CompletableFuture<Void> commitPageSinkAsync(ConnectorSession session, Con
return toCompletableFuture(stagingFileCommitter.commitFiles(session, handle.getSchemaName(), handle.getTableName(), getPartitionUpdates(fragments)));
}

@Override
public List<ConnectorMetadataUpdateHandle> getMetadataUpdateResults(List<ConnectorMetadataUpdateHandle> metadataUpdateRequests, QueryId queryId)
{
return hiveFileRenamer.getMetadataUpdateResults(metadataUpdateRequests, queryId);
}

@Override
public void doMetadataUpdateCleanup(QueryId queryId)
{
hiveFileRenamer.cleanup(queryId);
}

private List<GrantInfo> buildGrants(SchemaTableName tableName, PrestoPrincipal principal)
{
ImmutableList.Builder<GrantInfo> result = ImmutableList.builder();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,7 @@ public class HiveMetadataFactory
private final PartitionObjectBuilder partitionObjectBuilder;
private final HiveEncryptionInformationProvider encryptionInformationProvider;
private final HivePartitionStats hivePartitionStats;
private final HiveFileRenamer hiveFileRenamer;

@Inject
@SuppressWarnings("deprecation")
Expand All @@ -86,7 +87,8 @@ public HiveMetadataFactory(
NodeVersion nodeVersion,
PartitionObjectBuilder partitionObjectBuilder,
HiveEncryptionInformationProvider encryptionInformationProvider,
HivePartitionStats hivePartitionStats)
HivePartitionStats hivePartitionStats,
HiveFileRenamer hiveFileRenamer)
{
this(
metastore,
Expand Down Expand Up @@ -114,7 +116,8 @@ public HiveMetadataFactory(
nodeVersion.toString(),
partitionObjectBuilder,
encryptionInformationProvider,
hivePartitionStats);
hivePartitionStats,
hiveFileRenamer);
}

public HiveMetadataFactory(
Expand Down Expand Up @@ -143,7 +146,8 @@ public HiveMetadataFactory(
String prestoVersion,
PartitionObjectBuilder partitionObjectBuilder,
HiveEncryptionInformationProvider encryptionInformationProvider,
HivePartitionStats hivePartitionStats)
HivePartitionStats hivePartitionStats,
HiveFileRenamer hiveFileRenamer)
{
this.allowCorruptWritesForTesting = allowCorruptWritesForTesting;
this.skipDeletionForAlter = skipDeletionForAlter;
Expand Down Expand Up @@ -172,6 +176,7 @@ public HiveMetadataFactory(
this.partitionObjectBuilder = requireNonNull(partitionObjectBuilder, "partitionObjectBuilder is null");
this.encryptionInformationProvider = requireNonNull(encryptionInformationProvider, "encryptionInformationProvider is null");
this.hivePartitionStats = requireNonNull(hivePartitionStats, "hivePartitionStats is null");
this.hiveFileRenamer = requireNonNull(hiveFileRenamer, "hiveFileRenamer is null");

if (!allowCorruptWritesForTesting && !timeZone.equals(DateTimeZone.getDefault())) {
log.warn("Hive writes are disabled. " +
Expand Down Expand Up @@ -214,6 +219,7 @@ public HiveMetadata get()
zeroRowFileCreator,
partitionObjectBuilder,
encryptionInformationProvider,
hivePartitionStats);
hivePartitionStats,
hiveFileRenamer);
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -957,7 +957,8 @@ protected final void setup(String databaseName, HiveClientConfig hiveClientConfi
TEST_SERVER_VERSION,
new HivePartitionObjectBuilder(),
new HiveEncryptionInformationProvider(ImmutableList.of()),
new HivePartitionStats());
new HivePartitionStats(),
new HiveFileRenamer());
transactionManager = new HiveTransactionManager();
encryptionInformationProvider = new HiveEncryptionInformationProvider(ImmutableList.of());
splitManager = new HiveSplitManager(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -209,7 +209,8 @@ protected void setup(String host, int port, String databaseName, BiFunction<Hive
new NodeVersion("test_version"),
new HivePartitionObjectBuilder(),
new HiveEncryptionInformationProvider(ImmutableList.of()),
new HivePartitionStats());
new HivePartitionStats(),
new HiveFileRenamer());
transactionManager = new HiveTransactionManager();
splitManager = new HiveSplitManager(
transactionManager,
Expand Down
Loading

0 comments on commit a786699

Please sign in to comment.