From c707e234b008af62c1f416dc861933e764e71155 Mon Sep 17 00:00:00 2001
From: Nicolas Laval <nicolas.laval@making-sense.info>
Date: Wed, 24 Jan 2024 21:25:10 +0100
Subject: [PATCH] Return all ds in bindings

---
 .../java/fr/insee/trevas/lab/utils/Utils.java | 30 ++++++++++++++-----
 1 file changed, 23 insertions(+), 7 deletions(-)

diff --git a/src/main/java/fr/insee/trevas/lab/utils/Utils.java b/src/main/java/fr/insee/trevas/lab/utils/Utils.java
index b6a2782..5ae2c86 100644
--- a/src/main/java/fr/insee/trevas/lab/utils/Utils.java
+++ b/src/main/java/fr/insee/trevas/lab/utils/Utils.java
@@ -41,7 +41,14 @@ public static ScriptEngine initEngineWithSpark(Bindings bindings, SparkSession s
     public static Bindings getBindings(Bindings input) {
         Bindings output = new SimpleBindings();
         input.forEach((k, v) -> {
-            if (v instanceof PersistentDataset) output.put(k, v);
+            if (!k.startsWith("$")) {
+                if (v instanceof PersistentDataset) {
+                    output.put(k + "$PersistentDataset", v);
+                } else {
+                    output.put(k, v);
+                }
+            }
+            ;
         });
         return output;
     }
@@ -69,17 +76,26 @@ public static SparkConf loadSparkConfig(String stringPath) {
     public static Bindings getSparkBindings(Bindings input, Integer limit) {
         Bindings output = new SimpleBindings();
         input.forEach((k, v) -> {
-            if (v instanceof PersistentDataset) {
-                fr.insee.vtl.model.Dataset ds = ((PersistentDataset) v).getDelegate();
-                if (ds instanceof SparkDataset) {
-                    Dataset<Row> sparkDs = ((SparkDataset) ds).getSparkDataset();
+            if (!k.startsWith("$")) {
+                if ((v instanceof PersistentDataset) || (v instanceof SparkDataset)) {
+                    String name = k;
+                    SparkDataset spDs = null;
+                    if (v instanceof PersistentDataset) {
+                        fr.insee.vtl.model.Dataset ds = ((PersistentDataset) v).getDelegate();
+                        name = name + "$PersistentDataset";
+                        spDs = (SparkDataset) ds;
+                    }
+                    if (v instanceof SparkDataset) {
+                        spDs = (SparkDataset) v;
+                    }
+                    Dataset<Row> sparkDs = (spDs).getSparkDataset();
                     if (limit != null) {
                         SparkDataset sparkDataset = new SparkDataset(sparkDs.limit(limit));
                         InMemoryDataset im = new InMemoryDataset(
                                 sparkDataset.getDataPoints(),
                                 sparkDataset.getDataStructure());
-                        output.put(k, im);
-                    } else output.put(k, new SparkDataset(sparkDs)); // useless
+                        output.put(name, im);
+                    } else output.put(name, new SparkDataset(sparkDs)); // useless
                 }
             }
         });