public KuduScanEnumerator GetAsyncEnumerator(CancellationToken cancellationToken = default)
{
var partitionPruner = PartitionPruner.Create(
_table.Schema,
_table.PartitionSchema,
_predicates,
_lowerBoundPrimaryKey,
_upperBoundPrimaryKey,
_lowerBoundPartitionKey,
_upperBoundPartitionKey);
return(new KuduScanEnumerator(
_logger,
_client,
_table,
_projectedColumnsPb,
ProjectionSchema,
_orderMode,
ReadMode,
ReplicaSelection,
_isFaultTolerant,
_predicates,
_limit,
CacheBlocks,
_lowerBoundPrimaryKey,
_upperBoundPrimaryKey,
_startTimestamp,
_htTimestamp,
BatchSizeBytes,
partitionPruner,
cancellationToken));
}
/// <summary>
/// Checks the number of tablets and pruner ranges generated for a scan with
/// predicates and optional partition key bounds.
/// </summary>
/// <param name="expectedTablets">The expected number of tablets to satisfy the scan.</param>
/// <param name="expectedPrunerRanges">The expected number of generated partition pruner ranges.</param>
/// <param name="table">The table to scan.</param>
/// <param name="partitions">The partitions of the table.</param>
/// <param name="lowerBoundPartitionKey">An optional lower bound partition key.</param>
/// <param name="upperBoundPartitionKey">An optional upper bound partition key.</param>
/// <param name="predicates">The predicates to apply to the scan.</param>
private async Task CheckPartitionsAsync(
int expectedTablets,
int expectedPrunerRanges,
KuduTable table,
List <Partition> partitions,
byte[] lowerBoundPartitionKey,
byte[] upperBoundPartitionKey,
params KuduPredicate[] predicates)
{
// Partition key bounds can't be applied to the ScanTokenBuilder.
var scanBuilder = _client.NewScanBuilder(table);
foreach (var predicate in predicates)
{
scanBuilder.AddPredicate(predicate);
}
if (lowerBoundPartitionKey != null)
{
scanBuilder.LowerBoundPartitionKeyRaw(lowerBoundPartitionKey);
}
if (upperBoundPartitionKey != null)
{
scanBuilder.ExclusiveUpperBoundPartitionKeyRaw(upperBoundPartitionKey);
}
var pruner = PartitionPruner.Create(scanBuilder);
int scannedPartitions = 0;
foreach (var partition in partitions)
{
if (!pruner.ShouldPruneForTests(partition))
{
scannedPartitions++;
}
}
Assert.Equal(expectedTablets, scannedPartitions);
Assert.Equal(expectedPrunerRanges, pruner.NumRangesRemaining);
// Check that the scan token builder comes up with the same amount.
// The scan token builder does not allow for upper/lower partition keys.
if (lowerBoundPartitionKey == null && upperBoundPartitionKey == null)
{
var tokenBuilder = _client.NewScanTokenBuilder(table);
foreach (var predicate in predicates)
{
tokenBuilder.AddPredicate(predicate);
}
// Check that the number of ScanTokens built for the scan matches.
var tokens = await tokenBuilder.BuildAsync();
Assert.Equal(expectedTablets, tokens.Count);
}
}
public KuduScanEnumerator(
ILogger logger,
KuduClient client,
KuduTable table,
List <ColumnSchemaPB> projectedColumnsPb,
KuduSchema projectionSchema,
OrderMode orderMode,
ReadMode readMode,
ReplicaSelection replicaSelection,
bool isFaultTolerant,
Dictionary <string, KuduPredicate> predicates,
long limit,
bool cacheBlocks,
byte[] startPrimaryKey,
byte[] endPrimaryKey,
long startTimestamp,
long htTimestamp,
int batchSizeBytes,
PartitionPruner partitionPruner,
CancellationToken cancellationToken)
{
_logger = logger;
_client = client;
_table = table;
_partitionPruner = partitionPruner;
_orderMode = orderMode;
_readMode = readMode;
_columns = projectedColumnsPb;
_schema = projectionSchema;
_predicates = predicates;
_replicaSelection = replicaSelection;
_isFaultTolerant = isFaultTolerant;
_limit = limit;
_cacheBlocks = cacheBlocks;
_startPrimaryKey = startPrimaryKey ?? Array.Empty <byte>();
_endPrimaryKey = endPrimaryKey ?? Array.Empty <byte>();
_startTimestamp = startTimestamp;
SnapshotTimestamp = htTimestamp;
_batchSizeBytes = batchSizeBytes;
_scannerId = ByteString.Empty;
_lastPrimaryKey = ByteString.Empty;
_cancellationToken = cancellationToken;
ResourceMetrics = new ResourceMetrics();
// If the partition pruner has pruned all partitions, then the scan can be
// short circuited without contacting any tablet servers.
if (!_partitionPruner.HasMorePartitionKeyRanges)
{
_closed = true;
}
// For READ_YOUR_WRITES scan mode, get the latest observed timestamp
// and store it. Always use this one as the propagated timestamp for
// the duration of the scan to avoid unnecessary wait.
if (readMode == ReadMode.ReadYourWrites)
{
_lowerBoundPropagationTimestamp = client.LastPropagatedTimestamp;
}
}
/// <summary>
/// Counts the partitions touched by a scan with optional primary key bounds.
/// The table is assumed to have three INT8 columns as the primary key.
/// </summary>
/// <param name="expectedTablets">The expected number of tablets to satisfy the scan.</param>
/// <param name="table">The table to scan.</param>
/// <param name="partitions">The partitions of the table.</param>
/// <param name="lowerBoundPrimaryKey">The optional lower bound primary key.</param>
/// <param name="upperBoundPrimaryKey">The optional upper bound primary key.</param>
private async Task CheckPartitionsPrimaryKeyAsync(
int expectedTablets,
KuduTable table,
List <Partition> partitions,
sbyte[] lowerBoundPrimaryKey,
sbyte[] upperBoundPrimaryKey)
{
var scanBuilder = _client.NewScanTokenBuilder(table);
if (lowerBoundPrimaryKey != null)
{
var lower = new PartialRow(table.Schema);
for (int i = 0; i < 3; i++)
{
lower.SetSByte(i, lowerBoundPrimaryKey[i]);
}
scanBuilder.LowerBound(lower);
}
if (upperBoundPrimaryKey != null)
{
var upper = new PartialRow(table.Schema);
for (int i = 0; i < 3; i++)
{
upper.SetSByte(i, upperBoundPrimaryKey[i]);
}
scanBuilder.ExclusiveUpperBound(upper);
}
var pruner = PartitionPruner.Create(scanBuilder);
int scannedPartitions = 0;
foreach (var partition in partitions)
{
if (!pruner.ShouldPruneForTests(partition))
{
scannedPartitions++;
}
}
// Check that the number of ScanTokens built for the scan matches.
var tokens = await scanBuilder.BuildAsync();
Assert.Equal(expectedTablets, scannedPartitions);
Assert.Equal(scannedPartitions, tokens.Count);
Assert.Equal(expectedTablets == 0 ? 0 : 1, pruner.NumRangesRemaining);
}
public async ValueTask <List <KuduScanToken> > BuildAsync(
CancellationToken cancellationToken = default)
{
if (LowerBoundPartitionKey.Length != 0 ||
UpperBoundPartitionKey.Length != 0)
{
throw new ArgumentException(
"Partition key bounds may not be set on KuduScanTokenBuilder");
}
// If the scan is short-circuitable, then return no tokens.
foreach (var predicate in Predicates.Values)
{
if (predicate.Type == PredicateType.None)
{
return(new List <KuduScanToken>());
}
}
var proto = new ScanTokenPB();
if (_includeTableMetadata)
{
// Set the table metadata so that a call to the master is not needed when
// deserializing the token into a scanner.
var tableMetadataPb = new TableMetadataPB
{
TableId = Table.TableId,
TableName = Table.TableName,
NumReplicas = Table.NumReplicas,
Schema = Table.SchemaPb.Schema,
PartitionSchema = Table.SchemaPb.PartitionSchema
};
if (Table.Owner is not null)
{
tableMetadataPb.Owner = Table.Owner;
}
if (Table.Comment is not null)
{
tableMetadataPb.Comment = Table.Comment;
}
if (Table.ExtraConfig.Count > 0)
{
foreach (var kvp in Table.ExtraConfig)
{
tableMetadataPb.ExtraConfigs.Add(kvp.Key, kvp.Value);
}
}
proto.TableMetadata = tableMetadataPb;
// Only include the authz token if the table metadata is included.
// It is returned in the required GetTableSchema request otherwise.
var authzToken = Client.GetAuthzToken(Table.TableId);
if (authzToken is not null)
{
proto.AuthzToken = authzToken;
}
}
else
{
// If we add the table metadata, we don't need to set the old table id
// and table name. It is expected that the creation and use of a scan token
// will be on the same or compatible versions.
proto.TableId = Table.TableId;
proto.TableName = Table.TableName;
}
// Map the column names or indices to actual columns in the table schema.
// If the user did not set either projection, then scan all columns.
var schema = Table.Schema;
if (_includeTableMetadata)
{
// If the table metadata is included, then the column indexes can be
// used instead of duplicating the ColumnSchemaPBs in the serialized
// scan token.
if (ProjectedColumnNames is not null)
{
proto.ProjectedColumnIdx.Capacity = ProjectedColumnNames.Count;
foreach (var columnName in ProjectedColumnNames)
{
var columnIndex = schema.GetColumnIndex(columnName);
proto.ProjectedColumnIdx.Add(columnIndex);
}
}
else if (ProjectedColumnIndexes is not null)
{
proto.ProjectedColumnIdx.AddRange(ProjectedColumnIndexes);
}
else
{
var numColumns = schema.Columns.Count;
proto.ProjectedColumnIdx.Capacity = numColumns;
for (int i = 0; i < numColumns; i++)
{
proto.ProjectedColumnIdx.Add(i);
}
}
}
else
{
if (ProjectedColumnNames is not null)
{
proto.ProjectedColumns.Capacity = ProjectedColumnNames.Count;
foreach (var columnName in ProjectedColumnNames)
{
int columnIndex = schema.GetColumnIndex(columnName);
var columnSchema = Table.SchemaPb.Schema.Columns[columnIndex];
proto.ProjectedColumns.Add(columnSchema);
}
}
else if (ProjectedColumnIndexes is not null)
{
proto.ProjectedColumns.Capacity = ProjectedColumnIndexes.Count;
foreach (var columnIndex in ProjectedColumnIndexes)
{
var columnSchema = Table.SchemaPb.Schema.Columns[columnIndex];
proto.ProjectedColumns.Add(columnSchema);
}
}
else
{
proto.ProjectedColumns.AddRange(Table.SchemaPb.Schema.Columns);
}
}
foreach (var predicate in Predicates.Values)
{
proto.ColumnPredicates.Add(predicate.ToProtobuf());
}
if (LowerBoundPrimaryKey.Length > 0)
{
proto.LowerBoundPrimaryKey = UnsafeByteOperations.UnsafeWrap(LowerBoundPrimaryKey);
}
if (UpperBoundPrimaryKey.Length > 0)
{
proto.UpperBoundPrimaryKey = UnsafeByteOperations.UnsafeWrap(UpperBoundPrimaryKey);
}
proto.Limit = (ulong)Limit;
proto.ReadMode = (Protobuf.ReadMode)ReadMode;
proto.ReplicaSelection = (Protobuf.ReplicaSelection)ReplicaSelection;
// If the last propagated timestamp is set send it with the scan.
long lastPropagatedTimestamp = Client.LastPropagatedTimestamp;
if (lastPropagatedTimestamp != KuduClient.NoTimestamp)
{
proto.PropagatedTimestamp = (ulong)lastPropagatedTimestamp;
}
// If the mode is set to read on snapshot set the snapshot timestamps.
if (ReadMode == ReadMode.ReadAtSnapshot)
{
if (HtTimestamp != KuduClient.NoTimestamp)
{
proto.SnapTimestamp = (ulong)HtTimestamp;
}
if (StartTimestamp != KuduClient.NoTimestamp)
{
proto.SnapStartTimestamp = (ulong)StartTimestamp;
}
}
proto.CacheBlocks = CacheBlocks;
proto.FaultTolerant = IsFaultTolerant;
proto.BatchSizeBytes = (uint)BatchSizeBytes;
// TODO:
//proto.setScanRequestTimeoutMs(scanRequestTimeout);
//proto.setKeepAlivePeriodMs(keepAlivePeriodMs);
//proto.ScanRequestTimeoutMs = 30000;
//proto.KeepAlivePeriodMs = 15000;
var pruner = PartitionPruner.Create(
schema,
Table.PartitionSchema,
Predicates,
LowerBoundPrimaryKey,
UpperBoundPrimaryKey,
LowerBoundPartitionKey,
UpperBoundPartitionKey);
var keyRanges = new List <KeyRange>();
while (pruner.HasMorePartitionKeyRanges)
{
var partitionRange = pruner.NextPartitionKeyRange;
var newKeyRanges = await Client.GetTableKeyRangesAsync(
Table.TableId,
LowerBoundPrimaryKey,
UpperBoundPrimaryKey,
partitionRange.Lower.Length == 0?null : partitionRange.Lower,
partitionRange.Upper.Length == 0?null : partitionRange.Upper,
_fetchTabletsPerRangeLookup,
_splitSizeBytes,
cancellationToken).ConfigureAwait(false);
if (newKeyRanges.Count == 0)
{
pruner.RemovePartitionKeyRange(partitionRange.Upper);
}
else
{
pruner.RemovePartitionKeyRange(
newKeyRanges[newKeyRanges.Count - 1].PartitionKeyEnd);
}
keyRanges.AddRange(newKeyRanges);
}
var tokens = new List <KuduScanToken>(keyRanges.Count);
var nowMillis = _systemClock.CurrentMilliseconds;
foreach (var keyRange in keyRanges)
{
var token = proto.Clone();
token.LowerBoundPartitionKey = UnsafeByteOperations.UnsafeWrap(keyRange.PartitionKeyStart);
token.UpperBoundPartitionKey = UnsafeByteOperations.UnsafeWrap(keyRange.PartitionKeyEnd);
var primaryKeyStart = keyRange.PrimaryKeyStart;
if (primaryKeyStart.Length > 0)
{
token.LowerBoundPrimaryKey = UnsafeByteOperations.UnsafeWrap(primaryKeyStart);
}
var primaryKeyEnd = keyRange.PrimaryKeyEnd;
if (primaryKeyEnd.Length > 0)
{
token.UpperBoundPrimaryKey = UnsafeByteOperations.UnsafeWrap(primaryKeyEnd);
}
var tablet = keyRange.Tablet;
// Set the tablet metadata so that a call to the master is not needed to
// locate the tablet to scan when opening the scanner.
if (_includeTabletMetadata)
{
// TODO: It would be more efficient to pass the TTL in instead of
// looking it up again here.
var entry = Client.GetTableLocationEntry(
Table.TableId, tablet.Partition.PartitionKeyStart);
long ttl;
if (entry is not null &&
entry.IsCoveredRange &&
(ttl = entry.Expiration - nowMillis) > 0)
{
// Build the list of server and replica metadata.
var tabletServers = tablet.Servers;
var replicas = tablet.Replicas;
var numTabletServers = tabletServers.Count;
var tabletMetadataPb = new TabletMetadataPB
{
TabletId = tablet.TabletId,
Partition = ProtobufHelper.ToPartitionPb(tablet.Partition),
TtlMillis = (ulong)ttl
};
tabletMetadataPb.TabletServers.Capacity = numTabletServers;
tabletMetadataPb.Replicas.Capacity = numTabletServers;
for (int i = 0; i < numTabletServers; i++)
{
var serverInfo = tabletServers[i];
var replica = replicas[i];
var serverMetadataPb = ProtobufHelper.ToServerMetadataPb(serverInfo);
tabletMetadataPb.TabletServers.Add(serverMetadataPb);
var replicaMetadataPb = new ReplicaMetadataPB
{
TsIdx = (uint)i,
Role = (RaftPeerPB.Types.Role)replica.Role
};
if (replica.DimensionLabel is not null)
{
replicaMetadataPb.DimensionLabel = replica.DimensionLabel;
}
tabletMetadataPb.Replicas.Add(replicaMetadataPb);
}
token.TabletMetadata = tabletMetadataPb;
}
}
tokens.Add(new KuduScanToken(keyRange, token));
}
return(tokens);
}