Exposing DataFrame.Persist(StorageLevel) and DataFrame.StorageLevel APIs

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

@@ -651,8 +651,12 @@ public void TestSignaturesV2_3_X()
 
             Assert.IsType<DataFrame>(_df.Persist());
 
+            Assert.IsType<DataFrame>(_df.Persist(StorageLevel.DISK_ONLY));
+
             Assert.IsType<DataFrame>(_df.Cache());
 
+            Assert.IsType<StorageLevel>(_df.StorageLevel());
+
             Assert.IsType<DataFrame>(_df.Unpersist());
 
             _df.CreateTempView("view");

src/csharp/Microsoft.Spark.E2ETest/IpcTests/StorageLevelTests.cs

@@ -0,0 +1,60 @@
+using System;
+using System.Collections.Generic;
+using Microsoft.Spark.Sql;
+using Xunit;
+
+namespace Microsoft.Spark.E2ETest.IpcTests
+{
+    [Collection("Spark E2E Tests")]
+    public class StorageLevelTests
+    {
+        private readonly SparkSession _spark;
+        private readonly DataFrame _df;
+
+        public StorageLevelTests(SparkFixture fixture)
+        {
+            _spark = fixture.Spark;
+            _df = _spark.CreateDataFrame(new[] { "hello", "world" });
+        }
+
+        /// <summary>
+        /// Testing all public properties and methods of StorageLevel objects.
+        /// </summary>
+        [Fact]
+        public void TestStorageLevelProperties()
+        {
+            var storageLevels = new List<StorageLevel> {
+                StorageLevel.NONE,
+                StorageLevel.DISK_ONLY,
+                StorageLevel.DISK_ONLY_2,
+                StorageLevel.MEMORY_ONLY,
+                StorageLevel.MEMORY_ONLY_2,
+                StorageLevel.MEMORY_ONLY_SER,
+                StorageLevel.MEMORY_ONLY_SER_2,
+                StorageLevel.MEMORY_AND_DISK,
+                StorageLevel.MEMORY_AND_DISK_2,
+                StorageLevel.MEMORY_AND_DISK_SER,
+                StorageLevel.MEMORY_AND_DISK_SER_2,
+                StorageLevel.OFF_HEAP
+            };
+            foreach (StorageLevel expected in storageLevels)
+            {
+                _df.Persist(expected);
+                StorageLevel actual = _df.StorageLevel();
+                Assert.Equal(expected, actual);
+                // Needs to be unpersisted so other Persists can take effect.
+                _df.Unpersist();
+            }
+
+            StorageLevel storageLevel = StorageLevel.MEMORY_AND_DISK;
+            Assert.True(storageLevel.UseDisk);
+            Assert.True(storageLevel.UseMemory);
+            Assert.False(storageLevel.UseOffHeap);
+            Assert.True(storageLevel.Deserialized);
+            Assert.Equal(1, storageLevel.Replication);
+
+            Assert.IsType<string>(storageLevel.Description());
+            Assert.IsType<string>(storageLevel.ToString());
+        }
+    }
+}

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

