Fixed 'select *' case when ColumnarBatch is used (#263)

GoogleCloudDataproc · Nov 6, 2020 · 4c2ec94 · 4c2ec94
1 parent 4d3bd1b
commit 4c2ec94
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 1 deletion.
diff --git a/connector/src/main/java/com/google/cloud/spark/bigquery/v2/BigQueryDataSourceReader.java b/connector/src/main/java/com/google/cloud/spark/bigquery/v2/BigQueryDataSourceReader.java
@@ -150,14 +150,25 @@ public List<InputPartition<ColumnarBatch>> planBatchInputPartitions() {
         readSessionCreator.create(
             tableId, selectedFields, filter, readSessionCreatorConfig.getMaxParallelism());
     ReadSession readSession = readSessionResponse.getReadSession();
+
+    if (selectedFields.isEmpty()) {
+      // means select *
+      Schema tableSchema = readSessionResponse.getReadTableInfo().getDefinition().getSchema();
+      selectedFields =
+          tableSchema.getFields().stream()
+              .map(Field::getName)
+              .collect(ImmutableList.toImmutableList());
+    }
+
+    ImmutableList<String> partitionSelectedFields = selectedFields;
     return readSession.getStreamsList().stream()
         .map(
             stream ->
                 new ArrowInputPartition(
                     bigQueryReadClientFactory,
                     stream.getName(),
                     readSessionCreatorConfig.getMaxReadRowsRetries(),
-                    selectedFields,
+                    partitionSelectedFields,
                     readSessionResponse))
         .collect(Collectors.toList());
   }

diff --git a/...ctor/src/test/scala/com/google/cloud/spark/bigquery/it/SparkBigQueryEndToEndITSuite.scala b/...ctor/src/test/scala/com/google/cloud/spark/bigquery/it/SparkBigQueryEndToEndITSuite.scala
@@ -163,6 +163,18 @@ class SparkBigQueryEndToEndITSuite extends FunSuite
       assert(row(1).isInstanceOf[String])
     }
 
+    test("select all columns from a table. DataSource %s. Data Format %s"
+      .format(dataSourceFormat, dataFormat)) {
+      val row = spark.read.format(dataSourceFormat)
+        .option("table", SHAKESPEARE_TABLE)
+        .option("readDataFormat", dataFormat).load()
+        .select("word_count", "word", "corpus", "corpus_date").head
+      assert(row(0).isInstanceOf[Long])
+      assert(row(1).isInstanceOf[String])
+      assert(row(2).isInstanceOf[String])
+      assert(row(3).isInstanceOf[Long])
+    }
+
     test("cache data frame in DataSource %s. Data Format %s".format(dataSourceFormat, dataFormat)) {
       val allTypesTable = readAllTypesTable("bigquery")
       writeToBigQuery(allTypesTable, SaveMode.Overwrite, "avro")