Add Expression(schema) method and improve adq print method

r/R/dplyr.R

@@ -70,17 +70,20 @@ make_field_refs <- function(field_names) {
 #' @export
 print.arrow_dplyr_query <- function(x, ...) {
   schm <- x$.data$schema
-  fields <- map_chr(x$selected_columns, function(expr) {
+  types <- map_chr(x$selected_columns, function(expr) {
     name <- expr$field_name
     if (nzchar(name)) {
-      schm$GetFieldByName(name)$ToString()
+      # Just a field_ref, so look up in the schema
+      schm$GetFieldByName(name)$type$ToString()
     } else {
-      # It's "" because this is not a field_ref, it's a more complex expression
-      "expr"
+      # Expression, so get its type and append the expression
+      paste0(
+        expr$type(schm)$ToString(),
+        " (", expr$ToString(), ")"
+      )
     }
   })
-  # Strip off the field names as they are in the dataset and add the renamed ones
-  fields <- paste(names(fields), sub("^.*?: ", "", fields), sep = ": ", collapse = "\n")
+  fields <- paste(names(types), types, sep = ": ", collapse = "\n")
   cat(class(x$.data)[1], " (query)\n", sep = "")
   cat(fields, "\n", sep = "")
   cat("\n")

r/R/expression.R

@@ -250,6 +250,7 @@ print.array_expression <- function(x, ...) {
 Expression <- R6Class("Expression", inherit = ArrowObject,
   public = list(
     ToString = function() compute___expr__ToString(self),
+    type = function(schema) compute___expr__type(self, schema),
     cast = function(to_type, safe = TRUE, ...) {
       opts <- list(
         to_type = to_type,

r/src/expression.cpp

@@ -68,4 +68,12 @@ std::string compute___expr__ToString(const std::shared_ptr<compute::Expression>&
   return x->ToString();
 }
 
+// [[arrow::export]]
+std::shared_ptr<arrow::DataType> compute___expr__type(
+    const std::shared_ptr<compute::Expression>& x,
+    const std::shared_ptr<arrow::Schema>& schema) {
+  auto bound = ValueOrStop(x->Bind(*schema));
+  return bound.type();
+}
+
 #endif

r/tests/testthat/test-dataset.R

@@ -965,7 +965,7 @@ test_that("mutate()", {
 chr: string
 dbl: double
 int: int32
-twice: expr
+twice: double (multiply_checked(int, 2))
 
 * Filter: ((multiply_checked(dbl, 2) > 14) and (subtract_checked(dbl, 50) < 3))
 See $.data for the source Arrow object",
@@ -1120,7 +1120,7 @@ test_that("arrange()", {
 chr: string
 dbl: double
 int: int32
-twice: expr
+twice: double (multiply_checked(int, 2))
 
 * Filter: ((multiply_checked(dbl, 2) > 14) and (subtract_checked(dbl, 50) < 3))
 * Sorted by chr [asc], multiply_checked(int, 2) [desc], add_checked(dbl, int) [asc]

r/tests/testthat/test-dplyr-mutate.R

@@ -346,10 +346,11 @@ test_that("print a mutated table", {
       print(),
 'InMemoryDataset (query)
 int: int32
-twice: expr
+twice: double (multiply_checked(int, 2))
 
 See $.data for the source Arrow object',
-  fixed = TRUE)
+    fixed = TRUE
+  )
 
   # Handling non-expressions/edge cases
   skip("InMemoryDataset$Project() doesn't accept array (or could it?)")

r/tests/testthat/test-expression.R

@@ -76,6 +76,14 @@ test_that("C++ expressions", {
     'Expression\n(f > 4)',
     fixed = TRUE
   )
+  expect_type_equal(
+    f$type(schema(f = float64())),
+    float64()
+  )
+  expect_type_equal(
+    (f > 4)$type(schema(f = float64())),
+    bool()
+  )
   # Interprets that as a list type
   expect_r6_class(f == c(1L, 2L), "Expression")
 })