@@ -14,10 +14,6 @@
 
 namespace Microsoft.Spark.Sql
 {
-    /// TODO:
-    /// Missing APIs:
-    /// Persist() with "StorageLevel"
-
     /// <summary>
     ///  A distributed collection of data organized into named columns.
     /// </summary>
@@ -808,17 +804,34 @@ public DataFrame Coalesce(int numPartitions) =>
         public DataFrame Distinct() => WrapAsDataFrame(_jvmObject.Invoke("distinct"));
 
         /// <summary>
-        /// Persist this `DataFrame` with the default storage level (`MEMORY_AND_DISK`).
+        /// Persist this <see cref="DataFrame"/> with the default storage level MEMORY_AND_DISK.
         /// </summary>
         /// <returns>DataFrame object</returns>
         public DataFrame Persist() => WrapAsDataFrame(_jvmObject.Invoke("persist"));
 
         /// <summary>
-        /// Persist this `DataFrame` with the default storage level (`MEMORY_AND_DISK`).
+        /// Persist this <see cref="DataFrame"/> with the given storage level.
+        /// </summary>
+        /// <param name="storageLevel">
+        /// <see cref="StorageLevel"/> to persist the <see cref="DataFrame"/> to.
+        /// </param>
+        /// <returns>DataFrame object</returns>
+        public DataFrame Persist(StorageLevel storageLevel) =>
+            WrapAsDataFrame(_jvmObject.Invoke("persist", storageLevel));
+
+        /// <summary>
+        /// Persist this <see cref="DataFrame"/> with the default storage level MEMORY_AND_DISK.
         /// </summary>
         /// <returns>DataFrame object</returns>
         public DataFrame Cache() => WrapAsDataFrame(_jvmObject.Invoke("cache"));
 
+        /// <summary>
+        /// Get the <see cref="DataFrame"/>'s current <see cref="StorageLevel"/>.
+        /// </summary>
+        /// <returns><see cref="StorageLevel"/> object</returns>
+        public StorageLevel StorageLevel() =>
+            new StorageLevel((JvmObjectReference)_jvmObject.Invoke("storageLevel"));
+
         /// <summary>
         /// Mark the Dataset as non-persistent, and remove all blocks for it from memory and disk.
         /// </summary>

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

@@ -0,0 +1,239 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+using System;
+using Microsoft.Spark.Interop;
+using Microsoft.Spark.Interop.Ipc;
+
+namespace Microsoft.Spark.Sql
+{
+    /// <summary>
+    /// Flags for controlling the storage of an RDD. Each StorageLevel records whether to use
+    /// memory, whether to drop the RDD to disk if it falls out of memory, whether to keep the
+    /// data in memory in a JAVA-specific serialized format, and whether to replicate the RDD
+    /// partitions on multiple nodes. Also contains static properties for some commonly used
+    /// storage levels, MEMORY_ONLY.
+    /// </summary>
+    public sealed class StorageLevel : IJvmObjectReferenceProvider
+    {
+        private static readonly string s_storageLevelClassName =
+            "org.apache.spark.storage.StorageLevel";
+        private static StorageLevel s_none;
+        private static StorageLevel s_diskOnly;
+        private static StorageLevel s_diskOnly2;
+        private static StorageLevel s_memoryOnly;
+        private static StorageLevel s_memoryOnly2;
+        private static StorageLevel s_memoryOnlySer;
+        private static StorageLevel s_memoryOnlySer2;
+        private static StorageLevel s_memoryAndDisk;
+        private static StorageLevel s_memoryAndDisk2;
+        private static StorageLevel s_memoryAndDiskSer;
+        private static StorageLevel s_memoryAndDiskSer2;
+        private static StorageLevel s_offHeap;
+        private readonly JvmObjectReference _jvmObject;
+        private bool? _useDisk;
+        private bool? _useMemory;
+        private bool? _useOffHeap;
+        private bool? _deserialized;
+        private int? _replication;
+
+        internal StorageLevel(JvmObjectReference jvmObject)
+        {
+            _jvmObject = jvmObject;
+        }
+
+        public StorageLevel(
+            bool useDisk,
+            bool useMemory,
+            bool useOffHeap,
+            bool deserialized,
+            int replication = 1)
+            : this(
+                (JvmObjectReference)SparkEnvironment.JvmBridge.CallStaticJavaMethod(
+                    s_storageLevelClassName,
+                    "apply",
+                    useDisk,
+                    useMemory,
+                    useOffHeap,
+                    deserialized,
+                    replication))
+        {
+            _useDisk = useDisk;
+            _useMemory = useMemory;
+            _useOffHeap = useOffHeap;
+            _deserialized = deserialized;
+            _replication = replication;
+        }
+
+        JvmObjectReference IJvmObjectReferenceProvider.Reference => _jvmObject;
+
+        /// <summary>
+        /// Returns the StorageLevel object with all parameters set to false.
+        /// </summary>
+        public static StorageLevel NONE =>
+            s_none ??= new StorageLevel(
+                (JvmObjectReference)SparkEnvironment.JvmBridge.CallStaticJavaMethod(
+                    s_storageLevelClassName,
+                    "NONE"));
+
+        /// <summary>
+        /// Returns the StorageLevel to Disk, serialized and replicated once.
+        /// </summary>
+        public static StorageLevel DISK_ONLY =>
+            s_diskOnly ??= new StorageLevel(
+                (JvmObjectReference)SparkEnvironment.JvmBridge.CallStaticJavaMethod(
+                    s_storageLevelClassName,
+                    "DISK_ONLY"));
+
+        /// <summary>
+        /// Returns the StorageLevel to Disk, serialized and replicated twice.
+        /// </summary>
+        public static StorageLevel DISK_ONLY_2 =>
+            s_diskOnly2 ??= new StorageLevel(
+                (JvmObjectReference)SparkEnvironment.JvmBridge.CallStaticJavaMethod(
+                    s_storageLevelClassName,
+                    "DISK_ONLY_2"));
+
+        /// <summary>
+        /// Returns the StorageLevel to Memory, deserialized and replicated once.
+        /// </summary>
+        public static StorageLevel MEMORY_ONLY =>
+            s_memoryOnly ??= new StorageLevel(
+                (JvmObjectReference)SparkEnvironment.JvmBridge.CallStaticJavaMethod(
+                    s_storageLevelClassName,
+                    "MEMORY_ONLY"));
+
+        /// <summary>
+        /// Returns the StorageLevel to Memory, deserialized and replicated twice.
+        /// </summary>
+        public static StorageLevel MEMORY_ONLY_2 =>
+            s_memoryOnly2 ??= new StorageLevel(
+                (JvmObjectReference)SparkEnvironment.JvmBridge.CallStaticJavaMethod(
+                    s_storageLevelClassName,
+                    "MEMORY_ONLY_2"));
+
+        /// <summary>
+        /// Returns the StorageLevel to Memory, serialized and replicated once.
+        /// </summary>
+        public static StorageLevel MEMORY_ONLY_SER =>
+            s_memoryOnlySer ??= new StorageLevel(
+                (JvmObjectReference)SparkEnvironment.JvmBridge.CallStaticJavaMethod(
+                    s_storageLevelClassName,
+                    "MEMORY_ONLY_SER"));
+
+        /// <summary>
+        /// Returns the StorageLevel to Memory, serialized and replicated twice.
+        /// </summary>
+        public static StorageLevel MEMORY_ONLY_SER_2 =>
+            s_memoryOnlySer2 ??= new StorageLevel(
+                (JvmObjectReference)SparkEnvironment.JvmBridge.CallStaticJavaMethod(
+                    s_storageLevelClassName,
+                    "MEMORY_ONLY_SER_2"));
+
+        /// <summary>
+        /// Returns the StorageLevel to Disk and Memory, deserialized and replicated once.
+        /// </summary>
+        public static StorageLevel MEMORY_AND_DISK =>
+            s_memoryAndDisk ??= new StorageLevel(
+                (JvmObjectReference)SparkEnvironment.JvmBridge.CallStaticJavaMethod(
+                    s_storageLevelClassName,
+                    "MEMORY_AND_DISK"));
+
+        /// <summary>
+        /// Returns the StorageLevel to Disk and Memory, deserialized and replicated twice.
+        /// </summary>
+        public static StorageLevel MEMORY_AND_DISK_2 =>
+            s_memoryAndDisk2 ??= new StorageLevel(
+                (JvmObjectReference)SparkEnvironment.JvmBridge.CallStaticJavaMethod(
+                    s_storageLevelClassName,
+                    "MEMORY_AND_DISK_2"));
+
+        /// <summary>
+        /// Returns the StorageLevel to Disk and Memory, serialized and replicated once.
+        /// </summary>
+        public static StorageLevel MEMORY_AND_DISK_SER =>
+            s_memoryAndDiskSer ??= new StorageLevel(
+                (JvmObjectReference)SparkEnvironment.JvmBridge.CallStaticJavaMethod(
+                    s_storageLevelClassName,
+                    "MEMORY_AND_DISK_SER"));
+
+        /// <summary>
+        /// Returns the StorageLevel to Disk and Memory, serialized and replicated twice.
+        /// </summary>
+        public static StorageLevel MEMORY_AND_DISK_SER_2 =>
+            s_memoryAndDiskSer2 ??= new StorageLevel(
+                (JvmObjectReference)SparkEnvironment.JvmBridge.CallStaticJavaMethod(
+                    s_storageLevelClassName,
+                    "MEMORY_AND_DISK_SER_2"));
+
+        /// <summary>
+        /// Returns the StorageLevel to Disk, Memory and Offheap, serialized and replicated once.
+        /// </summary>
+        public static StorageLevel OFF_HEAP =>
+            s_offHeap ??= new StorageLevel(
+                (JvmObjectReference)SparkEnvironment.JvmBridge.CallStaticJavaMethod(
+                    s_storageLevelClassName,
+                    "OFF_HEAP"));
+
+        /// <summary>
+        /// Returns bool value of UseDisk of this StorageLevel.
+        /// </summary>
+        public bool UseDisk => _useDisk ??= (bool)_jvmObject.Invoke("useDisk");
+
+        /// <summary>
+        /// Returns bool value of UseMemory of this StorageLevel.
+        /// </summary>
+        public bool UseMemory => _useMemory ??= (bool)_jvmObject.Invoke("useMemory");
+
+        /// <summary>
+        /// Returns bool value of UseOffHeap of this StorageLevel.
+        /// </summary>
+        public bool UseOffHeap => _useOffHeap ??= (bool)_jvmObject.Invoke("useOffHeap");
+
+        /// <summary>
+        /// Returns bool value of Deserialized of this StorageLevel.
+        /// </summary>
+        public bool Deserialized => _deserialized ??= (bool)_jvmObject.Invoke("deserialized");
+
+        /// <summary>
+        /// Returns int value of Replication of this StorageLevel.
+        /// </summary>
+        public int Replication => _replication ??= (int)_jvmObject.Invoke("replication");
+
+        /// <summary>
+        /// Returns the description string of this StorageLevel.
+        /// </summary>
+        /// <returns>Description as string.</returns>
+        public string Description() => (string)_jvmObject.Invoke("description");
+
+        /// <summary>
+        /// Returns the string representation of this StorageLevel.
+        /// </summary>
+        /// <returns>representation as string value.</returns>
+        public override string ToString() => (string)_jvmObject.Invoke("toString");
+
+        /// <summary>
+        /// Checks if the given object is same as the current object.
+        /// </summary>
+        /// <param name="obj">Other object to compare against</param>
+        /// <returns>True if the other object is equal.</returns>
+        public override bool Equals(object obj)
+        {
+            if (!(obj is StorageLevel that))
+            {
+                return false;
+            }
+
+            return (UseDisk == that.UseDisk) && (UseMemory == that.UseMemory) &&
+                (UseOffHeap == that.UseOffHeap) && (Deserialized == that.Deserialized) &&
+                (Replication == that.Replication);
+        }
+
+        /// <summary>
+        /// Returns the hash code of the current object.
+        /// </summary>
+        /// <returns>The hash code of the current object</returns>
+        public override int GetHashCode() => base.GetHashCode();
+    }
+}