drill-commits mailing list archives

Site index · List index
Message view « Date » · « Thread »
Top « Date » · « Thread »
From j..@apache.org
Subject [06/13] drill git commit: DRILL-4735: ConvertCountToDirectScan rule enhancements
Date Tue, 15 Aug 2017 13:44:18 GMT
http://git-wip-us.apache.org/repos/asf/drill/blob/8b564235/exec/java-exec/src/main/java/org/apache/drill/exec/store/pojo/PojoWriters.java
----------------------------------------------------------------------
diff --git a/exec/java-exec/src/main/java/org/apache/drill/exec/store/pojo/PojoWriters.java
b/exec/java-exec/src/main/java/org/apache/drill/exec/store/pojo/PojoWriters.java
new file mode 100644
index 0000000..090f32f
--- /dev/null
+++ b/exec/java-exec/src/main/java/org/apache/drill/exec/store/pojo/PojoWriters.java
@@ -0,0 +1,296 @@
+/*
+ * 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.
+ */
+package org.apache.drill.exec.store.pojo;
+
+import io.netty.buffer.DrillBuf;
+
+import java.sql.Timestamp;
+
+import org.apache.drill.common.exceptions.ExecutionSetupException;
+import org.apache.drill.common.types.TypeProtos.MinorType;
+import org.apache.drill.common.types.Types;
+import org.apache.drill.exec.expr.holders.NullableVarCharHolder;
+import org.apache.drill.exec.vector.BigIntVector;
+import org.apache.drill.exec.vector.BitVector;
+import org.apache.drill.exec.vector.Float8Vector;
+import org.apache.drill.exec.vector.IntVector;
+import org.apache.drill.exec.vector.NullableBigIntVector;
+import org.apache.drill.exec.vector.NullableBitVector;
+import org.apache.drill.exec.vector.NullableFloat8Vector;
+import org.apache.drill.exec.vector.NullableIntVector;
+import org.apache.drill.exec.vector.NullableTimeStampVector;
+import org.apache.drill.exec.vector.NullableVarCharVector;
+
+import com.google.common.base.Charsets;
+
+public class PojoWriters {
+
+  /**
+   * Creates matching writer to the given field type.
+   *
+   * @param type field type
+   * @param fieldName field name
+   * @param buffer drill buffer
+   * @return pojo writer
+   * @throws ExecutionSetupException in case if writer was not found for the given type
+   */
+  public static PojoWriter getWriter(Class<?> type, String fieldName, DrillBuf buffer)
throws ExecutionSetupException {
+
+    if (type == Integer.class) {
+      return new NIntWriter(fieldName);
+    } else if (type == Long.class) {
+      return new NBigIntWriter(fieldName);
+    } else if (type == Boolean.class) {
+      return new NBooleanWriter(fieldName);
+    } else if (type == Double.class) {
+      return new NDoubleWriter(fieldName);
+    } else if (type.isEnum()) {
+      return new EnumWriter(fieldName, buffer);
+    } else if (type == String.class) {
+      return new StringWriter(fieldName, buffer);
+    } else if (type == Timestamp.class) {
+      return new NTimeStampWriter(fieldName);
+      // primitives
+    } else if (type == int.class) {
+      return new IntWriter(fieldName);
+    } else if (type == double.class) {
+      return new DoubleWriter(fieldName);
+    } else if (type == boolean.class) {
+      return new BitWriter(fieldName);
+    } else if (type == long.class) {
+      return new LongWriter(fieldName);
+    }
+
+    throw new ExecutionSetupException(String.format("PojoRecordReader doesn't yet support
conversions from the type [%s].", type));
+  }
+
+  /**
+   * Pojo writer for int. Does not expect to write null value.
+   */
+  public static class IntWriter extends AbstractPojoWriter<IntVector> {
+
+    public IntWriter(String fieldName) {
+      super(fieldName, Types.required(MinorType.INT));
+    }
+
+    @Override
+    public void writeField(Object value, int outboundIndex) {
+      vector.getMutator().setSafe(outboundIndex, (int) value);
+    }
+  }
+
+  /**
+   * Pojo writer for boolean. Does not expect to write null value.
+   */
+  public static class BitWriter extends AbstractPojoWriter<BitVector> {
+
+    public BitWriter(String fieldName) {
+      super(fieldName, Types.required(MinorType.BIT));
+    }
+
+    @Override
+    public void writeField(Object value, int outboundIndex) {
+      vector.getMutator().setSafe(outboundIndex, (boolean) value ? 1 : 0);
+    }
+
+  }
+
+  /**
+   * Pojo writer for long. Does not expect to write null value.
+   */
+  public static class LongWriter extends AbstractPojoWriter<BigIntVector> {
+
+    public LongWriter(String fieldName) {
+      super(fieldName, Types.required(MinorType.BIGINT));
+    }
+
+    @Override
+    public void writeField(Object value, int outboundIndex) {
+      vector.getMutator().setSafe(outboundIndex, (long) value);
+    }
+
+  }
+
+  /**
+   * Pojo writer for double. Does not expect to write null value.
+   */
+  public static class DoubleWriter extends AbstractPojoWriter<Float8Vector> {
+
+    public DoubleWriter(String fieldName) {
+      super(fieldName, Types.required(MinorType.FLOAT8));
+    }
+
+    @Override
+    public void writeField(Object value, int outboundIndex) {
+      vector.getMutator().setSafe(outboundIndex, (double) value);
+    }
+
+  }
+
+  /**
+   * Parent class for String and Enum writers. Writes data using nullable varchar holder.
+   */
+  private abstract static class AbstractStringWriter extends AbstractPojoWriter<NullableVarCharVector>
{
+    private DrillBuf data;
+    private final NullableVarCharHolder holder = new NullableVarCharHolder();
+
+    public AbstractStringWriter(String fieldName, DrillBuf managedBuf) {
+      super(fieldName, Types.optional(MinorType.VARCHAR));
+      this.data = managedBuf;
+      ensureLength(100);
+    }
+
+    void ensureLength(int len) {
+      data = data.reallocIfNeeded(len);
+    }
+
+    public void writeString(String s, int outboundIndex) {
+      holder.isSet = 1;
+      byte[] bytes = s.getBytes(Charsets.UTF_8);
+      ensureLength(bytes.length);
+      data.clear();
+      data.writeBytes(bytes);
+      holder.buffer = data;
+      holder.start = 0;
+      holder.end = bytes.length;
+      vector.getMutator().setSafe(outboundIndex, holder);
+    }
+  }
+
+  /**
+   * Pojo writer for Enum. If null is encountered does not write it.
+   */
+  public static class EnumWriter extends AbstractStringWriter{
+    public EnumWriter(String fieldName, DrillBuf managedBuf) {
+      super(fieldName, managedBuf);
+    }
+
+    @Override
+    public void writeField(Object value, int outboundIndex) {
+      if (value == null) {
+        return;
+      }
+      writeString(((Enum<?>) value).name(), outboundIndex);
+    }
+  }
+
+  /**
+   * Pojo writer for String. If null is encountered does not write it.
+   */
+  public static class StringWriter extends AbstractStringWriter {
+    public StringWriter(String fieldName, DrillBuf managedBuf) {
+      super(fieldName, managedBuf);
+    }
+
+    @Override
+    public void writeField(Object value, int outboundIndex) {
+      if (value != null) {
+        writeString((String) value, outboundIndex);
+      }
+    }
+  }
+
+  /**
+   * Pojo writer for Integer. If null is encountered does not write it.
+   */
+  public static class NIntWriter extends AbstractPojoWriter<NullableIntVector> {
+
+    public NIntWriter(String fieldName) {
+      super(fieldName, Types.optional(MinorType.INT));
+    }
+
+    @Override
+    public void writeField(Object value, int outboundIndex) {
+      if (value != null) {
+        vector.getMutator().setSafe(outboundIndex, (Integer) value);
+      }
+    }
+
+  }
+
+  /**
+   * Pojo writer for Long. If null is encountered does not write it.
+   */
+  public static class NBigIntWriter extends AbstractPojoWriter<NullableBigIntVector>
{
+
+    public NBigIntWriter(String fieldName) {
+      super(fieldName, Types.optional(MinorType.BIGINT));
+    }
+
+    @Override
+    public void writeField(Object value, int outboundIndex) {
+      if (value != null) {
+        vector.getMutator().setSafe(outboundIndex, (Long) value);
+      }
+    }
+
+  }
+
+  /**
+   * Pojo writer for Boolean. If null is encountered does not write it.
+   */
+  public static class NBooleanWriter extends AbstractPojoWriter<NullableBitVector>
{
+
+    public NBooleanWriter(String fieldName) {
+      super(fieldName, Types.optional(MinorType.BIT));
+    }
+
+    @Override
+    public void writeField(Object value, int outboundIndex) {
+      if (value != null) {
+        vector.getMutator().setSafe(outboundIndex, (Boolean) value ? 1 : 0);
+      }
+    }
+
+  }
+
+  /**
+   * Pojo writer for Double. If null is encountered does not write it.
+   */
+  public static class NDoubleWriter extends AbstractPojoWriter<NullableFloat8Vector>
{
+
+    public NDoubleWriter(String fieldName) {
+      super(fieldName, Types.optional(MinorType.FLOAT8));
+    }
+
+    @Override
+    public void writeField(Object value, int outboundIndex) {
+      if (value != null) {
+        vector.getMutator().setSafe(outboundIndex, (Double) value);
+      }
+    }
+
+  }
+
+  /**
+   * Pojo writer for Timestamp. If null is encountered does not write it.
+   */
+  public static class NTimeStampWriter extends AbstractPojoWriter<NullableTimeStampVector>
{
+
+    public NTimeStampWriter(String fieldName) {
+      super(fieldName, Types.optional(MinorType.TIMESTAMP));
+    }
+
+    @Override
+    public void writeField(Object value, int outboundIndex) {
+      if (value != null) {
+        vector.getMutator().setSafe(outboundIndex, ((Timestamp) value).getTime());
+      }
+    }
+  }
+}

http://git-wip-us.apache.org/repos/asf/drill/blob/8b564235/exec/java-exec/src/main/java/org/apache/drill/exec/store/pojo/Writers.java
----------------------------------------------------------------------
diff --git a/exec/java-exec/src/main/java/org/apache/drill/exec/store/pojo/Writers.java b/exec/java-exec/src/main/java/org/apache/drill/exec/store/pojo/Writers.java
deleted file mode 100644
index e52384e..0000000
--- a/exec/java-exec/src/main/java/org/apache/drill/exec/store/pojo/Writers.java
+++ /dev/null
@@ -1,274 +0,0 @@
-/**
- * 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.
- */
-package org.apache.drill.exec.store.pojo;
-
-import io.netty.buffer.DrillBuf;
-
-import java.lang.reflect.Field;
-import java.sql.Timestamp;
-
-import org.apache.drill.common.types.TypeProtos.MinorType;
-import org.apache.drill.common.types.Types;
-import org.apache.drill.exec.expr.holders.NullableVarCharHolder;
-import org.apache.drill.exec.vector.BigIntVector;
-import org.apache.drill.exec.vector.BitVector;
-import org.apache.drill.exec.vector.Float8Vector;
-import org.apache.drill.exec.vector.IntVector;
-import org.apache.drill.exec.vector.NullableBigIntVector;
-import org.apache.drill.exec.vector.NullableBitVector;
-import org.apache.drill.exec.vector.NullableFloat8Vector;
-import org.apache.drill.exec.vector.NullableIntVector;
-import org.apache.drill.exec.vector.NullableTimeStampVector;
-import org.apache.drill.exec.vector.NullableVarCharVector;
-
-import com.google.common.base.Charsets;
-
-public class Writers {
-  static final org.slf4j.Logger logger = org.slf4j.LoggerFactory.getLogger(Writers.class);
-
-  public static class IntWriter extends AbstractWriter<IntVector> {
-
-    public IntWriter(Field field) {
-      super(field, Types.required(MinorType.INT));
-      if (field.getType() != int.class) {
-        throw new IllegalStateException();
-      }
-    }
-
-    @Override
-    public void writeField(Object pojo, int outboundIndex) throws IllegalArgumentException,
IllegalAccessException {
-      int i = field.getInt(pojo);
-      vector.getMutator().setSafe(outboundIndex, i);
-    }
-
-  }
-
-  public static class BitWriter extends AbstractWriter<BitVector>{
-
-    public BitWriter(Field field) {
-      super(field, Types.required(MinorType.BIT));
-      if (field.getType() != boolean.class) {
-        throw new IllegalStateException();
-      }
-    }
-
-    @Override
-    public void writeField(Object pojo, int outboundIndex) throws IllegalArgumentException,
IllegalAccessException {
-      boolean b = field.getBoolean(pojo);
-      vector.getMutator().setSafe(outboundIndex, b ? 1 : 0);
-    }
-
-  }
-
-  public static class LongWriter extends AbstractWriter<BigIntVector>{
-
-    public LongWriter(Field field) {
-      super(field, Types.required(MinorType.BIGINT));
-      if (field.getType() != long.class) {
-        throw new IllegalStateException();
-      }
-    }
-
-    @Override
-    public void writeField(Object pojo, int outboundIndex) throws IllegalArgumentException,
IllegalAccessException {
-      long l = field.getLong(pojo);
-      vector.getMutator().setSafe(outboundIndex, l);
-    }
-
-  }
-
-  public static class DoubleWriter extends AbstractWriter<Float8Vector>{
-
-    public DoubleWriter(Field field) {
-      super(field, Types.required(MinorType.FLOAT8));
-      if (field.getType() != double.class) {
-        throw new IllegalStateException();
-      }
-    }
-
-    @Override
-    public void writeField(Object pojo, int outboundIndex) throws IllegalArgumentException,
IllegalAccessException {
-      double d = field.getDouble(pojo);
-
-      vector.getMutator().setSafe(outboundIndex, d);
-    }
-
-  }
-
-  private abstract static class AbstractStringWriter extends AbstractWriter<NullableVarCharVector>{
-    private DrillBuf data;
-    private final NullableVarCharHolder h = new NullableVarCharHolder();
-
-    public AbstractStringWriter(Field field, DrillBuf managedBuf) {
-      super(field, Types.optional(MinorType.VARCHAR));
-      this.data = managedBuf;
-      ensureLength(100);
-    }
-
-    void ensureLength(int len) {
-      data = data.reallocIfNeeded(len);
-    }
-
-    @Override
-    public void cleanup() {
-    }
-
-    public void writeString(String s, int outboundIndex) throws IllegalArgumentException,
IllegalAccessException {
-      if (s == null) {
-        return;
-      } else {
-        h.isSet = 1;
-        byte[] bytes = s.getBytes(Charsets.UTF_8);
-        ensureLength(bytes.length);
-        data.clear();
-        data.writeBytes(bytes);
-        h.buffer = data;
-        h.start = 0;
-        h.end = bytes.length;
-        vector.getMutator().setSafe(outboundIndex, h);
-      }
-    }
-
-  }
-
-  public static class EnumWriter extends AbstractStringWriter{
-    public EnumWriter(Field field, DrillBuf managedBuf) {
-      super(field, managedBuf);
-      if (!field.getType().isEnum()) {
-        throw new IllegalStateException();
-      }
-    }
-
-    @Override
-    public void writeField(Object pojo, int outboundIndex) throws IllegalArgumentException,
IllegalAccessException {
-      Enum<?> e= ((Enum<?>) field.get(pojo));
-      if (e == null) {
-        return;
-      }
-      writeString(e.name(), outboundIndex);
-    }
-  }
-
-  public static class StringWriter extends AbstractStringWriter {
-    public StringWriter(Field field, DrillBuf managedBuf) {
-      super(field, managedBuf);
-      if (field.getType() != String.class) {
-        throw new IllegalStateException();
-      }
-    }
-
-    @Override
-    public void writeField(Object pojo, int outboundIndex) throws IllegalArgumentException,
IllegalAccessException {
-      String s = (String) field.get(pojo);
-      writeString(s, outboundIndex);
-    }
-  }
-
-  public static class NIntWriter extends AbstractWriter<NullableIntVector>{
-
-    public NIntWriter(Field field) {
-      super(field, Types.optional(MinorType.INT));
-      if (field.getType() != Integer.class) {
-        throw new IllegalStateException();
-      }
-    }
-
-    @Override
-    public void writeField(Object pojo, int outboundIndex) throws IllegalArgumentException,
IllegalAccessException {
-      Integer i = (Integer) field.get(pojo);
-      if (i != null) {
-        vector.getMutator().setSafe(outboundIndex, i);
-      }
-    }
-
-  }
-
-  public static class NBigIntWriter extends AbstractWriter<NullableBigIntVector>{
-
-    public NBigIntWriter(Field field) {
-      super(field, Types.optional(MinorType.BIGINT));
-      if (field.getType() != Long.class) {
-        throw new IllegalStateException();
-      }
-    }
-
-    @Override
-    public void writeField(Object pojo, int outboundIndex) throws IllegalArgumentException,
IllegalAccessException {
-      Long o = (Long) field.get(pojo);
-      if (o != null) {
-        vector.getMutator().setSafe(outboundIndex, o);
-      }
-    }
-
-  }
-
-  public static class NBooleanWriter extends AbstractWriter<NullableBitVector>{
-
-    public NBooleanWriter(Field field) {
-      super(field, Types.optional(MinorType.BIT));
-      if (field.getType() != Boolean.class) {
-        throw new IllegalStateException();
-      }
-    }
-
-    @Override
-    public void writeField(Object pojo, int outboundIndex) throws IllegalArgumentException,
IllegalAccessException {
-      Boolean o = (Boolean) field.get(pojo);
-      if (o != null) {
-        vector.getMutator().setSafe(outboundIndex, o ? 1 : 0);
-      }
-    }
-
-  }
-  public static class NDoubleWriter extends AbstractWriter<NullableFloat8Vector>{
-
-    public NDoubleWriter(Field field) {
-      super(field, Types.optional(MinorType.FLOAT8));
-      if (field.getType() != Double.class) {
-        throw new IllegalStateException();
-      }
-    }
-
-    @Override
-    public void writeField(Object pojo, int outboundIndex) throws IllegalArgumentException,
IllegalAccessException {
-      Double o = (Double) field.get(pojo);
-      if (o != null) {
-        vector.getMutator().setSafe(outboundIndex, o);
-      }
-    }
-
-  }
-
-  public static class NTimeStampWriter extends AbstractWriter<NullableTimeStampVector>{
-
-    public NTimeStampWriter(Field field) {
-      super(field, Types.optional(MinorType.TIMESTAMP));
-      if (field.getType() != Timestamp.class) {
-        throw new IllegalStateException();
-      }
-    }
-
-    @Override
-    public void writeField(Object pojo, int outboundIndex) throws IllegalArgumentException,
IllegalAccessException {
-      Timestamp o = (Timestamp) field.get(pojo);
-      if (o != null) {
-        vector.getMutator().setSafe(outboundIndex, o.getTime());
-      }
-    }
-  }
-}

http://git-wip-us.apache.org/repos/asf/drill/blob/8b564235/exec/java-exec/src/main/java/org/apache/drill/exec/store/sys/SystemTableBatchCreator.java
----------------------------------------------------------------------
diff --git a/exec/java-exec/src/main/java/org/apache/drill/exec/store/sys/SystemTableBatchCreator.java
b/exec/java-exec/src/main/java/org/apache/drill/exec/store/sys/SystemTableBatchCreator.java
index 58bf433..2b0ef3f 100644
--- a/exec/java-exec/src/main/java/org/apache/drill/exec/store/sys/SystemTableBatchCreator.java
+++ b/exec/java-exec/src/main/java/org/apache/drill/exec/store/sys/SystemTableBatchCreator.java
@@ -1,4 +1,4 @@
-/**
+/*
  * 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
@@ -21,6 +21,7 @@ import java.util.Collections;
 import java.util.Iterator;
 import java.util.List;
 
+import com.google.common.collect.ImmutableList;
 import org.apache.drill.common.exceptions.ExecutionSetupException;
 import org.apache.drill.exec.ops.FragmentContext;
 import org.apache.drill.exec.physical.impl.BatchCreator;
@@ -35,7 +36,6 @@ import org.apache.drill.exec.store.pojo.PojoRecordReader;
  * Local system tables do not require a full-fledged query because these records are present
on every Drillbit.
  */
 public class SystemTableBatchCreator implements BatchCreator<SystemTableScan> {
-//  private static final org.slf4j.Logger logger = org.slf4j.LoggerFactory.getLogger(SystemTableBatchCreator.class);
 
   @SuppressWarnings({ "rawtypes", "unchecked" })
   @Override
@@ -44,7 +44,7 @@ public class SystemTableBatchCreator implements BatchCreator<SystemTableScan>
{
     throws ExecutionSetupException {
     final SystemTable table = scan.getTable();
     final Iterator<Object> iterator = table.getIterator(context);
-    final RecordReader reader = new PojoRecordReader(table.getPojoClass(), iterator);
+    final RecordReader reader = new PojoRecordReader(table.getPojoClass(), ImmutableList.copyOf(iterator));
 
     return new ScanBatch(scan, context, Collections.singleton(reader).iterator());
   }

http://git-wip-us.apache.org/repos/asf/drill/blob/8b564235/exec/java-exec/src/test/java/org/apache/drill/TestFunctionsWithTypeExpoQueries.java
----------------------------------------------------------------------
diff --git a/exec/java-exec/src/test/java/org/apache/drill/TestFunctionsWithTypeExpoQueries.java
b/exec/java-exec/src/test/java/org/apache/drill/TestFunctionsWithTypeExpoQueries.java
index 43b594b..46a4823 100644
--- a/exec/java-exec/src/test/java/org/apache/drill/TestFunctionsWithTypeExpoQueries.java
+++ b/exec/java-exec/src/test/java/org/apache/drill/TestFunctionsWithTypeExpoQueries.java
@@ -1,4 +1,4 @@
-/**
+/*
  * 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
@@ -256,7 +256,7 @@ public class TestFunctionsWithTypeExpoQueries extends BaseTestQuery {
         "where concat(a, 'asdf') = 'asdf'", root);
 
     // Validate the plan
-    final String[] expectedPlan = {"Scan.*a.parquet.*numFiles=1"};
+    final String[] expectedPlan = {"Scan.*a.parquet.*numFiles = 1"};
     final String[] excludedPlan = {"Filter"};
     PlanTestBase.testPlanMatchingPatterns(query, expectedPlan, excludedPlan);
 
@@ -265,7 +265,7 @@ public class TestFunctionsWithTypeExpoQueries extends BaseTestQuery {
         .sqlQuery(query)
         .ordered()
         .baselineColumns("col")
-        .baselineValues(1l)
+        .baselineValues(1L)
         .build()
         .run();
   }

http://git-wip-us.apache.org/repos/asf/drill/blob/8b564235/exec/java-exec/src/test/java/org/apache/drill/exec/planner/logical/TestConvertCountToDirectScan.java
----------------------------------------------------------------------
diff --git a/exec/java-exec/src/test/java/org/apache/drill/exec/planner/logical/TestConvertCountToDirectScan.java
b/exec/java-exec/src/test/java/org/apache/drill/exec/planner/logical/TestConvertCountToDirectScan.java
index 21b4c79..04fe913 100644
--- a/exec/java-exec/src/test/java/org/apache/drill/exec/planner/logical/TestConvertCountToDirectScan.java
+++ b/exec/java-exec/src/test/java/org/apache/drill/exec/planner/logical/TestConvertCountToDirectScan.java
@@ -1,4 +1,4 @@
-/**
+/*
  * 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
@@ -18,15 +18,16 @@
 package org.apache.drill.exec.planner.logical;
 
 import org.apache.drill.PlanTestBase;
+import org.apache.drill.exec.ExecConstants;
 import org.junit.Test;
 
 public class TestConvertCountToDirectScan extends PlanTestBase {
-  static final org.slf4j.Logger logger = org.slf4j.LoggerFactory.getLogger(TestConvertCountToDirectScan.class);
 
   @Test
   public void ensureCaseDoesntConvertToDirectScan() throws Exception {
     testPlanMatchingPatterns(
-        "select count(case when n_name = 'ALGERIA' and n_regionkey = 2 then n_nationkey else
null end) as cnt from dfs.`${WORKING_PATH}/src/test/resources/directcount.parquet`",
+        "select count(case when n_name = 'ALGERIA' and n_regionkey = 2 then n_nationkey else
null end) as cnt\n" +
+            "from dfs.`${WORKING_PATH}/src/test/resources/directcount.parquet`",
         new String[] { "CASE" },
         new String[]{});
   }
@@ -36,7 +37,7 @@ public class TestConvertCountToDirectScan extends PlanTestBase {
     final String sql = "select count(*) as cnt from cp.`tpch/nation.parquet`";
     testPlanMatchingPatterns(
         sql,
-        new String[] { "PojoRecordReader" },
+        new String[] { "DynamicPojoRecordReader" },
         new String[]{});
 
     testBuilder()
@@ -45,7 +46,6 @@ public class TestConvertCountToDirectScan extends PlanTestBase {
         .baselineColumns("cnt")
         .baselineValues(25L)
         .go();
-
   }
 
   @Test
@@ -53,7 +53,7 @@ public class TestConvertCountToDirectScan extends PlanTestBase {
     final String sql = "select count(100) as cnt from cp.`tpch/nation.parquet`";
     testPlanMatchingPatterns(
         sql,
-        new String[] { "PojoRecordReader" },
+        new String[] { "DynamicPojoRecordReader" },
         new String[]{});
 
     testBuilder()
@@ -62,7 +62,6 @@ public class TestConvertCountToDirectScan extends PlanTestBase {
         .baselineColumns("cnt")
         .baselineValues(25L)
         .go();
-
   }
 
   @Test
@@ -70,7 +69,39 @@ public class TestConvertCountToDirectScan extends PlanTestBase {
     final String sql = "select count(1 + 2) as cnt from cp.`tpch/nation.parquet`";
     testPlanMatchingPatterns(
         sql,
-        new String[] { "PojoRecordReader" },
+        new String[] { "DynamicPojoRecordReader" },
+        new String[]{});
+
+    testBuilder()
+        .sqlQuery(sql)
+        .unOrdered()
+        .baselineColumns("cnt")
+        .baselineValues(25L)
+        .go();
+  }
+
+  @Test
+  public void ensureDoesNotConvertForDirectoryColumns() throws Exception {
+    final String sql = "select count(dir0) as cnt from cp.`tpch/nation.parquet`";
+    testPlanMatchingPatterns(
+        sql,
+        new String[] { "ParquetGroupScan" },
+        new String[]{});
+
+    testBuilder()
+        .sqlQuery(sql)
+        .unOrdered()
+        .baselineColumns("cnt")
+        .baselineValues(0L)
+        .go();
+  }
+
+  @Test
+  public void ensureConvertForImplicitColumns() throws Exception {
+    final String sql = "select count(fqn) as cnt from cp.`tpch/nation.parquet`";
+    testPlanMatchingPatterns(
+        sql,
+        new String[] { "DynamicPojoRecordReader" },
         new String[]{});
 
     testBuilder()
@@ -79,7 +110,42 @@ public class TestConvertCountToDirectScan extends PlanTestBase {
         .baselineColumns("cnt")
         .baselineValues(25L)
         .go();
+  }
+
+  @Test
+  public void ensureConvertForSeveralColumns() throws Exception {
+    test("use %s", TEMP_SCHEMA);
+    final String tableName = "parquet_table_counts";
+
+    try {
+      final String newFqnColumnName = "new_fqn";
+      test("alter session set `%s` = '%s'", ExecConstants.IMPLICIT_FQN_COLUMN_LABEL, newFqnColumnName);
+      test("create table %s as select * from cp.`parquet/alltypes_optional.parquet`", tableName);
+      test("refresh table metadata %s", tableName);
+
+      final String sql = String.format("select\n" +
+          "count(%s) as implicit_count,\n" +
+          "count(*) as star_count,\n" +
+          "count(col_int) as int_column_count,\n" +
+          "count(col_vrchr) as vrchr_column_count\n" +
+          "from %s", newFqnColumnName, tableName);
+
+      testPlanMatchingPatterns(
+          sql,
+          new String[] { "DynamicPojoRecordReader" },
+          new String[]{});
+
+      testBuilder()
+          .sqlQuery(sql)
+          .unOrdered()
+          .baselineColumns("implicit_count", "star_count", "int_column_count", "vrchr_column_count")
+          .baselineValues(6L, 6L, 2L, 3L)
+          .go();
 
+    } finally {
+      test("alter session reset `%s`", ExecConstants.IMPLICIT_FQN_COLUMN_LABEL);
+      test("drop table if exists %s", tableName);
+    }
   }
 
 }


Mime
View raw message