diff --git a/src/main/java/fr/insee/trevas/jupyter/VtlKernel.java b/src/main/java/fr/insee/trevas/jupyter/VtlKernel.java index 0371844..588a0cd 100644 --- 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 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 {