feat: support ddl for branch operation (#576)

## Summary

This PR adds DDL support for branch operations in Lance Spark.

Previously, branch-related operations did not support DDL workflows
properly. With this change, users can perform DDL on branches more
consistently, improving branch management usability and aligning branch
behavior with expected table operation semantics.

## Changes

- Add DDL support for branch operations
- Enable branch-related workflows to work with DDL statements
- Improve consistency of branch management behavior in Lance Spark

## Motivation

Branch operations are an important part of versioned data workflows.
Supporting DDL on branches makes these workflows more complete and
practical, especially for users who need to manage schema or
table-related changes in isolated branches before merging.

---------

Co-authored-by: fangbo <fangbo.0511@bytedance.com>

Author: fangbo

docs/src/operations/ddl/.pages

@@ -12,5 +12,8 @@ nav:
   - drop-table.md
   - create-index.md
   - show-indexes.md
+  - create-branch.md
+  - drop-branch.md
+  - show-branches.md
   - optimize.md
   - vacuum.md

docs/src/operations/ddl/create-branch.md

@@ -0,0 +1,101 @@
+# CREATE BRANCH
+
+Create a named branch for a Lance table.
+
+!!! warning "Spark Extension Required"
+    This feature requires the Lance Spark SQL extension to be enabled. See [Spark SQL Extensions](../../config.md#spark-sql-extensions) for configuration details.
+
+## Overview
+
+The `CREATE BRANCH` command creates a new branch that points to a specific table state. A branch can be created from:
+
+- the latest version of `main`
+- a specific version on `main`
+- the head of another branch
+- a specific version on another branch
+- a tag
+
+Creating a branch records a new reference in table metadata. It does not duplicate the table's data files.
+
+## Syntax
+
+=== "SQL"
+
+    ```sql
+    ALTER TABLE <table> CREATE BRANCH [IF NOT EXISTS] <branch_name>;
+
+    ALTER TABLE <table> CREATE BRANCH [IF NOT EXISTS] <branch_name>
+    AS OF VERSION <version>;
+
+    ALTER TABLE <table> CREATE BRANCH [IF NOT EXISTS] <branch_name>
+    AS OF BRANCH <source_branch>;
+
+    ALTER TABLE <table> CREATE BRANCH [IF NOT EXISTS] <branch_name>
+    AS OF BRANCH <source_branch> VERSION <version>;
+
+    ALTER TABLE <table> CREATE BRANCH [IF NOT EXISTS] <branch_name>
+    AS OF TAG <tag_name>;
+    ```
+
+If the `AS OF` clause is omitted, the new branch is created from the latest version of `main`.
+
+## Examples
+
+### Create a branch from the latest `main`
+
+=== "SQL"
+    ```sql
+    ALTER TABLE lance.db.users CREATE BRANCH feature_x;
+    ```
+
+### Create a branch from a specific `main` version
+
+=== "SQL"
+    ```sql
+    ALTER TABLE lance.db.users CREATE BRANCH IF NOT EXISTS snapshot_v5
+    AS OF VERSION 5;
+    ```
+
+### Create a branch from another branch head
+
+=== "SQL"
+    ```sql
+    ALTER TABLE lance.db.users CREATE BRANCH experiment_b
+    AS OF BRANCH experiment_a;
+    ```
+
+### Create a branch from a specific version on another branch
+
+=== "SQL"
+    ```sql
+    ALTER TABLE lance.db.users CREATE BRANCH experiment_b_v3
+    AS OF BRANCH experiment_a VERSION 3;
+    ```
+
+### Create a branch from a tag
+
+=== "SQL"
+    ```sql
+    ALTER TABLE lance.db.users CREATE BRANCH release_fix
+    AS OF TAG release_candidate;
+    ```
+
+## Output
+
+The `CREATE BRANCH` command returns:
+
+| Column | Type   | Description                |
+|--------|--------|----------------------------|
+| `name` | String | The name of the new branch |
+
+## Notes and Limitations
+
+- `CREATE BRANCH` is implemented as a Spark SQL extension command.
+- The referenced table must be a Lance table.
+- Creating a branch from a non-existent branch, tag, or version returns an error.
+- The branch name is returned even when the command result contains only a single row.
+
+## See Also
+
+- [SHOW BRANCHES](./show-branches.md)
+- [DROP BRANCH](./drop-branch.md)

docs/src/operations/ddl/drop-branch.md

@@ -0,0 +1,44 @@
+# DROP BRANCH
+
+Remove a branch from a Lance table.
+
+!!! warning "Spark Extension Required"
+    This feature requires the Lance Spark SQL extension to be enabled. See [Spark SQL Extensions](../../config.md#spark-sql-extensions) for configuration details.
+
+## Overview
+
+The `DROP BRANCH` command deletes a named branch reference from a Lance table. This removes the branch metadata entry, but does not rewrite data on `main` or on any other existing branch.
+
+## Syntax
+
+=== "SQL"
+    ```sql
+    ALTER TABLE <table> DROP BRANCH [IF EXISTS] <branch_name>;
+    ```
+
+## Example
+
+=== "SQL"
+    ```sql
+    ALTER TABLE lance.db.users DROP BRANCH IF EXISTS feature_x;
+    ```
+
+## Output
+
+The `DROP BRANCH` command returns:
+
+| Column | Type   | Description                    |
+|--------|--------|--------------------------------|
+| `name` | String | The name of the dropped branch |
+
+## Notes and Limitations
+
+- `DROP BRANCH` is implemented as a Spark SQL extension command.
+- The target table must be a Lance table.
+- `IF EXISTS` can be used when dropping a branch conditionally.
+- Use [SHOW BRANCHES](./show-branches.md) to inspect existing branches before deleting one.
+
+## See Also
+
+- [CREATE BRANCH](./create-branch.md)
+- [SHOW BRANCHES](./show-branches.md)

docs/src/operations/ddl/show-branches.md

@@ -0,0 +1,62 @@
+# SHOW BRANCHES
+
+List branches defined on a Lance table.
+
+!!! warning "Spark Extension Required"
+    This feature requires the Lance Spark SQL extension to be enabled. See [Spark SQL Extensions](../../config.md#spark-sql-extensions) for configuration details.
+
+## Overview
+
+The `SHOW BRANCHES` command returns one row for each branch on a Lance table. It is useful for inspecting branch lineage, verifying branch creation, and understanding which version each branch points to.
+
+## Syntax
+
+=== "SQL"
+    ```sql
+    SHOW BRANCHES FROM <table>;
+    SHOW BRANCHES IN <table>;
+    SHOW BRANCH FROM <table>;
+    SHOW BRANCH IN <table>;
+    ```
+
+Both plural and singular forms are accepted. Both `FROM` and `IN` are accepted.
+
+## Examples
+
+### List all branches on a table
+
+=== "SQL"
+    ```sql
+    SHOW BRANCHES FROM lance.db.users;
+    ```
+
+### Use the singular alias
+
+=== "SQL"
+    ```sql
+    SHOW BRANCH IN lance.db.users;
+    ```
+
+## Output
+
+The `SHOW BRANCHES` command returns the following columns:
+
+| Column | Type | Description                                                                      |
+|--------|------|----------------------------------------------------------------------------------|
+| `name` | String | Branch name                                                                      |
+| `parent_branch` | String | Source branch name if the branch was created from another branch; otherwise NULL |
+| `parent_version` | Long | Version used as the branch starting point                                        |
+| `created_at` | Long | Branch creation timestamp as returned by Lance metadata                          |
+| `manifest_size` | Integer | Manifest size recorded for the branch                                            |
+
+## Notes
+
+- `SHOW BRANCHES` is implemented as a Spark SQL extension command.
+- The target table must be a Lance table.
+- The result includes one row per branch currently registered in the table metadata.
+- Use [CREATE BRANCH](./create-branch.md) to add a branch and [DROP BRANCH](./drop-branch.md) to remove one.
+
+## See Also
+
+- [CREATE BRANCH](./create-branch.md)
+- [DROP BRANCH](./drop-branch.md)

lance-spark-3.4_2.12/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensionsAstBuilder.scala

@@ -13,9 +13,10 @@
  */
 package org.apache.spark.sql.catalyst.parser.extensions
 
+import org.antlr.v4.runtime.ParserRuleContext
 import org.apache.spark.sql.catalyst.analysis.{UnresolvedIdentifier, UnresolvedRelation}
-import org.apache.spark.sql.catalyst.parser.ParserInterface
-import org.apache.spark.sql.catalyst.plans.logical.{AddColumnsBackfill, AddIndex, LanceDropIndex, LanceNamedArgument, LogicalPlan, Optimize, SetUnenforcedPrimaryKey, ShowIndexes, UpdateColumnsBackfill, Vacuum}
+import org.apache.spark.sql.catalyst.parser.{ParseException, ParserInterface}
+import org.apache.spark.sql.catalyst.plans.logical.{AddColumnsBackfill, AddIndex, LanceCreateBranch, LanceDropBranch, LanceDropIndex, LanceNamedArgument, LanceShowBranches, LogicalPlan, Optimize, SetUnenforcedPrimaryKey, ShowIndexes, UpdateColumnsBackfill, Vacuum}
 import org.lance.spark.utils.{FieldPathUtils, ParserUtils}
 
 import scala.collection.JavaConverters._
@@ -115,6 +116,71 @@ class LanceSqlExtensionsAstBuilder(delegate: ParserInterface)
     LanceDropIndex(table, indexName)
   }
 
+  override def visitCreateBranchRefMain(ctx: LanceSqlExtensionsParser.CreateBranchRefMainContext)
+      : LanceCreateBranch = {
+    val table = UnresolvedIdentifier(visitMultipartIdentifier(ctx.multipartIdentifier()))
+    val branchName = cleanIdentifier(ctx.branchName.getText)
+    val ifNotExists = ctx.EXISTS() != null
+    if (ctx.refMainVersion == null)
+      LanceCreateBranch(table, branchName, org.lance.Ref.ofMain(), ifNotExists)
+    else LanceCreateBranch(
+      table,
+      branchName,
+      org.lance.Ref.ofMain(_parseVersion(ctx, ctx.refMainVersion.getText)),
+      ifNotExists)
+  }
+
+  override def visitCreateBranchRefBranch(
+      ctx: LanceSqlExtensionsParser.CreateBranchRefBranchContext): LanceCreateBranch = {
+    val table = UnresolvedIdentifier(visitMultipartIdentifier(ctx.multipartIdentifier()))
+    val branchName = cleanIdentifier(ctx.branchName.getText)
+    val refBranchName = cleanIdentifier(ctx.refBranchName.getText)
+    val ifNotExists = ctx.EXISTS() != null
+    if (ctx.refBranchVersion == null)
+      LanceCreateBranch(table, branchName, org.lance.Ref.ofBranch(refBranchName), ifNotExists)
+    else LanceCreateBranch(
+      table,
+      branchName,
+      org.lance.Ref.ofBranch(refBranchName, _parseVersion(ctx, ctx.refBranchVersion.getText)),
+      ifNotExists)
+  }
+
+  private def _parseVersion(ctx: ParserRuleContext, value: String): Long = {
+    try {
+      java.lang.Long.valueOf(value)
+    } catch {
+      case _: NumberFormatException =>
+        throw new ParseException(
+          errorClass = "INVALID_TYPED_LITERAL",
+          messageParameters = Map(
+            "valueType" -> "LONG",
+            "value" -> value),
+          ctx)
+    }
+  }
+
+  override def visitCreateBranchRefTag(ctx: LanceSqlExtensionsParser.CreateBranchRefTagContext)
+      : LanceCreateBranch = {
+    val table = UnresolvedIdentifier(visitMultipartIdentifier(ctx.multipartIdentifier()))
+    val branchName = cleanIdentifier(ctx.branchName.getText)
+    val refTagName = cleanIdentifier(ctx.refTagName.getText)
+    val ifNotExists = ctx.EXISTS() != null
+    LanceCreateBranch(table, branchName, org.lance.Ref.ofTag(refTagName), ifNotExists)
+  }
+
+  override def visitDropBranch(ctx: LanceSqlExtensionsParser.DropBranchContext): LanceDropBranch = {
+    val table = UnresolvedIdentifier(visitMultipartIdentifier(ctx.multipartIdentifier()))
+    val branchName = cleanIdentifier(ctx.branchName.getText)
+    val ifExists = ctx.EXISTS() != null
+    LanceDropBranch(table, branchName, ifExists)
+  }
+
+  override def visitShowBranches(ctx: LanceSqlExtensionsParser.ShowBranchesContext)
+      : LanceShowBranches = {
+    val table = UnresolvedIdentifier(visitMultipartIdentifier(ctx.multipartIdentifier()))
+    LanceShowBranches(table)
+  }
+
   override def visitSetUnenforcedPrimaryKey(
       ctx: LanceSqlExtensionsParser.SetUnenforcedPrimaryKeyContext): SetUnenforcedPrimaryKey = {
     val table = UnresolvedIdentifier(visitMultipartIdentifier(ctx.multipartIdentifier()))

lance-spark-3.4_2.12/src/test/java/org/lance/spark/branch/BranchDDLTest.java

@@ -0,0 +1,16 @@
+/*
+ * 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 org.lance.spark.branch;
+
+public class BranchDDLTest extends BaseBranchDDLTest {}

lance-spark-3.4_2.12/src/test/java/org/lance/spark/update/LanceSqlExtensionsAstBuilderTest.java

@@ -189,4 +189,67 @@ public void testVacuumWithBacktickedTableName() {
     assertEquals(
         List.of("my-catalog", "my-table"), JavaConverters.seqAsJavaList(table.nameParts()));
   }
+
+  @Test
+  public void testCreateBranchFromMain() {
+    LanceSqlExtensionsParser parser =
+        createParser(
+            "ALTER TABLE `my-catalog`.`my-table` CREATE BRANCH IF NOT EXISTS `branch-a` AS OF VERSION 7");
+    org.apache.spark.sql.catalyst.plans.logical.LanceCreateBranch plan =
+        (org.apache.spark.sql.catalyst.plans.logical.LanceCreateBranch)
+            astBuilder.visitSingleStatement(parser.singleStatement());
+
+    UnresolvedIdentifier table = (UnresolvedIdentifier) plan.table();
+    assertEquals(
+        List.of("my-catalog", "my-table"), JavaConverters.seqAsJavaList(table.nameParts()));
+    assertEquals("branch-a", plan.branchName());
+    assertEquals(true, plan.ifNotExists());
+    assertTrue(plan.ref().getBranchName().isEmpty());
+    assertEquals(7, plan.ref().getVersionNumber().get());
+  }
+
+  @Test
+  public void testCreateBranchFromBranch() {
+    LanceSqlExtensionsParser parser =
+        createParser(
+            "ALTER TABLE `my-catalog`.`my-table` CREATE BRANCH IF NOT EXISTS `branch-a` AS OF BRANCH `source-branch` VERSION 7");
+    org.apache.spark.sql.catalyst.plans.logical.LanceCreateBranch plan =
+        (org.apache.spark.sql.catalyst.plans.logical.LanceCreateBranch)
+            astBuilder.visitSingleStatement(parser.singleStatement());
+
+    UnresolvedIdentifier table = (UnresolvedIdentifier) plan.table();
+    assertEquals(
+        List.of("my-catalog", "my-table"), JavaConverters.seqAsJavaList(table.nameParts()));
+    assertEquals("branch-a", plan.branchName());
+    assertEquals(true, plan.ifNotExists());
+    assertEquals("source-branch", plan.ref().getBranchName().get());
+    assertEquals(7, plan.ref().getVersionNumber().get());
+  }
+
+  @Test
+  public void testDropBranch() {
+    LanceSqlExtensionsParser parser =
+        createParser("ALTER TABLE `my-catalog`.`my-table` DROP BRANCH IF EXISTS `branch-a`");
+    org.apache.spark.sql.catalyst.plans.logical.LanceDropBranch plan =
+        (org.apache.spark.sql.catalyst.plans.logical.LanceDropBranch)
+            astBuilder.visitSingleStatement(parser.singleStatement());
+
+    UnresolvedIdentifier table = (UnresolvedIdentifier) plan.table();
+    assertEquals(
+        List.of("my-catalog", "my-table"), JavaConverters.seqAsJavaList(table.nameParts()));
+    assertEquals("branch-a", plan.branchName());
+    assertEquals(true, plan.ifExists());
+  }
+
+  @Test
+  public void testShowBranch() {
+    LanceSqlExtensionsParser parser = createParser("SHOW BRANCH IN `my-catalog`.`my-table`");
+    org.apache.spark.sql.catalyst.plans.logical.LanceShowBranches plan =
+        (org.apache.spark.sql.catalyst.plans.logical.LanceShowBranches)
+            astBuilder.visitSingleStatement(parser.singleStatement());
+
+    UnresolvedIdentifier table = (UnresolvedIdentifier) plan.table();
+    assertEquals(
+        List.of("my-catalog", "my-table"), JavaConverters.seqAsJavaList(table.nameParts()));
+  }
 }