Fix datasetToDisplay

InseeFrLab · Jul 10, 2024 · 832e744 · 832e744
1 parent 42470eb
commit 832e744
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 18 deletions.
diff --git a/src/main/java/fr/insee/trevas/jupyter/DatasetUtils.java b/src/main/java/fr/insee/trevas/jupyter/DatasetUtils.java
@@ -10,12 +10,21 @@
 
 public class DatasetUtils {
 
-	public static String datasetToDisplay(Dataset dataset) {
+	public static String datasetToDisplay(Dataset dataset) throws ClassNotFoundException {
 		Map<String, Dataset.Role> roles = dataset.getDataStructure().getRoles();
-		Dataset reducedDataset = new SparkDataset(
-				((SparkDataset) ((PersistentDataset) dataset).getDelegate())
-						.getSparkDataset()
-						.limit(50), roles);
+		SparkDataset sparkDs;
+		if (dataset instanceof PersistentDataset) {
+			sparkDs = ((SparkDataset) ((PersistentDataset) dataset).getDelegate());
+		} else if (dataset instanceof SparkDataset) {
+			sparkDs = (SparkDataset) dataset;
+		} else {
+			throw new ClassNotFoundException(dataset.getClass().getName() +
+					" unhandled in datasetToDisplay method");
+		}
+		var reducedDataset = new SparkDataset(
+				sparkDs.getSparkDataset().limit(50),
+				roles
+		);
 		var b = new StringBuilder();
 		b.append("<table id='dataset_").append(reducedDataset.hashCode()).append("' class='display'>");
 		b.append("<thead>");

diff --git a/src/main/java/fr/insee/trevas/jupyter/VtlKernel.java b/src/main/java/fr/insee/trevas/jupyter/VtlKernel.java
@@ -107,22 +107,16 @@ public static String getSize(Dataset ds) {
 		return "Dataset size: " + sparkDataset.getDataPoints().size();
 	}
 
-	public static Object show(Object o) {
+	public static Object show(Object o) throws ClassNotFoundException {
 		if (o instanceof Dataset) {
-			SparkDataset dataset = asSparkDataset((Dataset) o);
-			var roles =
-					dataset.getDataStructure().entrySet().stream()
-							.collect(
-									Collectors.toMap(
-											Map.Entry::getKey, e -> e.getValue().getRole()));
-			showDataset(new SparkDataset(dataset.getSparkDataset().limit(50), roles));
+			showDataset((SparkDataset) o);
 		} else {
 			displayData.putText(o.toString());
 		}
 		return o;
 	}
 
-	private static void showDataset(Dataset dataset) {
+	private static void showDataset(Dataset dataset) throws ClassNotFoundException {
 		displayData.putHTML(DatasetUtils.datasetToDisplay(dataset));
 	}
 
@@ -209,10 +203,14 @@ public static void runSDMX(String path, String data) {
 
 		results.forEach(
 				(k, v) -> {
-					result.append("<h2>")
-							.append(k)
-							.append("</h2>")
-							.append(DatasetUtils.datasetToDisplay(v));
+					try {
+						result.append("<h2>")
+								.append(k)
+								.append("</h2>")
+								.append(DatasetUtils.datasetToDisplay(v));
+					} catch (ClassNotFoundException e) {
+						throw new RuntimeException(e);
+					}
 				});
 
 		displayData.putHTML(result.toString());