Introducing an Item UDT

src/main/java/org/rumbledb/items/parsing/ItemUserDefinedType.java

@@ -0,0 +1,54 @@
+package org.rumbledb.items.parsing;
+
+import org.apache.spark.sql.types.DataType;
+import org.apache.spark.sql.types.DataTypes;
+import org.apache.spark.sql.types.UserDefinedType;
+import org.rumbledb.api.Item;
+import org.rumbledb.exceptions.OurBadException;
+import org.rumbledb.runtime.flwor.udfs.DataFrameContext;
+
+public class ItemUserDefinedType extends UserDefinedType<Item> {
+
+    DataFrameContext dataFrameContext;
+
+    public void initialize() {
+        this.dataFrameContext = new DataFrameContext();
+    }
+
+    @Override
+    public DataType sqlType() {
+        return DataTypes.BinaryType;
+    }
+
+    @Override
+    public byte[] serialize(Item item) {
+        if (dataFrameContext == null) {
+            initialize();
+        }
+        this.dataFrameContext.getOutput().clear();
+        this.dataFrameContext.getKryo().writeClassAndObject(this.dataFrameContext.getOutput(), item);
+        byte[] s = this.dataFrameContext.getOutput().toBytes();
+        return s;
+    }
+
+    @Override
+    public Item deserialize(Object object) {
+        if (dataFrameContext == null) {
+            initialize();
+        }
+        if (object instanceof byte[] b) {
+            this.dataFrameContext.getInput().setBuffer(b);
+            Object result = this.dataFrameContext.getKryo().readClassAndObject(this.dataFrameContext.getInput());
+            if (result instanceof Item) {
+                return (Item) result;
+            }
+            throw new OurBadException("Expected an item but got " + result.getClass());
+        }
+        throw new OurBadException("Expected a byte array but got " + object.getClass());
+    }
+
+    @Override
+    public Class<Item> userClass() {
+        return Item.class;
+    }
+}

src/main/java/org/rumbledb/runtime/HybridRuntimeIterator.java

@@ -21,6 +21,7 @@
 package org.rumbledb.runtime;
 
 import org.apache.spark.api.java.JavaRDD;
+import org.apache.spark.sql.Dataset;
 import org.apache.spark.sql.Row;
 import org.rumbledb.api.Item;
 import org.rumbledb.config.RumbleRuntimeConfiguration;
@@ -31,6 +32,7 @@
 import org.rumbledb.exceptions.IteratorFlowException;
 import org.rumbledb.exceptions.MoreThanOneItemException;
 import org.rumbledb.exceptions.NoItemException;
+import org.rumbledb.exceptions.OurBadException;
 import org.rumbledb.expressions.ExecutionMode;
 import org.rumbledb.items.parsing.RowToItemMapper;
 import org.rumbledb.items.structured.JSoundDataFrame;
@@ -57,6 +59,10 @@ protected boolean implementsDataFrames() {
         return false;
     }
 
+    protected boolean implementsItemDataFrames() {
+        return false;
+    }
+
     protected boolean implementsLocal() {
         return true;
     }
