diff --git a/native/src/exchange.rs b/native/src/exchange.rs index b48978b1f..531a66bad 100644 --- a/native/src/exchange.rs +++ b/native/src/exchange.rs @@ -35,7 +35,13 @@ pub(crate) fn partition_batch( .collect(); out.push(( key_group, - RecordBatch::try_new(batch.schema(), columns).expect("sub batch"), + // A key-only projection carries no columns, and Arrow cannot infer a row count from none. + RecordBatch::try_new_with_options( + batch.schema(), + columns, + &arrow::record_batch::RecordBatchOptions::new().with_row_count(Some(indices.len())), + ) + .expect("sub batch"), )); } out @@ -74,7 +80,13 @@ pub(crate) fn partition_batch_by_key_group( Some(( key_group, ordinals, - RecordBatch::try_new(batch.schema(), columns).expect("key-group sub-batch"), + RecordBatch::try_new_with_options( + batch.schema(), + columns, + &arrow::record_batch::RecordBatchOptions::new() + .with_row_count(Some(indices.len())), + ) + .expect("key-group sub-batch"), )) }) .collect() @@ -245,3 +257,34 @@ pub extern "system" fn Java_tech_streamfusion_Native_concatBatches<'local>( export_record_batch(merged, out_array_address, out_schema_address); }) } + +#[cfg(test)] +mod tests { + use super::*; + use arrow::datatypes::Schema; + use std::sync::Arc; + + fn zero_column_batch(rows: usize) -> RecordBatch { + RecordBatch::try_new_with_options( + Arc::new(Schema::empty()), + vec![], + &arrow::record_batch::RecordBatchOptions::new().with_row_count(Some(rows)), + ) + .expect("zero-column batch") + } + + /// A projection can leave no columns behind, and Arrow refuses to infer a row count from none. + #[test] + fn partitions_a_zero_column_batch() { + let batch = zero_column_batch(4); + let parts = partition_batch(&batch, &[], &[], 128, 2); + assert_eq!(parts.iter().map(|(_, b)| b.num_rows()).sum::(), 4); + } + + #[test] + fn partitions_a_zero_column_batch_by_key_group() { + let batch = zero_column_batch(3); + let parts = partition_batch_by_key_group(&batch, &[], &[], 128); + assert_eq!(parts.iter().map(|(_, _, b)| b.num_rows()).sum::(), 3); + } +} diff --git a/native/src/expr.rs b/native/src/expr.rs index a2017af4a..1c194c14a 100644 --- a/native/src/expr.rs +++ b/native/src/expr.rs @@ -33,6 +33,8 @@ pub(crate) fn build_expr( 7 => logical_lit(longs[arg] as i32), 8 => logical_lit(longs[arg] as i16), 9 => logical_lit(longs[arg] as i8), + // Likewise for FLOAT, which the host evaluates in single precision. + 22 => logical_lit(doubles[arg] as f32), 11 => { // A widening numeric cast: build the single child, then wrap it. `arg` is the target code. let child = build_expr( diff --git a/src/main/java/tech/streamfusion/planner/RexExpression.java b/src/main/java/tech/streamfusion/planner/RexExpression.java index 071e6858f..773ccaf0c 100644 --- a/src/main/java/tech/streamfusion/planner/RexExpression.java +++ b/src/main/java/tech/streamfusion/planner/RexExpression.java @@ -88,6 +88,9 @@ final class RexExpression { private static final int KIND_DECIMAL_DIVIDE = 20; private static final int KIND_DECIMAL_MOD = 21; + // A single-precision literal, so a FLOAT expression is not widened to double by its constants. + private static final int KIND_LIT_FLOAT = 22; + // Cast target type codes, mirrored on the native side. private static final int CAST_TINYINT = 0; private static final int CAST_SMALLINT = 1; @@ -406,6 +409,17 @@ private boolean emitLiteral(RexLiteral literal) { } case FLOAT: case REAL: + { + Double value = literal.getValueAs(Double.class); + if (value == null) { + return false; + } + // The host evaluates FLOAT arithmetic in single precision. A double literal would widen the + // whole expression, changing both the result type the plan promised and its last bits. + add(KIND_LIT_FLOAT, doubles.size(), 0); + doubles.add(value); + return true; + } case DOUBLE: { Double value = literal.getValueAs(Double.class); diff --git a/src/test/java/tech/streamfusion/FlinkCalcSqlHarnessTest.java b/src/test/java/tech/streamfusion/FlinkCalcSqlHarnessTest.java index d83c6f936..0fdae4efc 100644 --- a/src/test/java/tech/streamfusion/FlinkCalcSqlHarnessTest.java +++ b/src/test/java/tech/streamfusion/FlinkCalcSqlHarnessTest.java @@ -28,6 +28,25 @@ void multiplyProjectionMatchesHost() throws Exception { NativeParity.assertParity(FlinkCalcSqlHarnessTest::environment, "SELECT v * 2 FROM f"); } + @Test + void floatLiteralArithmeticMatchesHost() throws Exception { + // A double-typed constant would widen the whole expression, so the projection would come back + // as DOUBLE where the plan promised FLOAT. + NativeParity.assertParity( + FlinkCalcSqlHarnessTest::environment, "SELECT f4 * CAST(2 AS FLOAT) FROM f"); + } + + @Test + void floatColumnArithmeticMatchesHost() throws Exception { + NativeParity.assertParity(FlinkCalcSqlHarnessTest::environment, "SELECT f4 + f4, f4 - f4 FROM f"); + } + + @Test + void floatCastProjectionMatchesHost() throws Exception { + NativeParity.assertParity( + FlinkCalcSqlHarnessTest::environment, "SELECT CAST(v AS FLOAT) * f4 FROM f"); + } + @Test void mixedComputedAndColumnProjectionMatchesHost() throws Exception { NativeParity.assertParity(FlinkCalcSqlHarnessTest::environment, "SELECT v + k, s, k FROM f"); @@ -474,11 +493,16 @@ private static TableEnvironment environment() { StreamTableEnvironment tEnv = StreamTableEnvironment.create(env); DataStream source = env.fromData( - Types.ROW_NAMED(new String[] {"k", "v", "s"}, Types.LONG, Types.INT, Types.STRING), - Row.of(1L, 10, "a"), - Row.of(2L, 30, "b"), - Row.of(3L, 20, "c"), - Row.of(4L, 40, "d")); + Types.ROW_NAMED( + new String[] {"k", "v", "s", "f4"}, + Types.LONG, + Types.INT, + Types.STRING, + Types.FLOAT), + Row.of(1L, 10, "a", 1.5f), + Row.of(2L, 30, "b", 2.25f), + Row.of(3L, 20, "c", 3.75f), + Row.of(4L, 40, "d", 4.125f)); tEnv.createTemporaryView( "f", source, @@ -486,6 +510,7 @@ private static TableEnvironment environment() { .column("k", DataTypes.BIGINT()) .column("v", DataTypes.INT()) .column("s", DataTypes.STRING()) + .column("f4", DataTypes.FLOAT()) .build()); return tEnv; }