Spark 3.1.0 APIs - Column

src/csharp/Microsoft.Spark.E2ETest/IpcTests/Sql/ColumnTests.cs

@@ -2,6 +2,7 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 // See the LICENSE file in the project root for more information.
 
+using Microsoft.Spark.E2ETest.Utils;
 using Microsoft.Spark.Sql;
 using Xunit;
 using static Microsoft.Spark.Sql.Expressions.Window;
@@ -143,5 +144,18 @@ public void TestSignaturesV2_3_X()
             Assert.Equal("col2", col2.ToString());
         }
 
+        /// <summary>
+        /// Test signatures for APIs introduced in Spark 3.1.*.
+        /// </summary>
+        [SkipIfSparkVersionIsLessThan(Versions.V3_1_0)]
+        public void TestSignaturesV3_1_X()
+        {
+            Column col = Column("col");
+
+            Assert.IsType<Column>(col.WithField("col2", Lit(3)));
+
+            Assert.IsType<Column>(col.DropFields("col"));
+            Assert.IsType<Column>(col.DropFields("col", "col2"));
+        }
     }
 }

src/csharp/Microsoft.Spark/Sql/Column.cs

@@ -450,6 +450,31 @@ public Column GetItem(object key)
             return ApplyMethod("getItem", key);
         }
 
+        /// <summary>
+        /// An expression that adds/replaces field in <see cref="Types.StructType"/> by name.
+        /// </summary>
+        /// <param name="fieldName">The name of the field</param>
+        /// <param name="column">Column to assign to the field</param>
+        /// <returns>
+        /// New column after adding/replacing field in <see cref="Types.StructType"/> by name.
+        /// </returns>
+        [Since(Versions.V3_1_0)]
+        public Column WithField(string fieldName, Column column)
+        {
+            return ApplyMethod("withField", fieldName, column);
+        }
+
+        /// <summary>
+        /// An expression that drops fields in <see cref="Types.StructType"/> by name.
+        /// </summary>
+        /// <param name="fieldNames">Name of fields to drop.</param>
+        /// <returns>New column after after dropping fields.</returns>
+        [Since(Versions.V3_1_0)]
+        public Column DropFields(params string[] fieldNames)
+        {
+            return ApplyMethod("dropFields", fieldNames);
+        }
+
         /// <summary>
         /// An expression that gets a field by name in a `StructType`.
         /// </summary>