fixed up zframe selects

zinggAI · Oct 23, 2023 · 008e95d · 008e95d
1 parent ab9ade1
commit 008e95d
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 4 deletions.
diff --git a/common/client/src/main/java/zingg/common/client/ZFrame.java b/common/client/src/main/java/zingg/common/client/ZFrame.java
@@ -16,8 +16,7 @@ public interface ZFrame<D, R, C> {
     public String[] columns();
     public ZFrame<D, R, C> select(String... string);
     public ZFrame<D, R, C> select(List<C> cols);
-    public ZFrame<D, R, C> select(String col, String... cols);
-    public ZFrame<D, R, C> select(String col);
+    public ZFrame<D, R, C> select(C col);
     public ZFrame<D, R, C> selectExpr(String... col);
     public ZFrame <D, R, C> distinct();
     public List<R> collectAsList();

diff --git a/spark/client/src/main/java/zingg/spark/client/SparkFrame.java b/spark/client/src/main/java/zingg/spark/client/SparkFrame.java
@@ -44,14 +44,18 @@ public ZFrame<Dataset<Row>, Row, Column> select(Column... cols) {
         return new SparkFrame(df.select(cols));
     }
 
+    public ZFrame<Dataset<Row>, Row, Column> select(Column col) {
+        return new SparkFrame(df.select(col));
+    }
+
 
     public ZFrame<Dataset<Row>, Row, Column> select(List<Column> cols){
         return new SparkFrame(df.select(JavaConverters.asScalaIteratorConverter(cols.iterator()).asScala().toSeq()));
     }
 
 
-    public ZFrame<Dataset<Row>, Row, Column> select(String col) {
-        return new SparkFrame(df.select(col));
+    public ZFrame<Dataset<Row>, Row, Column> select(String... col) {
+        return toDF(col);
     }
 
     public ZFrame<Dataset<Row>, Row, Column> selectExpr(String... col) {