Skip to content

Commit

Permalink
feat: Implement compilation for column expressions
Browse files Browse the repository at this point in the history
  • Loading branch information
@clflushopt
clflushopt committed Feb 3, 2025
1 parent 4728de8 commit e1f9d62
Show file tree
Hide file tree
Showing 3 changed files with 110 additions and 14 deletions.
104 changes: 91 additions & 13 deletions glint/src/main/java/co/clflushopt/glint/query/compiler/QueryCompiler.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,8 +8,11 @@
import org.codehaus.commons.compiler.ISimpleCompiler;
import org.codehaus.janino.CompilerFactory;

import co.clflushopt.glint.query.logical.expr.LogicalColumnExpr;
import co.clflushopt.glint.query.logical.plan.LogicalPlan;
import co.clflushopt.glint.query.logical.plan.Scan;
import co.clflushopt.glint.types.ArrowTypes;
import co.clflushopt.glint.types.Field;

/**
* Core Query Compiler using Janino for runtime code generation. This class
Expand Down Expand Up @@ -113,30 +116,48 @@ public Object compile(LogicalPlan logicalPlan) throws Exception {
}

/**
* Generates source code for a given logical plan.
* Compile a logical plan to Java code.
*
* @param logicalPlan The logical plan to generate code for
* @return Generated Java source code as a string
*/
private String generateSourceCode(LogicalPlan logicalPlan) {
return String.format(
"""
package co.clflushopt.glint.generated;
return String.format("""
package co.clflushopt.glint.generated;
import co.clflushopt.glint.query.logical.plan.LogicalPlan;
import co.clflushopt.glint.query.logical.plan.Scan;
import co.clflushopt.glint.query.logical.plan.Filter;
import co.clflushopt.glint.datasource.DataSource;
import co.clflushopt.glint.types.RecordBatch;
import co.clflushopt.glint.types.Schema;
import co.clflushopt.glint.types.RecordBatch;
import co.clflushopt.glint.types.ArrowTypes;
import java.util.List;
import co.clflushopt.glint.query.logical.plan.LogicalPlan;
import co.clflushopt.glint.query.logical.plan.Scan;
import co.clflushopt.glint.datasource.DataSource;
import co.clflushopt.glint.types.RecordBatch;
import co.clflushopt.glint.types.Schema;
import java.util.List;
%s
""", compileScanPlan(logicalPlan));
}

/**
* Compile scan operator.
*
*
* @param expr
* @param plan
* @return
*/
private String compileScanPlan(LogicalPlan plan) {
return String.format(
"""
public class LogicalPlanExecutor {
public long execute() {
return executePlan(%s);
}
private long executePlan(LogicalPlan plan) {
// Handle different logical plan types
// TODO: Handle different logical plan types
if (plan instanceof Scan) {
Scan scan = (Scan) plan;
return executeScan(scan);
Expand All @@ -154,14 +175,71 @@ private long executeScan(Scan scan) {
for (RecordBatch batch : dataSource.scan(projections)) {
totalRecords += batch.getRowCount();
System.out.println("Scan batch size: " + batch.getRowCount() +
" Path: " + scan.getPath());
" Path: " + scan.getPath());
}
return totalRecords;
}
}""",
generatePlanArgument(plan));

}

/**
* Generates source code specifically for column expressions.
*
* @param columnExpr The column expression to generate code for
* @param plan The logical plan context
* @return Generated Java source code for the column expression
*/
private String compileColumnExpr(LogicalColumnExpr columnExpr, LogicalPlan plan) {
// Retrieve the field information for the column
Field field = columnExpr.toField(plan);

return String.format(
"""
public class ExpressionExecutor {
// Add method to find column index
private int findColumnIndex(RecordBatch batch) {
Schema schema = batch.getSchema();
return IntStream.range(0, schema.getFields().size())
.filter(idx -> schema.getFields().get(idx).name().equals(columnExpr.getName()))
.findFirst()
.orElseThrow(() -> new IllegalArgumentException("Column '" + columnExpr.getName() + "' not found in schema"));
}
/**
* Extracts the value of the column '%s' from a record batch.
*
* @param batch The record batch to extract the value from
* @param rowIndex The index of the row to extract
* @return The value of the column
*/
public Object getValue(RecordBatch batch, int rowIndex) {
int columnIndex = findColumnIndex(batch);
return batch.getField(rowIndex);
}
/**
* Gets the field metadata for the column.
*
* @return Field metadata
*/
public Field getFieldMetadata() {
return new Field("%s", %s);
}
@Override
public String toString() {
return "%s";
}
}
""",
generatePlanArgument(logicalPlan));
columnExpr.getName(), // Column name for comments
columnExpr.getName(), // Method to get column value
columnExpr.getName(), // Field name
ArrowTypes.compile(field.dataType()), // Field type
columnExpr.toString() // String representation
);
}

/**
Expand Down
18 changes: 18 additions & 0 deletions glint/src/main/java/co/clflushopt/glint/types/ArrowTypes.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -50,4 +50,22 @@ public static String toString(ArrowType dataType) {
return dataType.accept(new PrettyArrowTypeVisitor());
}

/**
* Compiled representation of a reference to a wrapped arrow type.
*
*/
public static String compile(ArrowType type) {
switch (type.getTypeID()) {
case Bool:
return "ArrowTypes.BooleanType";
case Int:
return "ArrowTypes.Int64Type";
case FloatingPoint:
return "ArrowTypes.DoubleType";
default:
break;
}
return "ArrowTypes.StringType";
}

}
2 changes: 1 addition & 1 deletion glint/src/main/java/co/clflushopt/glint/types/Field.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@ public record Field(String name, ArrowType dataType) {

/**
* Transform an internal `Field` type to Arrow Field type.
*
*
* @return
*/
public org.apache.arrow.vector.types.pojo.Field toArrow() {
Expand Down

0 comments on commit e1f9d62

Please sign in to comment.