Numeric array support for columnar frames (#15917)

Columnar frames used in subquery materialization and window functions now support numeric arrays.

processing/src/main/java/org/apache/druid/frame/read/columnar/DoubleArrayFrameColumnReader.java

@@ -0,0 +1,64 @@
+/*
+ * 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.druid.frame.read.columnar;
+
+import com.google.common.math.LongMath;
+import org.apache.datasketches.memory.Memory;
+import org.apache.druid.frame.Frame;
+import org.apache.druid.frame.write.columnar.FrameColumnWriters;
+import org.apache.druid.segment.column.ColumnType;
+
+/**
+ * Reader for columns written by {@link org.apache.druid.frame.write.columnar.DoubleArrayFrameColumnWriter}
+ *
+ * @see NumericArrayFrameColumnReader
+ * @see org.apache.druid.frame.write.columnar.NumericArrayFrameColumnWriter for column's layout in memory
+ */
+public class DoubleArrayFrameColumnReader extends NumericArrayFrameColumnReader
+{
+  public DoubleArrayFrameColumnReader(int columnNumber)
+  {
+    super(FrameColumnWriters.TYPE_DOUBLE_ARRAY, ColumnType.DOUBLE_ARRAY, columnNumber);
+  }
+
+  @Override
+  NumericArrayFrameColumn column(Frame frame, Memory memory, ColumnType columnType)
+  {
+    return new DoubleArrayFrameColumn(frame, memory, columnType);
+  }
+
+  private static class DoubleArrayFrameColumn extends NumericArrayFrameColumn
+  {
+    public DoubleArrayFrameColumn(
+        Frame frame,
+        Memory memory,
+        ColumnType columnType
+    )
+    {
+      super(frame, memory, columnType);
+    }
+
+    @Override
+    Number getElement(Memory memory, long rowDataOffset, int cumulativeIndex)
+    {
+      return memory.getDouble(LongMath.checkedAdd(rowDataOffset, (long) cumulativeIndex * Double.BYTES));
+    }
+  }
+}

processing/src/main/java/org/apache/druid/frame/read/columnar/FloatArrayFrameColumnReader.java

@@ -0,0 +1,64 @@
+/*
+ * 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.druid.frame.read.columnar;
+
+import com.google.common.math.LongMath;
+import org.apache.datasketches.memory.Memory;
+import org.apache.druid.frame.Frame;
+import org.apache.druid.frame.write.columnar.FrameColumnWriters;
+import org.apache.druid.segment.column.ColumnType;
+
+/**
+ * Reader for columns written by {@link org.apache.druid.frame.write.columnar.FloatArrayFrameColumnWriter}
+ *
+ * @see NumericArrayFrameColumnReader
+ * @see org.apache.druid.frame.write.columnar.NumericArrayFrameColumnWriter for column's layout in memory
+ */
+public class FloatArrayFrameColumnReader extends NumericArrayFrameColumnReader
+{
+  public FloatArrayFrameColumnReader(int columnNumber)
+  {
+    super(FrameColumnWriters.TYPE_FLOAT_ARRAY, ColumnType.FLOAT_ARRAY, columnNumber);
+  }
+
+  @Override
+  NumericArrayFrameColumn column(Frame frame, Memory memory, ColumnType columnType)
+  {
+    return new FloatArrayFrameColumn(frame, memory, columnType);
+  }
+
+  private static class FloatArrayFrameColumn extends NumericArrayFrameColumn
+  {
+    public FloatArrayFrameColumn(
+        Frame frame,
+        Memory memory,
+        ColumnType columnType
+    )
+    {
+      super(frame, memory, columnType);
+    }
+
+    @Override
+    Number getElement(Memory memory, long rowDataOffset, int cumulativeIndex)
+    {
+      return memory.getFloat(LongMath.checkedAdd(rowDataOffset, (long) cumulativeIndex * Float.BYTES));
+    }
+  }
+}

processing/src/main/java/org/apache/druid/frame/read/columnar/FrameColumnReaderUtils.java

@@ -0,0 +1,66 @@
+/*
+ * 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.druid.frame.read.columnar;
+
+import org.apache.datasketches.memory.Memory;
+
+public class FrameColumnReaderUtils
+{
+  /**
+   * Adjusts a negative cumulative row length from {@link #getCumulativeRowLength(Memory, long, int)} to be the actual
+   * positive length.
+   */
+  public static int adjustCumulativeRowLength(final int cumulativeRowLength)
+  {
+    if (cumulativeRowLength < 0) {
+      return -(cumulativeRowLength + 1);
+    } else {
+      return cumulativeRowLength;
+    }
+  }
+
+  /**
+   * Returns cumulative row length, if the row is not null itself, or -(cumulative row length) - 1 if the row is
+   * null itself.
+   *
+   * To check if the return value from this function indicate a null row, use {@link #isNullRow(int)}
+   *
+   * To get the actual cumulative row length, use {@link FrameColumnReaderUtils#adjustCumulativeRowLength(int)}.
+   */
+  public static int getCumulativeRowLength(final Memory memory, final long offset, final int physicalRow)
+  {
+    // Note: only valid to call this if multiValue = true.
+    return memory.getInt(offset + (long) Integer.BYTES * physicalRow);
+  }
+
+  public static int getAdjustedCumulativeRowLength(final Memory memory, final long offset, final int physicalRow)
+  {
+    return adjustCumulativeRowLength(getCumulativeRowLength(memory, offset, physicalRow));
+  }
+
+  /**
+   * When given a return value from {@link FrameColumnReaderUtils#getCumulativeRowLength(Memory, long, int)}, returns whether the row is
+   * null itself (i.e. a null array).
+   */
+  public static boolean isNullRow(final int cumulativeRowLength)
+  {
+    return cumulativeRowLength < 0;
+  }
+}

