com.facebook.presto.redis.RedisRecordCursor Maven / Gradle / Ivy
/*
* 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.redis;
import com.facebook.airlift.log.Logger;
import com.facebook.presto.common.type.Type;
import com.facebook.presto.decoder.DecoderColumnHandle;
import com.facebook.presto.decoder.FieldValueProvider;
import com.facebook.presto.decoder.RowDecoder;
import com.facebook.presto.spi.ColumnHandle;
import com.facebook.presto.spi.RecordCursor;
import io.airlift.slice.Slice;
import redis.clients.jedis.Jedis;
import redis.clients.jedis.JedisPool;
import redis.clients.jedis.ScanParams;
import redis.clients.jedis.ScanResult;
import java.nio.charset.StandardCharsets;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.concurrent.atomic.AtomicBoolean;
import static com.facebook.presto.decoder.FieldValueProviders.booleanValueProvider;
import static com.facebook.presto.decoder.FieldValueProviders.bytesValueProvider;
import static com.facebook.presto.decoder.FieldValueProviders.longValueProvider;
import static com.google.common.base.Preconditions.checkArgument;
import static java.lang.String.format;
import static redis.clients.jedis.ScanParams.SCAN_POINTER_START;
public class RedisRecordCursor
implements RecordCursor
{
private static final Logger log = Logger.get(RedisRecordCursor.class);
private static final byte[] EMPTY_BYTE_ARRAY = new byte[0];
private final RowDecoder keyDecoder;
private final RowDecoder valueDecoder;
private final RedisSplit split;
private final List columnHandles;
private final RedisJedisManager redisJedisManager;
private final JedisPool jedisPool;
private final ScanParams scanParams;
private ScanResult redisCursor;
private Iterator keysIterator;
private final AtomicBoolean reported = new AtomicBoolean();
private FieldValueProvider[] fieldValueProviders;
private String valueString;
private Map valueMap;
private long totalBytes;
private long totalValues;
private final FieldValueProvider[] currentRowValues;
RedisRecordCursor(
RowDecoder keyDecoder,
RowDecoder valueDecoder,
RedisSplit split,
List columnHandles,
RedisJedisManager redisJedisManager)
{
this.keyDecoder = keyDecoder;
this.valueDecoder = valueDecoder;
this.split = split;
this.columnHandles = columnHandles;
this.redisJedisManager = redisJedisManager;
this.jedisPool = redisJedisManager.getJedisPool(split.getNodes().get(0));
this.scanParams = setScanParams();
this.currentRowValues = new FieldValueProvider[columnHandles.size()];
fetchKeys();
}
@Override
public long getCompletedBytes()
{
return totalBytes;
}
@Override
public long getReadTimeNanos()
{
return 0;
}
@Override
public Type getType(int field)
{
checkArgument(field < columnHandles.size(), "Invalid field index");
return columnHandles.get(field).getType();
}
public boolean hasUnscannedData()
{
if (redisCursor == null) {
return false;
}
// no more keys are unscanned when
// when redis scan command
// returns 0 string cursor
return (!redisCursor.getStringCursor().equals("0"));
}
@Override
public boolean advanceNextPosition()
{
while (!keysIterator.hasNext()) {
if (!hasUnscannedData()) {
return endOfData();
}
fetchKeys();
}
return nextRow(keysIterator.next());
}
private boolean endOfData()
{
if (!reported.getAndSet(true)) {
log.debug("Read a total of %d values with %d bytes.", totalValues, totalBytes);
}
return false;
}
private boolean nextRow(String keyString)
{
fetchData(keyString);
byte[] keyData = keyString.getBytes(StandardCharsets.UTF_8);
byte[] valueData = EMPTY_BYTE_ARRAY;
if (valueString != null) {
valueData = valueString.getBytes(StandardCharsets.UTF_8);
}
totalBytes += valueData.length;
totalValues++;
Optional> decodedKey = keyDecoder.decodeRow(
keyData,
null);
Optional> decodedValue = valueDecoder.decodeRow(
valueData,
valueMap);
Map currentRowValuesMap = new HashMap<>();
for (DecoderColumnHandle columnHandle : columnHandles) {
if (columnHandle.isInternal()) {
RedisInternalFieldDescription fieldDescription = RedisInternalFieldDescription.forColumnName(columnHandle.getName());
switch (fieldDescription) {
case KEY_FIELD:
currentRowValuesMap.put(columnHandle, bytesValueProvider(keyData));
break;
case VALUE_FIELD:
currentRowValuesMap.put(columnHandle, bytesValueProvider(valueData));
break;
case KEY_LENGTH_FIELD:
currentRowValuesMap.put(columnHandle, longValueProvider(keyData.length));
break;
case VALUE_LENGTH_FIELD:
currentRowValuesMap.put(columnHandle, longValueProvider(valueData.length));
break;
case KEY_CORRUPT_FIELD:
currentRowValuesMap.put(columnHandle, booleanValueProvider(!decodedKey.isPresent()));
break;
case VALUE_CORRUPT_FIELD:
currentRowValuesMap.put(columnHandle, booleanValueProvider(!decodedValue.isPresent()));
break;
default:
throw new IllegalArgumentException("unknown internal field " + fieldDescription);
}
}
}
decodedKey.ifPresent(currentRowValuesMap::putAll);
decodedValue.ifPresent(currentRowValuesMap::putAll);
for (int i = 0; i < columnHandles.size(); i++) {
ColumnHandle columnHandle = columnHandles.get(i);
currentRowValues[i] = currentRowValuesMap.get(columnHandle);
}
return true;
}
@Override
public boolean getBoolean(int field)
{
return getFieldValueProvider(field, boolean.class).getBoolean();
}
@Override
public long getLong(int field)
{
return getFieldValueProvider(field, long.class).getLong();
}
@Override
public double getDouble(int field)
{
return getFieldValueProvider(field, double.class).getDouble();
}
@Override
public Slice getSlice(int field)
{
return getFieldValueProvider(field, Slice.class).getSlice();
}
@Override
public boolean isNull(int field)
{
checkArgument(field < columnHandles.size(), "Invalid field index");
return currentRowValues == null || currentRowValues[field].isNull();
}
@Override
public Object getObject(int field)
{
checkArgument(field < columnHandles.size(), "Invalid field index");
throw new IllegalArgumentException(format("Type %s is not supported", getType(field)));
}
private FieldValueProvider getFieldValueProvider(int field, Class expectedType)
{
checkArgument(field < columnHandles.size(), "Invalid field index");
checkFieldType(field, expectedType);
return currentRowValues[field];
}
private void checkFieldType(int field, Class expected)
{
Class actual = getType(field).getJavaType();
checkArgument(actual == expected, "Expected field %s to be type %s but is %s", field, expected, actual);
}
@Override
public void close()
{
}
private ScanParams setScanParams()
{
if (split.getKeyDataType() == RedisDataType.STRING) {
ScanParams scanParams = new ScanParams();
scanParams.count(redisJedisManager.getRedisConnectorConfig().getRedisScanCount());
// when Redis key string follows "schema:table:*" format
// scan command can efficiently query tables
// by returning matching keys
// the alternative is to set key-prefix-schema-table to false
// and treat entire redis as single schema , single table
// redis Hash/Set types are to be supported - they can also be
// used to filter out table data
// "default" schema is not prefixed to the key
if (redisJedisManager.getRedisConnectorConfig().isKeyPrefixSchemaTable()) {
String keyMatch = "";
if (!split.getSchemaName().equals("default")) {
keyMatch = split.getSchemaName() + redisJedisManager.getRedisConnectorConfig().getRedisKeyDelimiter();
}
keyMatch = keyMatch + split.getTableName() + redisJedisManager.getRedisConnectorConfig().getRedisKeyDelimiter() + "*";
scanParams.match(keyMatch);
}
return scanParams;
}
return null;
}
// Redis keys can be contained in the user-provided ZSET
// Otherwise they need to be found by scanning Redis
private boolean fetchKeys()
{
try (Jedis jedis = jedisPool.getResource()) {
switch (split.getKeyDataType()) {
case STRING: {
String cursor = SCAN_POINTER_START;
if (redisCursor != null) {
cursor = redisCursor.getStringCursor();
}
log.debug("Scanning new Redis keys from cursor %s . %d values read so far", cursor, totalValues);
redisCursor = jedis.scan(cursor, scanParams);
List keys = redisCursor.getResult();
keysIterator = keys.iterator();
}
break;
case ZSET:
Set keys = jedis.zrange(split.getKeyName(), split.getStart(), split.getEnd());
keysIterator = keys.iterator();
break;
default:
log.debug("Redis type of key %s is unsupported", split.getKeyDataFormat());
return false;
}
}
return true;
}
private boolean fetchData(String keyString)
{
valueString = null;
valueMap = null;
// Redis connector supports two types of Redis
// values: STRING and HASH
// HASH types requires hash row decoder to
// fill in the columns
// whereas for the STRING type decoders are optional
try (Jedis jedis = jedisPool.getResource()) {
switch (split.getValueDataType()) {
case STRING:
valueString = jedis.get(keyString);
if (valueString == null) {
log.warn("Redis data modified while query was running, string value at key %s deleted", keyString);
return false;
}
break;
case HASH:
valueMap = jedis.hgetAll(keyString);
if (valueMap == null) {
log.warn("Redis data modified while query was running, hash value at key %s deleted", keyString);
return false;
}
break;
default:
log.debug("Redis type for key %s is unsupported", keyString);
return false;
}
}
return true;
}
}
© 2015 - 2025 Weber Informatics LLC | Privacy Policy