AssociationMap to query Dataset across levels?

Question

Is there a way to operate on Dataset across levels to apply a function f to a key -> value pair in an Association as follows:

pivotApply[f_, key_ -> value_] :=  key -> f[key, value]

For example, take the Titanic dataset and pivot up "class":

titanicClass = titanic[GroupBy[Key@"class"], KeyDrop[Key@"class"]];

For this problem, assume only titanicClass. "class" key has been deleted and so can't referenced. The motivation is tree-structured data where it would be inefficient to insert the keys {"1st"...} at the deeper level.

Here foo represents a "client" function to be applied to deeper level data but depenent on class values.

EDIT // #age >= 18 &

foo["1st", data_] := Dataset[data][All, #age >= 18 &];
foo["2nd", data_] := Dataset[data][All, "sex"];
foo["3rd", data_] := Dataset[data][All, "survived"];

Is there a way to avoid inefficient downcasting to Normal and back to Dataset to use AssociationMap? Also note there are Datasets nested in a larger one, which should also be cast to Lists (since there is an outer Dataset).

AssociationMap[pivotApply[foo, #] &, titanicClass // Normal] // Dataset

Is there a way to query Datasets across levels like this? This seems extremely roundabout, but haven't figured an implementation using, say Keys and Values and related Dataset operators.

Answer 1

Note: the question has evolved, as has my understanding of it. My response has been changed completely, although I have kept my earlier responses below.

If we can drop the requirement to use AssociationMap, pivotApply and foo, then we can generate the desired output directly using subqueries:

titanicClass @
  { "1st" -> Query[All, #age > 18&]
  , "2nd" -> Query[All, "sex"]
  , "3rd" -> Query[All, "survived"]
  }

As requested, the inner lists are not wrapped in Dataset.

---

Original Responses

If we are willing to dispense with the requirement to use AssociationMap, here are a couple of options to "query across levels" in the described manner.

If we introduce bar as a slightly rewritten form of the function foo from the question...

bar[a_] := bar[a, a@"class"]
bar[a_, "1st"] := a@"age"
bar[a_, "2nd"] := a@"sex"
bar[a_, "3rd"] := a@"survived"

... then the function can be called directly within the dataset query after grouping:

titanic[GroupBy @ Key @ "class", All, bar]

Alternatively, we could dispense with the function altogether and extract the various components with an inline Switch statement.

titanic[
  GroupBy @ Key @ "class"
, All
, Switch[#class, "1st", #age, "2nd", #sex, "3rd", #survived]&
]

---

Edit: Starting From titanicClass

My original response mistakenly assumed that titanic was the input data. The question has been revised to emphasize that titanicClass is actually the input data. Under that assumption, we can generate the described output as follows...

First, we supplement the definition of pivotApply to include a one-argument "lifting" form of the operator:

pivotApply[f_] := pivotApply[f, #]&

Next, we create the function baz which is yet another variation upon foo:

baz["1st", a_] := a[[All, "age"]]
baz["2nd", a_] := a[[All, "sex"]]
baz["3rd", a_] := a[[All, "survived"]]

Now, we can generate the required output using AssociationMap:

titanicClass[AssociationMap @ pivotApply @ baz]

Once again, we can dispense with the function baz and express the query inline if desired:

titanicClass @ AssociationMap @ pivotApply[
  #2[[All, # /. {"1st" -> "age", "2nd" -> "sex", "3rd" -> "survived"}]]&
]

Answer 2

Using this helper function

ClearAll[keyStrip]
keyStrip[Key[k_]|{Key[k_]}] := k;
keyStrip[expr_] := expr;

you can achieve what you asked for as

titanicClass[MapIndexed[foo[keyStrip[First@#2], #1] &]]

The general issue about being able to reference some elements of the parent structures inside child structures during the traversal (in the queries) is more complicated. As of now, Dataset query language does not, as far as I know, easily allow this, but such capabilities may be added in the future.

EDIT

Here is a version which does not use the foo function, but uses Lookup and specifies the lookup explicitly via an association:

titanicClass[
  MapIndexed[
    Lookup[
      #1, 
      <|"1st" -> "age", "2nd" -> "sex", "3rd" -> "survived"|>[keyStrip[#2]]
    ] &
  ]
]

Answer 3

titanic[GroupBy[Key@"class" -> (Switch[#class,
    "1st", #age,
    "2nd", #sex,
    "3rd", #survived] &)
]] // Normal // Dataset