feat(glue): support partition index on tables

packages/@aws-cdk/aws-glue/README.md

@@ -194,7 +194,7 @@ new glue.Table(this, 'MyTable', {
 
 By default, an S3 bucket will be created to store the table's data and stored in the bucket root. You can also manually pass the `bucket` and `s3Prefix`:
 
-### Partitions
+### Partition Keys
 
 To improve query performance, a table can specify `partitionKeys` on which data is stored and queried separately. For example, you might partition a table by `year` and `month` to optimize queries based on a time window:
 
@@ -218,6 +218,52 @@ new glue.Table(this, 'MyTable', {
 });
 ```
 
+### Partition Indexes
+
+Another way to improve query performance is to specify partition indexes. If no partition indexes are
+present on the table, AWS Glue loads all partitions of the table and filters the loaded partitions using
+the query expression. The query takes more time to run as the number of partitions increase. With an
+index, the query will try to fetch a subset of the partitions instead of loading all partitions of the
+table.
+
+The keys of a partition index must be a subset of the partition keys of the table. You can have a
+maximum of 3 partition indexes per table. To specify a partition index, you can use the `partitionIndexes`
+property:
+
+```ts
+declare const myDatabase: glue.Database;
+new glue.Table(this, 'MyTable', {
+ database: myDatabase,
+ tableName: 'my_table',
+ columns: [{
+ name: 'col1',
+ type: glue.Schema.STRING,
+ }],
+ partitionKeys: [{
+ name: 'year',
+ type: glue.Schema.SMALL_INT,
+ }, {
+ name: 'month',
+ type: glue.Schema.SMALL_INT,
+ }],
+ partitionIndexes: [{
+ indexName: 'my-index', // optional
+ keyNames: ['year'],
+ }], // supply up to 3 indexes
+ dataFormat: glue.DataFormat.JSON,
+});
+```
+
+Alternatively, you can call the `addPartitionIndex()` function on a table:
+
+```ts
+declare const myTable: glue.Table;
+myTable.addPartitionIndex({
+ indexName: 'my-index',
+ keyNames: ['year'],
+});
+```
+
 ## [Encryption](https://docs.aws.amazon.com/athena/latest/ug/encryption.html)
 
 You can enable encryption on a Table's data:

packages/@aws-cdk/aws-glue/lib/table.ts

@@ -1,13 +1,33 @@
 import * as iam from '@aws-cdk/aws-iam';
 import * as kms from '@aws-cdk/aws-kms';
 import * as s3 from '@aws-cdk/aws-s3';
-import { ArnFormat, Fn, IResource, Resource, Stack } from '@aws-cdk/core';
+import { ArnFormat, Fn, IResource, Names, Resource, Stack } from '@aws-cdk/core';
+import * as cr from '@aws-cdk/custom-resources';
+import { AwsCustomResource } from '@aws-cdk/custom-resources';
 import { Construct } from 'constructs';
 import { DataFormat } from './data-format';
 import { IDatabase } from './database';
 import { CfnTable } from './glue.generated';
 import { Column } from './schema';
 
+/**
+ * Properties of a Partition Index.
+ */
+export interface PartitionIndex {
+ /**
+ * The name of the partition index.
+ *
+ * @default - a name will be generated for you.
+ */
+ readonly indexName?: string;
+
+ /**
+ * The partition key names that comprise the partition
+ * index. The names must correspond to a name in the
+ * table's partition keys.
+ */
+ readonly keyNames: string[];
+}
 export interface ITable extends IResource {
  /**
  * @attribute
@@ -102,7 +122,16 @@ export interface TableProps {
  *
  * @default table is not partitioned
  */
- readonly partitionKeys?: Column[]
+ readonly partitionKeys?: Column[];
+
+ /**
+ * Partition indexes on the table. A maximum of 3 indexes
+ * are allowed on a table. Keys in the index must be part
+ * of the table's partition keys.
+ *
+ * @default table has no partition indexes
+ */
+ readonly partitionIndexes?: PartitionIndex[];
 
  /**
  * Storage type of the table's data.
@@ -230,6 +259,18 @@ export class Table extends Resource implements ITable {
  */
  public readonly partitionKeys?: Column[];
 
+ /**
+ * This table's partition indexes.
+ */
+ public readonly partitionIndexes?: PartitionIndex[];
+
+ /**
+ * Partition indexes must be created one at a time. To avoid
+ * race conditions, we store the resource and add dependencies
+ * each time a new partition index is created.
+ */
+ private partitionIndexCustomResources: AwsCustomResource[] = [];
+
  constructor(scope: Construct, id: string, props: TableProps) {
  super(scope, id, {
  physicalName: props.tableName,
@@ -287,6 +328,77 @@ export class Table extends Resource implements ITable {
  resourceName: `${this.database.databaseName}/${this.tableName}`,
  });
  this.node.defaultChild = tableResource;
+
+ // Partition index creation relies on created table.
+ if (props.partitionIndexes) {
+ this.partitionIndexes = props.partitionIndexes;
+ this.partitionIndexes.forEach((index) => this.addPartitionIndex(index));
+ }
+ }
+
+ /**
+ * Add a partition index to the table. You can have a maximum of 3 partition
+ * indexes to a table. Partition index keys must be a subset of the table's
+ * partition keys.
+ *
+ * @see https://docs.aws.amazon.com/glue/latest/dg/partition-indexes.html
+ */
+ public addPartitionIndex(index: PartitionIndex) {
+ const numPartitions = this.partitionIndexCustomResources.length;
+ if (numPartitions >= 3) {
+ throw new Error('Maximum number of partition indexes allowed is 3');
+ }
+ this.validatePartitionIndex(index);
+
+ const indexName = index.indexName ?? this.generateIndexName(index.keyNames);
+ const partitionIndexCustomResource = new cr.AwsCustomResource(this, `partition-index-${indexName}`, {
+ onCreate: {
+ service: 'Glue',
+ action: 'createPartitionIndex',
+ parameters: {
+ DatabaseName: this.database.databaseName,
+ TableName: this.tableName,
+ PartitionIndex: {
+ IndexName: indexName,
+ Keys: index.keyNames,
+ },
+ },
+ physicalResourceId: cr.PhysicalResourceId.of(
+ indexName,
+ ),
+ },
+ policy: cr.AwsCustomResourcePolicy.fromSdkCalls({
+ resources: cr.AwsCustomResourcePolicy.ANY_RESOURCE,
+ }),
+ });
+ this.grantToUnderlyingResources(partitionIndexCustomResource, ['glue:UpdateTable']);
+
+ // Depend on previous partition index if possible, to avoid race condition
+ if (numPartitions > 0) {
+ this.partitionIndexCustomResources[numPartitions-1].node.addDependency(partitionIndexCustomResource);
+ }
+ this.partitionIndexCustomResources.push(partitionIndexCustomResource);
+ }
+
+ private generateIndexName(keys: string[]): string {
+ const prefix = keys.join('-') + '-';
+ const uniqueId = Names.uniqueId(this);
+ const maxIndexLength = 80; // arbitrarily specified
+ const startIndex = Math.max(0, uniqueId.length - (maxIndexLength - prefix.length));
+ return prefix + uniqueId.substring(startIndex);
+ }
+
+ private validatePartitionIndex(index: PartitionIndex) {
+ if (index.indexName !== undefined && (index.indexName.length < 1 || index.indexName.length > 255)) {
+ throw new Error(`Index name must be between 1 and 255 characters, but got ${index.indexName.length}`);
+ }
+ if (!this.partitionKeys || this.partitionKeys.length === 0) {
+ throw new Error('The table must have partition keys to create a partition index');
+ }
+ const keyNames = this.partitionKeys.map(pk => pk.name);
+ if (!index.keyNames.every(k => keyNames.includes(k))) {
+ throw new Error(`All index keys must also be partition keys. Got ${index.keyNames} but partition key names are ${keyNames}`);
+ }
  }
 
  /**
@@ -336,6 +448,22 @@ export class Table extends Resource implements ITable {
  });
  }
 
+ /**
+ * Grant the given identity custom permissions to ALL underlying resources of the table.
+ * Permissions will be granted to the catalog, the database, and the table.
+ */
+ public grantToUnderlyingResources(grantee: iam.IGrantable, actions: string[]) {
+ return iam.Grant.addToPrincipal({
+ grantee,
+ resourceArns: [
+ this.tableArn,
+ this.database.catalogArn,
+ this.database.databaseArn,
+ ],
+ actions,
+ });
+ }
+
  private getS3PrefixForGrant() {
  return this.s3Prefix + '*';
  }

packages/@aws-cdk/aws-glue/package.json

@@ -99,6 +99,7 @@
  "@aws-cdk/aws-s3": "0.0.0",
  "@aws-cdk/aws-s3-assets": "0.0.0",
  "@aws-cdk/core": "0.0.0",
+ "@aws-cdk/custom-resources": "0.0.0",
  "constructs": "^3.3.69"
  },
  "homepage": "https://github.com/aws/aws-cdk",
@@ -113,6 +114,7 @@
  "@aws-cdk/aws-s3": "0.0.0",
  "@aws-cdk/aws-s3-assets": "0.0.0",
  "@aws-cdk/core": "0.0.0",
+ "@aws-cdk/custom-resources": "0.0.0",
  "constructs": "^3.3.69"
  },
  "engines": {