processing/src/main/java/org/apache/druid/frame/read/columnar/FrameColumnReaders.java

@@ -20,7 +20,6 @@
 package org.apache.druid.frame.read.columnar;
 
 import org.apache.druid.segment.column.ColumnType;
-import org.apache.druid.segment.column.ValueType;
 
 /**
  * Creates {@link FrameColumnReader}  corresponding to a given column type and number.
@@ -58,12 +57,18 @@ public static FrameColumnReader create(
         return new ComplexFrameColumnReader(columnNumber);
 
       case ARRAY:
-        if (columnType.getElementType().getType() == ValueType.STRING) {
-          return new StringFrameColumnReader(columnNumber, true);
-        } else {
-          return new UnsupportedColumnTypeFrameColumnReader(columnName, columnType);
+        switch (columnType.getElementType().getType()) {
+          case STRING:
+            return new StringFrameColumnReader(columnNumber, true);
+          case LONG:
+            return new LongArrayFrameColumnReader(columnNumber);
+          case FLOAT:
+            return new FloatArrayFrameColumnReader(columnNumber);
+          case DOUBLE:
+            return new DoubleArrayFrameColumnReader(columnNumber);
+          default:
+            return new UnsupportedColumnTypeFrameColumnReader(columnName, columnType);
         }
-
       default:
         return new UnsupportedColumnTypeFrameColumnReader(columnName, columnType);
     }

processing/src/main/java/org/apache/druid/frame/read/columnar/LongArrayFrameColumnReader.java

@@ -0,0 +1,64 @@
+/*
+ * 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.druid.frame.read.columnar;
+
+import com.google.common.math.LongMath;
+import org.apache.datasketches.memory.Memory;
+import org.apache.druid.frame.Frame;
+import org.apache.druid.frame.write.columnar.FrameColumnWriters;
+import org.apache.druid.segment.column.ColumnType;
+
+/**
+ * Reader for columns written by {@link org.apache.druid.frame.write.columnar.LongArrayFrameColumnWriter}
+ *
+ * @see NumericArrayFrameColumnReader
+ * @see org.apache.druid.frame.write.columnar.NumericArrayFrameColumnWriter for column's layout in memory
+ */
+public class LongArrayFrameColumnReader extends NumericArrayFrameColumnReader
+{
+  public LongArrayFrameColumnReader(int columnNumber)
+  {
+    super(FrameColumnWriters.TYPE_LONG_ARRAY, ColumnType.LONG_ARRAY, columnNumber);
+  }
+
+  @Override
+  NumericArrayFrameColumn column(Frame frame, Memory memory, ColumnType columnType)
+  {
+    return new LongArrayFrameColumn(frame, memory, columnType);
+  }
+
+  private static class LongArrayFrameColumn extends NumericArrayFrameColumn
+  {
+    public LongArrayFrameColumn(
+        Frame frame,
+        Memory memory,
+        ColumnType columnType
+    )
+    {
+      super(frame, memory, columnType);
+    }
+
+    @Override
+    Number getElement(Memory memory, long rowDataOffset, int cumulativeIndex)
+    {
+      return memory.getLong(LongMath.checkedAdd(rowDataOffset, (long) cumulativeIndex * Long.BYTES));
+    }
+  }
+}