You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@iceberg.apache.org by bl...@apache.org on 2020/12/18 23:11:53 UTC
[iceberg] branch master updated: Spark: Add SQL commands evolve
partition specs (#1948)
This is an automated email from the ASF dual-hosted git repository.
blue pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/iceberg.git
The following commit(s) were added to refs/heads/master by this push:
new d2e5a07 Spark: Add SQL commands evolve partition specs (#1948)
d2e5a07 is described below
commit d2e5a078299477f9387dab78e0a8eae9d39c7ef2
Author: Ryan Blue <bl...@apache.org>
AuthorDate: Fri Dec 18 15:11:41 2020 -0800
Spark: Add SQL commands evolve partition specs (#1948)
---
.../IcebergSqlExtensions.g4 | 24 +-
.../extensions/IcebergSparkSessionExtensions.scala | 2 +-
.../IcebergSparkSqlExtensionsParser.scala | 9 +-
.../IcebergSqlExtensionsAstBuilder.scala | 78 +++++-
.../catalyst/plans/logical/AddPartitionField.scala | 33 +++
.../plans/logical/DropPartitionField.scala | 33 +++
.../datasources/v2/AddPartitionFieldExec.scala | 56 +++++
.../datasources/v2/DropPartitionFieldExec.scala | 67 +++++
.../v2/ExtendedDataSourceV2Strategy.scala | 45 +++-
.../extensions/TestAlterTablePartitionFields.java | 275 +++++++++++++++++++++
.../java/org/apache/iceberg/spark/Spark3Util.java | 27 ++
11 files changed, 639 insertions(+), 10 deletions(-)
diff --git a/spark3-extensions/src/main/antlr/org.apache.spark.sql.catalyst.parser.extensions/IcebergSqlExtensions.g4 b/spark3-extensions/src/main/antlr/org.apache.spark.sql.catalyst.parser.extensions/IcebergSqlExtensions.g4
index db3712e..86ec87d 100644
--- a/spark3-extensions/src/main/antlr/org.apache.spark.sql.catalyst.parser.extensions/IcebergSqlExtensions.g4
+++ b/spark3-extensions/src/main/antlr/org.apache.spark.sql.catalyst.parser.extensions/IcebergSqlExtensions.g4
@@ -66,7 +66,9 @@ singleStatement
;
statement
- : CALL multipartIdentifier '(' (callArgument (',' callArgument)*)? ')' #call
+ : CALL multipartIdentifier '(' (callArgument (',' callArgument)*)? ')' #call
+ | ALTER TABLE multipartIdentifier ADD PARTITION FIELD transform (AS name=identifier)? #addPartitionField
+ | ALTER TABLE multipartIdentifier DROP PARTITION FIELD transform #dropPartitionField
;
callArgument
@@ -74,6 +76,17 @@ callArgument
| identifier '=>' expression #namedArgument
;
+transform
+ : multipartIdentifier #identityTransform
+ | transformName=identifier
+ '(' arguments+=transformArgument (',' arguments+=transformArgument)* ')' #applyTransform
+ ;
+
+transformArgument
+ : multipartIdentifier
+ | constant
+ ;
+
expression
: constant
| stringMap
@@ -121,12 +134,19 @@ quotedIdentifier
;
nonReserved
- : CALL
+ : ADD | ALTER | AS | CALL | DROP | FIELD | PARTITION | TABLE
| TRUE | FALSE
| MAP
;
+ADD: 'ADD';
+ALTER: 'ALTER';
+AS: 'AS';
CALL: 'CALL';
+DROP: 'DROP';
+FIELD: 'FIELD';
+PARTITION: 'PARTITION';
+TABLE: 'TABLE';
TRUE: 'TRUE';
FALSE: 'FALSE';
diff --git a/spark3-extensions/src/main/scala/org/apache/iceberg/spark/extensions/IcebergSparkSessionExtensions.scala b/spark3-extensions/src/main/scala/org/apache/iceberg/spark/extensions/IcebergSparkSessionExtensions.scala
index 2778be3..2f36978 100644
--- a/spark3-extensions/src/main/scala/org/apache/iceberg/spark/extensions/IcebergSparkSessionExtensions.scala
+++ b/spark3-extensions/src/main/scala/org/apache/iceberg/spark/extensions/IcebergSparkSessionExtensions.scala
@@ -37,6 +37,6 @@ class IcebergSparkSessionExtensions extends (SparkSessionExtensions => Unit) {
// TODO: PullupCorrelatedPredicates should handle row-level operations
extensions.injectOptimizerRule { _ => PullupCorrelatedPredicatesInRowLevelOperations }
extensions.injectOptimizerRule { _ => RewriteDelete }
- extensions.injectPlannerStrategy { _ => ExtendedDataSourceV2Strategy }
+ extensions.injectPlannerStrategy { spark => ExtendedDataSourceV2Strategy(spark) }
}
}
diff --git a/spark3-extensions/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/IcebergSparkSqlExtensionsParser.scala b/spark3-extensions/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/IcebergSparkSqlExtensionsParser.scala
index 5680b50..d45d6cb 100644
--- a/spark3-extensions/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/IcebergSparkSqlExtensionsParser.scala
+++ b/spark3-extensions/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/IcebergSparkSqlExtensionsParser.scala
@@ -94,13 +94,20 @@ class IcebergSparkSqlExtensionsParser(delegate: ParserInterface) extends ParserI
*/
override def parsePlan(sqlText: String): LogicalPlan = {
val sqlTextAfterSubstitution = substitutor.substitute(sqlText)
- if (sqlTextAfterSubstitution.toLowerCase(Locale.ROOT).trim().startsWith("call")) {
+ if (isIcebergCommand(sqlTextAfterSubstitution)) {
parse(sqlTextAfterSubstitution) { parser => astBuilder.visit(parser.singleStatement()) }.asInstanceOf[LogicalPlan]
} else {
delegate.parsePlan(sqlText)
}
}
+ private def isIcebergCommand(sqlText: String): Boolean = {
+ val normalized = sqlText.toLowerCase(Locale.ROOT).trim()
+ normalized.startsWith("call") ||
+ (normalized.startsWith("alter table") && (
+ normalized.contains("add partition field") || normalized.contains("drop partition field")))
+ }
+
protected def parse[T](command: String)(toResult: IcebergSqlExtensionsParser => T): T = {
val lexer = new IcebergSqlExtensionsLexer(new UpperCaseCharStream(CharStreams.fromString(command)))
lexer.removeErrorListeners()
diff --git a/spark3-extensions/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/IcebergSqlExtensionsAstBuilder.scala b/spark3-extensions/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/IcebergSqlExtensionsAstBuilder.scala
index bc27797..f67c749 100644
--- a/spark3-extensions/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/IcebergSqlExtensionsAstBuilder.scala
+++ b/spark3-extensions/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/IcebergSqlExtensionsAstBuilder.scala
@@ -20,12 +20,27 @@
package org.apache.spark.sql.catalyst.parser.extensions
import org.antlr.v4.runtime._
-import org.antlr.v4.runtime.tree.{ParseTree, TerminalNode}
+import org.antlr.v4.runtime.tree.ParseTree
+import org.antlr.v4.runtime.tree.TerminalNode
import org.apache.spark.sql.catalyst.expressions.Expression
+import org.apache.spark.sql.catalyst.expressions.Literal
+import org.apache.spark.sql.catalyst.parser.ParseException
import org.apache.spark.sql.catalyst.parser.ParserInterface
import org.apache.spark.sql.catalyst.parser.ParserUtils._
import org.apache.spark.sql.catalyst.parser.extensions.IcebergSqlExtensionsParser._
-import org.apache.spark.sql.catalyst.plans.logical.{CallArgument, CallStatement, LogicalPlan, NamedArgument, PositionalArgument}
+import org.apache.spark.sql.catalyst.plans.logical.AddPartitionField
+import org.apache.spark.sql.catalyst.plans.logical.CallArgument
+import org.apache.spark.sql.catalyst.plans.logical.CallStatement
+import org.apache.spark.sql.catalyst.plans.logical.DropPartitionField
+import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
+import org.apache.spark.sql.catalyst.plans.logical.NamedArgument
+import org.apache.spark.sql.catalyst.plans.logical.PositionalArgument
+import org.apache.spark.sql.connector.expressions
+import org.apache.spark.sql.connector.expressions.ApplyTransform
+import org.apache.spark.sql.connector.expressions.FieldReference
+import org.apache.spark.sql.connector.expressions.IdentityTransform
+import org.apache.spark.sql.connector.expressions.LiteralValue
+import org.apache.spark.sql.connector.expressions.Transform
import scala.collection.JavaConverters._
class IcebergSqlExtensionsAstBuilder(delegate: ParserInterface) extends IcebergSqlExtensionsBaseVisitor[AnyRef] {
@@ -40,6 +55,61 @@ class IcebergSqlExtensionsAstBuilder(delegate: ParserInterface) extends IcebergS
}
/**
+ * Create an ADD PARTITION FIELD logical command.
+ */
+ override def visitAddPartitionField(ctx: AddPartitionFieldContext): AddPartitionField = withOrigin(ctx) {
+ AddPartitionField(
+ typedVisit[Seq[String]](ctx.multipartIdentifier),
+ typedVisit[Transform](ctx.transform),
+ Option(ctx.name).map(_.getText))
+ }
+
+ /**
+ * Create a DROP PARTITION FIELD logical command.
+ */
+ override def visitDropPartitionField(ctx: DropPartitionFieldContext): DropPartitionField = {
+ DropPartitionField(
+ typedVisit[Seq[String]](ctx.multipartIdentifier),
+ typedVisit[Transform](ctx.transform))
+ }
+
+ /**
+ * Create an IdentityTransform for a column reference.
+ */
+ override def visitIdentityTransform(ctx: IdentityTransformContext): Transform = withOrigin(ctx) {
+ IdentityTransform(FieldReference(typedVisit[Seq[String]](ctx.multipartIdentifier())))
+ }
+
+ /**
+ * Create a named Transform from argument expressions.
+ */
+ override def visitApplyTransform(ctx: ApplyTransformContext): Transform = withOrigin(ctx) {
+ val args = ctx.arguments.asScala.map(typedVisit[expressions.Expression])
+ ApplyTransform(ctx.transformName.getText, args)
+ }
+
+ /**
+ * Create a transform argument from a column reference or a constant.
+ */
+ override def visitTransformArgument(ctx: TransformArgumentContext): expressions.Expression = withOrigin(ctx) {
+ val reference = Option(ctx.multipartIdentifier())
+ .map(typedVisit[Seq[String]])
+ .map(FieldReference(_))
+ val literal = Option(ctx.constant)
+ .map(visitConstant)
+ .map(lit => LiteralValue(lit.value, lit.dataType))
+ reference.orElse(literal)
+ .getOrElse(throw new ParseException(s"Invalid transform argument", ctx))
+ }
+
+ /**
+ * Return a multi-part identifier as Seq[String].
+ */
+ override def visitMultipartIdentifier(ctx: MultipartIdentifierContext): Seq[String] = withOrigin(ctx) {
+ ctx.parts.asScala.map(_.getText)
+ }
+
+ /**
* Create a positional argument in a stored procedure call.
*/
override def visitPositionalArgument(ctx: PositionalArgumentContext): CallArgument = withOrigin(ctx) {
@@ -60,6 +130,10 @@ class IcebergSqlExtensionsAstBuilder(delegate: ParserInterface) extends IcebergS
visit(ctx.statement).asInstanceOf[LogicalPlan]
}
+ def visitConstant(ctx: ConstantContext): Literal = {
+ delegate.parseExpression(ctx.getText).asInstanceOf[Literal]
+ }
+
override def visitExpression(ctx: ExpressionContext): Expression = {
// reconstruct the SQL string and parse it using the main Spark parser
// while we can avoid the logic to build Spark expressions, we still have to parse them
diff --git a/spark3-extensions/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/AddPartitionField.scala b/spark3-extensions/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/AddPartitionField.scala
new file mode 100644
index 0000000..ae5a239
--- /dev/null
+++ b/spark3-extensions/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/AddPartitionField.scala
@@ -0,0 +1,33 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.spark.sql.catalyst.plans.logical
+
+import org.apache.spark.sql.catalyst.expressions.Attribute
+import org.apache.spark.sql.connector.expressions.Transform
+
+case class AddPartitionField(table: Seq[String], transform: Transform, name: Option[String]) extends Command {
+ import org.apache.spark.sql.connector.catalog.CatalogV2Implicits._
+
+ override lazy val output: Seq[Attribute] = Nil
+
+ override def simpleString(maxFields: Int): String = {
+ s"AddPartitionField ${table.quoted} ${name.map(n => s"$n=").getOrElse("")}${transform.describe}"
+ }
+}
diff --git a/spark3-extensions/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/DropPartitionField.scala b/spark3-extensions/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/DropPartitionField.scala
new file mode 100644
index 0000000..5fe8f92
--- /dev/null
+++ b/spark3-extensions/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/DropPartitionField.scala
@@ -0,0 +1,33 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.spark.sql.catalyst.plans.logical
+
+import org.apache.spark.sql.catalyst.expressions.Attribute
+import org.apache.spark.sql.connector.expressions.Transform
+
+case class DropPartitionField(table: Seq[String], transform: Transform) extends Command {
+ import org.apache.spark.sql.connector.catalog.CatalogV2Implicits._
+
+ override lazy val output: Seq[Attribute] = Nil
+
+ override def simpleString(maxFields: Int): String = {
+ s"DropPartitionField ${table.quoted} ${transform.describe}"
+ }
+}
diff --git a/spark3-extensions/src/main/scala/org/apache/spark/sql/execution/datasources/v2/AddPartitionFieldExec.scala b/spark3-extensions/src/main/scala/org/apache/spark/sql/execution/datasources/v2/AddPartitionFieldExec.scala
new file mode 100644
index 0000000..b33de89
--- /dev/null
+++ b/spark3-extensions/src/main/scala/org/apache/spark/sql/execution/datasources/v2/AddPartitionFieldExec.scala
@@ -0,0 +1,56 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.spark.sql.execution.datasources.v2
+
+import org.apache.iceberg.spark.Spark3Util
+import org.apache.iceberg.spark.source.SparkTable
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.Attribute
+import org.apache.spark.sql.connector.catalog.Identifier
+import org.apache.spark.sql.connector.catalog.TableCatalog
+import org.apache.spark.sql.connector.expressions.Transform
+
+case class AddPartitionFieldExec(
+ catalog: TableCatalog,
+ ident: Identifier,
+ transform: Transform,
+ name: Option[String]) extends V2CommandExec {
+ import org.apache.spark.sql.connector.catalog.CatalogV2Implicits._
+
+ override lazy val output: Seq[Attribute] = Nil
+
+ override protected def run(): Seq[InternalRow] = {
+ catalog.loadTable(ident) match {
+ case iceberg: SparkTable =>
+ iceberg.table.updateSpec()
+ .addField(name.orNull, Spark3Util.toIcebergTerm(transform))
+ .commit()
+
+ case table =>
+ throw new UnsupportedOperationException(s"Cannot add partition field to non-Iceberg table: $table")
+ }
+
+ Nil
+ }
+
+ override def simpleString(maxFields: Int): String = {
+ s"AddPartitionField ${catalog.name}.${ident.quoted} ${name.map(n => s"$n=").getOrElse("")}${transform.describe}"
+ }
+}
diff --git a/spark3-extensions/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DropPartitionFieldExec.scala b/spark3-extensions/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DropPartitionFieldExec.scala
new file mode 100644
index 0000000..76128c5
--- /dev/null
+++ b/spark3-extensions/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DropPartitionFieldExec.scala
@@ -0,0 +1,67 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.spark.sql.execution.datasources.v2
+
+import org.apache.iceberg.spark.Spark3Util
+import org.apache.iceberg.spark.source.SparkTable
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.Attribute
+import org.apache.spark.sql.connector.catalog.Identifier
+import org.apache.spark.sql.connector.catalog.TableCatalog
+import org.apache.spark.sql.connector.expressions.FieldReference
+import org.apache.spark.sql.connector.expressions.IdentityTransform
+import org.apache.spark.sql.connector.expressions.Transform
+
+case class DropPartitionFieldExec(
+ catalog: TableCatalog,
+ ident: Identifier,
+ transform: Transform) extends V2CommandExec {
+ import org.apache.spark.sql.connector.catalog.CatalogV2Implicits._
+
+ override lazy val output: Seq[Attribute] = Nil
+
+ override protected def run(): Seq[InternalRow] = {
+ catalog.loadTable(ident) match {
+ case iceberg: SparkTable =>
+ val schema = iceberg.table.schema
+ transform match {
+ case IdentityTransform(FieldReference(parts)) if parts.size == 1 && schema.findField(parts.head) == null =>
+ // the name is not present in the Iceberg schema, so it must be a partition field name, not a column name
+ iceberg.table.updateSpec()
+ .removeField(parts.head)
+ .commit()
+
+ case _ =>
+ iceberg.table.updateSpec()
+ .removeField(Spark3Util.toIcebergTerm(transform))
+ .commit()
+ }
+
+ case table =>
+ throw new UnsupportedOperationException(s"Cannot drop partition field in non-Iceberg table: $table")
+ }
+
+ Nil
+ }
+
+ override def simpleString(maxFields: Int): String = {
+ s"DropPartitionField ${catalog.name}.${ident.quoted} ${transform.describe}"
+ }
+}
diff --git a/spark3-extensions/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ExtendedDataSourceV2Strategy.scala b/spark3-extensions/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ExtendedDataSourceV2Strategy.scala
index 5354405..f085865 100644
--- a/spark3-extensions/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ExtendedDataSourceV2Strategy.scala
+++ b/spark3-extensions/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ExtendedDataSourceV2Strategy.scala
@@ -19,18 +19,39 @@
package org.apache.spark.sql.execution.datasources.v2
+import org.apache.iceberg.spark.Spark3Util
+import org.apache.iceberg.spark.SparkCatalog
+import org.apache.iceberg.spark.SparkSessionCatalog
+import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.Strategy
import org.apache.spark.sql.catalyst.InternalRow
-import org.apache.spark.sql.catalyst.expressions.{Expression, GenericInternalRow}
-import org.apache.spark.sql.catalyst.plans.logical.{Call, DynamicFileFilter, LogicalPlan, ReplaceData}
-import org.apache.spark.sql.execution.{ProjectExec, SparkPlan}
+import org.apache.spark.sql.catalyst.expressions.Expression
+import org.apache.spark.sql.catalyst.expressions.GenericInternalRow
+import org.apache.spark.sql.catalyst.plans.logical.AddPartitionField
+import org.apache.spark.sql.catalyst.plans.logical.Call
+import org.apache.spark.sql.catalyst.plans.logical.DropPartitionField
+import org.apache.spark.sql.catalyst.plans.logical.DynamicFileFilter
+import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
+import org.apache.spark.sql.catalyst.plans.logical.ReplaceData
+import org.apache.spark.sql.connector.catalog.Identifier
+import org.apache.spark.sql.connector.catalog.TableCatalog
+import org.apache.spark.sql.execution.ProjectExec
+import org.apache.spark.sql.execution.SparkPlan
+import scala.collection.JavaConverters._
-object ExtendedDataSourceV2Strategy extends Strategy {
+case class ExtendedDataSourceV2Strategy(spark: SparkSession) extends Strategy {
override def apply(plan: LogicalPlan): Seq[SparkPlan] = plan match {
case c @ Call(procedure, args) =>
val input = buildInternalRow(args)
CallExec(c.output, procedure, input) :: Nil
+
+ case AddPartitionField(IcebergCatalogAndIdentifier(catalog, ident), transform, name) =>
+ AddPartitionFieldExec(catalog, ident, transform, name) :: Nil
+
+ case DropPartitionField(IcebergCatalogAndIdentifier(catalog, ident), transform) =>
+ DropPartitionFieldExec(catalog, ident, transform) :: Nil
+
case DynamicFileFilter(scanRelation, fileFilterPlan) =>
// we don't use planLater here as we need ExtendedBatchScanExec, not BatchScanExec
val scanExec = ExtendedBatchScanExec(scanRelation.output, scanRelation.scan)
@@ -41,8 +62,10 @@ object ExtendedDataSourceV2Strategy extends Strategy {
// add a projection to ensure we have UnsafeRows required by some operations
ProjectExec(scanRelation.output, dynamicFileFilter) :: Nil
}
+
case ReplaceData(_, batchWrite, query) =>
ReplaceDataExec(batchWrite, planLater(query)) :: Nil
+
case _ => Nil
}
@@ -53,4 +76,18 @@ object ExtendedDataSourceV2Strategy extends Strategy {
}
new GenericInternalRow(values)
}
+
+ private object IcebergCatalogAndIdentifier {
+ def unapply(identifier: Seq[String]): Option[(TableCatalog, Identifier)] = {
+ val catalogAndIdentifier = Spark3Util.catalogAndIdentifier(spark, identifier.asJava)
+ catalogAndIdentifier.catalog match {
+ case icebergCatalog: SparkCatalog =>
+ Some((icebergCatalog, catalogAndIdentifier.identifier))
+ case icebergCatalog: SparkSessionCatalog[_] =>
+ Some((icebergCatalog, catalogAndIdentifier.identifier))
+ case _ =>
+ None
+ }
+ }
+ }
}
diff --git a/spark3-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestAlterTablePartitionFields.java b/spark3-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestAlterTablePartitionFields.java
new file mode 100644
index 0000000..80206b8
--- /dev/null
+++ b/spark3-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestAlterTablePartitionFields.java
@@ -0,0 +1,275 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.iceberg.spark.extensions;
+
+import java.util.Map;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.Table;
+import org.junit.After;
+import org.junit.Assert;
+import org.junit.Test;
+
+public class TestAlterTablePartitionFields extends SparkExtensionsTestBase {
+ public TestAlterTablePartitionFields(String catalogName, String implementation, Map<String, String> config) {
+ super(catalogName, implementation, config);
+ }
+
+ @After
+ public void removeTable() {
+ sql("DROP TABLE IF EXISTS %s", tableName);
+ }
+
+ @Test
+ public void testAddIdentityPartition() {
+ sql("CREATE TABLE %s (id bigint NOT NULL, category string, ts timestamp, data string) USING iceberg", tableName);
+ Table table = validationCatalog.loadTable(tableIdent);
+
+ Assert.assertTrue("Table should start unpartitioned", table.spec().isUnpartitioned());
+
+ sql("ALTER TABLE %s ADD PARTITION FIELD category", tableName);
+
+ table.refresh();
+
+ PartitionSpec expected = PartitionSpec.builderFor(table.schema())
+ .withSpecId(1)
+ .identity("category")
+ .build();
+
+ Assert.assertEquals("Should have new spec field", expected, table.spec());
+ }
+
+ @Test
+ public void testAddBucketPartition() {
+ sql("CREATE TABLE %s (id bigint NOT NULL, category string, ts timestamp, data string) USING iceberg", tableName);
+ Table table = validationCatalog.loadTable(tableIdent);
+
+ Assert.assertTrue("Table should start unpartitioned", table.spec().isUnpartitioned());
+
+ sql("ALTER TABLE %s ADD PARTITION FIELD bucket(16, id)", tableName);
+
+ table.refresh();
+
+ PartitionSpec expected = PartitionSpec.builderFor(table.schema())
+ .withSpecId(1)
+ .bucket("id", 16, "id_bucket_16")
+ .build();
+
+ Assert.assertEquals("Should have new spec field", expected, table.spec());
+ }
+
+ @Test
+ public void testAddTruncatePartition() {
+ sql("CREATE TABLE %s (id bigint NOT NULL, category string, ts timestamp, data string) USING iceberg", tableName);
+ Table table = validationCatalog.loadTable(tableIdent);
+
+ Assert.assertTrue("Table should start unpartitioned", table.spec().isUnpartitioned());
+
+ sql("ALTER TABLE %s ADD PARTITION FIELD truncate(data, 4)", tableName);
+
+ table.refresh();
+
+ PartitionSpec expected = PartitionSpec.builderFor(table.schema())
+ .withSpecId(1)
+ .truncate("data", 4, "data_trunc_4")
+ .build();
+
+ Assert.assertEquals("Should have new spec field", expected, table.spec());
+ }
+
+ @Test
+ public void testAddYearsPartition() {
+ sql("CREATE TABLE %s (id bigint NOT NULL, category string, ts timestamp, data string) USING iceberg", tableName);
+ Table table = validationCatalog.loadTable(tableIdent);
+
+ Assert.assertTrue("Table should start unpartitioned", table.spec().isUnpartitioned());
+
+ sql("ALTER TABLE %s ADD PARTITION FIELD years(ts)", tableName);
+
+ table.refresh();
+
+ PartitionSpec expected = PartitionSpec.builderFor(table.schema())
+ .withSpecId(1)
+ .year("ts")
+ .build();
+
+ Assert.assertEquals("Should have new spec field", expected, table.spec());
+ }
+
+ @Test
+ public void testAddMonthsPartition() {
+ sql("CREATE TABLE %s (id bigint NOT NULL, category string, ts timestamp, data string) USING iceberg", tableName);
+ Table table = validationCatalog.loadTable(tableIdent);
+
+ Assert.assertTrue("Table should start unpartitioned", table.spec().isUnpartitioned());
+
+ sql("ALTER TABLE %s ADD PARTITION FIELD months(ts)", tableName);
+
+ table.refresh();
+
+ PartitionSpec expected = PartitionSpec.builderFor(table.schema())
+ .withSpecId(1)
+ .month("ts")
+ .build();
+
+ Assert.assertEquals("Should have new spec field", expected, table.spec());
+ }
+
+ @Test
+ public void testAddDaysPartition() {
+ sql("CREATE TABLE %s (id bigint NOT NULL, category string, ts timestamp, data string) USING iceberg", tableName);
+ Table table = validationCatalog.loadTable(tableIdent);
+
+ Assert.assertTrue("Table should start unpartitioned", table.spec().isUnpartitioned());
+
+ sql("ALTER TABLE %s ADD PARTITION FIELD days(ts)", tableName);
+
+ table.refresh();
+
+ PartitionSpec expected = PartitionSpec.builderFor(table.schema())
+ .withSpecId(1)
+ .day("ts")
+ .build();
+
+ Assert.assertEquals("Should have new spec field", expected, table.spec());
+ }
+
+ @Test
+ public void testAddHoursPartition() {
+ sql("CREATE TABLE %s (id bigint NOT NULL, category string, ts timestamp, data string) USING iceberg", tableName);
+ Table table = validationCatalog.loadTable(tableIdent);
+
+ Assert.assertTrue("Table should start unpartitioned", table.spec().isUnpartitioned());
+
+ sql("ALTER TABLE %s ADD PARTITION FIELD hours(ts)", tableName);
+
+ table.refresh();
+
+ PartitionSpec expected = PartitionSpec.builderFor(table.schema())
+ .withSpecId(1)
+ .hour("ts")
+ .build();
+
+ Assert.assertEquals("Should have new spec field", expected, table.spec());
+ }
+
+ @Test
+ public void testAddNamedPartition() {
+ sql("CREATE TABLE %s (id bigint NOT NULL, category string, ts timestamp, data string) USING iceberg", tableName);
+ Table table = validationCatalog.loadTable(tableIdent);
+
+ Assert.assertTrue("Table should start unpartitioned", table.spec().isUnpartitioned());
+
+ sql("ALTER TABLE %s ADD PARTITION FIELD bucket(16, id) AS shard", tableName);
+
+ table.refresh();
+
+ PartitionSpec expected = PartitionSpec.builderFor(table.schema())
+ .withSpecId(1)
+ .bucket("id", 16, "shard")
+ .build();
+
+ Assert.assertEquals("Should have new spec field", expected, table.spec());
+ }
+
+ @Test
+ public void testDropIdentityPartition() {
+ sql("CREATE TABLE %s (id bigint NOT NULL, category string, data string) USING iceberg PARTITIONED BY (category)",
+ tableName);
+ Table table = validationCatalog.loadTable(tableIdent);
+
+ Assert.assertEquals("Table should start with 1 partition field", 1, table.spec().fields().size());
+
+ sql("ALTER TABLE %s DROP PARTITION FIELD category", tableName);
+
+ table.refresh();
+
+ PartitionSpec expected = PartitionSpec.builderFor(table.schema())
+ .withSpecId(1)
+ .alwaysNull("category", "category")
+ .build();
+
+ Assert.assertEquals("Should have new spec field", expected, table.spec());
+ }
+
+ @Test
+ public void testDropDaysPartition() {
+ sql("CREATE TABLE %s (id bigint NOT NULL, ts timestamp, data string) USING iceberg PARTITIONED BY (days(ts))",
+ tableName);
+ Table table = validationCatalog.loadTable(tableIdent);
+
+ Assert.assertEquals("Table should start with 1 partition field", 1, table.spec().fields().size());
+
+ sql("ALTER TABLE %s DROP PARTITION FIELD days(ts)", tableName);
+
+ table.refresh();
+
+ PartitionSpec expected = PartitionSpec.builderFor(table.schema())
+ .withSpecId(1)
+ .alwaysNull("ts", "ts_day")
+ .build();
+
+ Assert.assertEquals("Should have new spec field", expected, table.spec());
+ }
+
+ @Test
+ public void testDropBucketPartition() {
+ sql("CREATE TABLE %s (id bigint NOT NULL, data string) USING iceberg PARTITIONED BY (bucket(16, id))",
+ tableName);
+ Table table = validationCatalog.loadTable(tableIdent);
+
+ Assert.assertEquals("Table should start with 1 partition field", 1, table.spec().fields().size());
+
+ sql("ALTER TABLE %s DROP PARTITION FIELD bucket(16, id)", tableName);
+
+ table.refresh();
+
+ PartitionSpec expected = PartitionSpec.builderFor(table.schema())
+ .withSpecId(1)
+ .alwaysNull("id", "id_bucket")
+ .build();
+
+ Assert.assertEquals("Should have new spec field", expected, table.spec());
+ }
+
+ @Test
+ public void testDropPartitionByName() {
+ sql("CREATE TABLE %s (id bigint NOT NULL, category string, ts timestamp, data string) USING iceberg", tableName);
+ Table table = validationCatalog.loadTable(tableIdent);
+
+ Assert.assertTrue("Table should start unpartitioned", table.spec().isUnpartitioned());
+
+ sql("ALTER TABLE %s ADD PARTITION FIELD bucket(16, id) AS shard", tableName);
+
+ table.refresh();
+
+ Assert.assertEquals("Table should have 1 partition field", 1, table.spec().fields().size());
+
+ sql("ALTER TABLE %s DROP PARTITION FIELD shard", tableName);
+
+ table.refresh();
+
+ PartitionSpec expected = PartitionSpec.builderFor(table.schema())
+ .withSpecId(2)
+ .alwaysNull("id", "shard")
+ .build();
+
+ Assert.assertEquals("Should have new spec field", expected, table.spec());
+ }
+}
diff --git a/spark3/src/main/java/org/apache/iceberg/spark/Spark3Util.java b/spark3/src/main/java/org/apache/iceberg/spark/Spark3Util.java
index 62045be..6134b9a 100644
--- a/spark3/src/main/java/org/apache/iceberg/spark/Spark3Util.java
+++ b/spark3/src/main/java/org/apache/iceberg/spark/Spark3Util.java
@@ -34,6 +34,7 @@ import org.apache.iceberg.UpdateProperties;
import org.apache.iceberg.UpdateSchema;
import org.apache.iceberg.expressions.BoundPredicate;
import org.apache.iceberg.expressions.ExpressionVisitors;
+import org.apache.iceberg.expressions.Term;
import org.apache.iceberg.expressions.UnboundPredicate;
import org.apache.iceberg.hadoop.HadoopInputFile;
import org.apache.iceberg.io.FileIO;
@@ -265,6 +266,32 @@ public class Spark3Util {
return transforms.toArray(new Transform[0]);
}
+ public static Term toIcebergTerm(Transform transform) {
+ Preconditions.checkArgument(transform.references().length == 1,
+ "Cannot convert transform with more than one column reference: %s", transform);
+ String colName = DOT.join(transform.references()[0].fieldNames());
+ switch (transform.name()) {
+ case "identity":
+ return org.apache.iceberg.expressions.Expressions.ref(colName);
+ case "bucket":
+ return org.apache.iceberg.expressions.Expressions.bucket(colName, findWidth(transform));
+ case "years":
+ return org.apache.iceberg.expressions.Expressions.year(colName);
+ case "months":
+ return org.apache.iceberg.expressions.Expressions.month(colName);
+ case "date":
+ case "days":
+ return org.apache.iceberg.expressions.Expressions.day(colName);
+ case "date_hour":
+ case "hours":
+ return org.apache.iceberg.expressions.Expressions.hour(colName);
+ case "truncate":
+ return org.apache.iceberg.expressions.Expressions.truncate(colName, findWidth(transform));
+ default:
+ throw new UnsupportedOperationException("Transform is not supported: " + transform);
+ }
+ }
+
/**
* Converts Spark transforms into a {@link PartitionSpec}.
*