Handle PersistentDataset

InseeFrLab · Jan 24, 2024 · 2eeefe1 · 2eeefe1
1 parent e1c7ae0
commit 2eeefe1
Showing 1 changed file with 10 additions and 2 deletions.
diff --git a/src/main/java/fr/insee/trevas/jupyter/VtlKernel.java b/src/main/java/fr/insee/trevas/jupyter/VtlKernel.java
@@ -2,6 +2,7 @@
 
 import fr.insee.vtl.engine.VtlScriptEngine;
 import fr.insee.vtl.model.Dataset;
+import fr.insee.vtl.model.PersistentDataset;
 import fr.insee.vtl.model.Structured;
 import fr.insee.vtl.spark.SparkDataset;
 import io.github.spencerpark.jupyter.channels.JupyterConnection;
@@ -58,9 +59,16 @@ private static Map<String, Dataset.Role> getRoleMap(fr.insee.vtl.model.Dataset d
     private static SparkDataset asSparkDataset(Dataset dataset) {
         if (dataset instanceof SparkDataset) {
             return (SparkDataset) dataset;
-        } else {
-            return new SparkDataset(dataset, getRoleMap(dataset), spark);
         }
+        if (dataset instanceof PersistentDataset) {
+            fr.insee.vtl.model.Dataset ds = ((PersistentDataset) dataset).getDelegate();
+            if (ds instanceof SparkDataset) {
+                return (SparkDataset) ds;
+            } else {
+                return new SparkDataset(ds, getRoleMap(dataset), spark);
+            }
+        }
+        throw new IllegalArgumentException("Unknow dataset type");
     }
 
     public static SparkDataset loadParquet(String path) throws Exception {