@@ -66,7 +72,9 @@ protected boolean implementsRDD() {
     }
 
     protected void fallbackToRDDIfDFNotImplemented(ExecutionMode executionMode) {
-        if (executionMode == ExecutionMode.DATAFRAME && !this.implementsDataFrames()) {
+        if (
+            executionMode == ExecutionMode.DATAFRAME && !this.implementsDataFrames() && !this.implementsItemDataFrames()
+        ) {
             this.staticContext = this.staticContext.withExecutionMode(ExecutionMode.RDD);
         }
     }
@@ -106,15 +114,7 @@ public boolean hasNext() {
         }
         if (this.result == null) {
             this.currentResultIndex = 0;
-            JavaRDD<Item> rdd = null;
-            if (!isRDD() && implementsDataFrames()) {
-                rdd = dataFrameToRDDOfItems(
-                    this.getDataFrame(this.currentDynamicContextForLocalExecution),
-                    this.getMetadata()
-                );
-            } else {
-                rdd = this.getRDDAux(this.currentDynamicContextForLocalExecution);
-            }
+            JavaRDD<Item> rdd = this.getRDD(this.currentDynamicContextForLocalExecution);
             this.result = collectRDDwithLimit(rdd, this.getConfiguration(), this.getMetadata());
             this.hasNext = !this.result.isEmpty();
         }
@@ -153,7 +153,12 @@ public JavaRDD<Item> getRDD(DynamicContext context) {
             return dataFrameToRDDOfItems(df, getMetadata());
         }
         if (isRDDOrDataFrame()) {
-            return getRDDAux(context);
+            if (implementsItemDataFrames()) {
+                Dataset<Row> dataFrame = getItemDataFrame(context);
+                return dataFrame.select("item").toJavaRDD().map(row -> (Item) row.getAs("item"));
+            } else {
+                return getRDDAux(context);
+            }
         }
         List<Item> contents = this.materialize(context);
         return SparkSessionManager.getInstance().getJavaSparkContext().parallelize(contents);
@@ -261,7 +266,9 @@ public Item materializeAtMostOneItemOrNull(
         throw new MoreThanOneItemException();
     }
 
-    protected abstract JavaRDD<Item> getRDDAux(DynamicContext context);
+    protected JavaRDD<Item> getRDDAux(DynamicContext context) {
+        throw new OurBadException("getRDDAux not implemented for " + this.getClass().getSimpleName());
+    }
 
     protected abstract void openLocal();
 

src/main/java/org/rumbledb/runtime/ItemDataFrameRuntimeIterator.java

@@ -0,0 +1,79 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ * Authors: Stefan Irimescu, Can Berker Cikis
+ *
+ */
+
+package org.rumbledb.runtime;
+
+import org.rumbledb.api.Item;
+import org.rumbledb.context.RuntimeStaticContext;
+import org.rumbledb.exceptions.OurBadException;
+
+import java.util.List;
+
+public abstract class ItemDataFrameRuntimeIterator extends HybridRuntimeIterator {
+
+    private static final long serialVersionUID = 1L;
+
+    protected ItemDataFrameRuntimeIterator(
+            List<RuntimeIterator> children,
+            RuntimeStaticContext staticContext
+    ) {
+        super(children, staticContext);
+    }
+
+    @Override
+    protected boolean implementsItemDataFrames() {
+        return true;
+    }
+
+    @Override
+    protected boolean implementsLocal() {
+        return false;
+    }
+
+    @Override
+    protected boolean implementsRDD() {
+        return false;
+    }
+
+    @Override
+    protected void openLocal() {
+        throw new OurBadException("Local evaluation are not implemented for the iterator", getMetadata());
+    }
+
+    @Override
+    protected void closeLocal() {
+        throw new OurBadException("Local evaluation are not implemented for the iterator", getMetadata());
+    }
+
+    @Override
+    protected void resetLocal() {
+        throw new OurBadException("Local evaluation are not implemented for the iterator", getMetadata());
+    }
+
+    @Override
+    protected boolean hasNextLocal() {
+        throw new OurBadException("Local evaluation are not implemented for the iterator", getMetadata());
+    }
+
+    @Override
+    protected Item nextLocal() {
+        throw new OurBadException("Local evaluation are not implemented for the iterator", getMetadata());
+    }
+}

src/main/java/org/rumbledb/runtime/RuntimeIterator.java

@@ -34,6 +34,10 @@
 import java.util.TreeMap;
 
 import org.apache.spark.api.java.JavaRDD;
+import org.apache.spark.sql.Dataset;
+import org.apache.spark.sql.Row;
+import org.apache.spark.sql.RowFactory;
+import org.apache.spark.sql.types.StructType;
 import org.rumbledb.api.Item;
 import org.rumbledb.config.RumbleRuntimeConfiguration;
 import org.rumbledb.context.DynamicContext;
@@ -51,6 +55,7 @@
 import org.rumbledb.exceptions.RumbleException;
 import org.rumbledb.expressions.ExecutionMode;
 import org.rumbledb.expressions.comparison.ComparisonExpression.ComparisonOperator;
+import org.rumbledb.items.parsing.ItemUserDefinedType;
 import org.rumbledb.items.structured.JSoundDataFrame;
 import org.rumbledb.runtime.flwor.NativeClauseContext;
 import org.rumbledb.runtime.misc.ComparisonIterator;
@@ -65,9 +70,16 @@
 import com.esotericsoftware.kryo.io.Input;
 import com.esotericsoftware.kryo.io.Output;
 
+import sparksoniq.spark.SparkSessionManager;
+
 public abstract class RuntimeIterator implements RuntimeIteratorInterface, KryoSerializable {
 
     protected static final String FLOW_EXCEPTION_MESSAGE = "Invalid next() call; ";
+    private static final StructType ITEM_DATA_FRAME_SCHEMA = new StructType().add(
+        "item",
+        new ItemUserDefinedType(),
+        false
+    );
     private static final long serialVersionUID = 1L;
     protected transient boolean hasNext;
     protected transient boolean isOpen;
@@ -340,6 +352,15 @@ public JSoundDataFrame getDataFrame(DynamicContext context) {
         );
     }
 
+    public Dataset<Row> getItemDataFrame(DynamicContext context) {
+        JavaRDD<Row> rdd = getRDD(context).map(RowFactory::create);;
+        return SparkSessionManager.getInstance().getOrCreateSession().createDataFrame(rdd, getItemDataFrameSchema());
+    }
+
+    protected static StructType getItemDataFrameSchema() {
+        return ITEM_DATA_FRAME_SCHEMA;
+    }
+
     /**
      * Gets the output as a DataFrame. If necessary and possible, forcibly converts the items to a DataFrame.
      * 

src/main/java/org/rumbledb/runtime/functions/input/JsonLinesFunctionIterator.java

@@ -21,6 +21,9 @@
 package org.rumbledb.runtime.functions.input;
 
 import org.apache.spark.api.java.JavaRDD;
+import org.apache.spark.sql.Dataset;
+import org.apache.spark.sql.Row;
+import org.apache.spark.sql.RowFactory;
 import org.rumbledb.api.Item;
 import org.rumbledb.context.DynamicContext;
 import org.rumbledb.context.RuntimeStaticContext;
@@ -65,7 +68,17 @@ public JsonLinesFunctionIterator(
     }
 
     @Override
-    public JavaRDD<Item> getRDDAux(DynamicContext context) {
+    protected boolean implementsItemDataFrames() {
+        return true;
+    }
+
+    @Override
+    protected boolean implementsRDD() {
+        return false;
+    }
+
+    @Override
+    public Dataset<Row> getItemDataFrame(DynamicContext context) {
         String url = this.children.get(0).materializeFirstItemOrNull(context).getStringValue();
         URI uri = FileSystemUtil.resolveURI(this.staticURI, url, getMetadata());
 
@@ -123,9 +136,12 @@ public JavaRDD<Item> getRDDAux(DynamicContext context) {
                     );
             }
         }
-        return strings.mapPartitions(
+        JavaRDD<Row> rows = strings.mapPartitions(
             new JSONSyntaxToItemMapper(getMetadata(), this.getRuntimeStaticContext().isQuerySideEffecting())
-        );
+        ).map(RowFactory::create);
+        return SparkSessionManager.getInstance()
+            .getOrCreateSession()
+            .createDataFrame(rows, getItemDataFrameSchema());
     }
 
     protected void init() {

src/main/java/org/rumbledb/runtime/functions/input/UnparsedTextLinesFunctionIterator.java

@@ -21,12 +21,15 @@
 package org.rumbledb.runtime.functions.input;
 
 import org.apache.spark.api.java.JavaRDD;
+import org.apache.spark.sql.Dataset;
+import org.apache.spark.sql.Row;
+import org.apache.spark.sql.RowFactory;
 import org.rumbledb.api.Item;
 import org.rumbledb.context.DynamicContext;
 import org.rumbledb.context.RuntimeStaticContext;
 import org.rumbledb.exceptions.CannotRetrieveResourceException;
-import org.rumbledb.items.parsing.StringToStringItemMapper;
-import org.rumbledb.runtime.RDDRuntimeIterator;
+import org.rumbledb.items.ItemFactory;
+import org.rumbledb.runtime.ItemDataFrameRuntimeIterator;
 import org.rumbledb.runtime.RuntimeIterator;
 
 import sparksoniq.spark.SparkSessionManager;
@@ -39,7 +42,7 @@
 import java.util.ArrayList;
 import java.util.List;
 
-public class UnparsedTextLinesFunctionIterator extends RDDRuntimeIterator {
+public class UnparsedTextLinesFunctionIterator extends ItemDataFrameRuntimeIterator {
 
     private static final long serialVersionUID = 1L;
     public static final int MIN_PARTITIONS = 10;
@@ -52,13 +55,16 @@ public UnparsedTextLinesFunctionIterator(
     }
 
     @Override
-    public JavaRDD<Item> getRDDAux(DynamicContext context) {
+    public Dataset<Row> getItemDataFrame(DynamicContext context) {
         RuntimeIterator urlIterator = this.children.get(0);
         Item url = urlIterator.materializeFirstItemOrNull(context);
         if (url == null) {
-            return SparkSessionManager.getInstance()
+            JavaRDD<Row> emptyRows = SparkSessionManager.getInstance()
                 .getJavaSparkContext()
                 .emptyRDD();
+            return SparkSessionManager.getInstance()
+                .getOrCreateSession()
+                .createDataFrame(emptyRows, getItemDataFrameSchema());
         }
         URI uri = FileSystemUtil.resolveURI(this.staticURI, url.getStringValue(), getMetadata());
         int partitions = -1;
@@ -116,6 +122,9 @@ public JavaRDD<Item> getRDDAux(DynamicContext context) {
                 partitionsIterator.close();
             }
         }
-        return strings.mapPartitions(new StringToStringItemMapper());
+        JavaRDD<Row> rows = strings.map(line -> RowFactory.create(ItemFactory.getInstance().createStringItem(line)));
+        return SparkSessionManager.getInstance()
+            .getOrCreateSession()
+            .createDataFrame(rows, getItemDataFrameSchema());
     }
 }

src/main/java/org/rumbledb/runtime/functions/sequences/aggregate/CountFunctionIterator.java

@@ -120,8 +120,8 @@ private static Item computeRDD(
             DynamicContext context,
             ExceptionMetadata metadata
     ) {
-        long count = iterator.getRDD(context).count();
-        if (count > (long) Integer.MAX_VALUE) {
+        long count = iterator.getItemDataFrame(context).count();
+        if (count > Long.MAX_VALUE) {
             throw new OurBadException("The count value is too big to convert to integer type.");
         } else {
             return ItemFactory.getInstance().createLongItem(count);
@@ -134,7 +134,7 @@ private static Item computeDataFrame(
             ExceptionMetadata metadata
     ) {
         long count = iterator.getDataFrame(context).count();
-        if (count > (long) Integer.MAX_VALUE) {
+        if (count > Long.MAX_VALUE) {
             throw new OurBadException("The count value is too big to convert to integer type.");
         } else {
             return ItemFactory.getInstance().createLongItem(count);