cpp/src/arrow/dataset/partition.cc - arrow - Git at Google

 // 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.

 #include "arrow/dataset/partition.h"

 #include <algorithm>
 #include <chrono>
 #include <map>
 #include <memory>
 #include <stack>
 #include <utility>
 #include <vector>

 #include "arrow/dataset/dataset_internal.h"
 #include "arrow/dataset/file_base.h"
 #include "arrow/dataset/filter.h"
 #include "arrow/dataset/scanner.h"
 #include "arrow/filesystem/filesystem.h"
 #include "arrow/filesystem/path_util.h"
 #include "arrow/scalar.h"
 #include "arrow/util/iterator.h"
 #include "arrow/util/range.h"
 #include "arrow/util/sort.h"
 #include "arrow/util/string_view.h"

 namespace arrow {
 namespace dataset {

 using util::string_view;

 using internal::checked_cast;

 Result<std::shared_ptr<Expression>> Partitioning::Parse(const std::string& path) const {
   ExpressionVector expressions;
   int i = 0;

   for (auto segment : fs::internal::SplitAbstractPath(path)) {
     ARROW_ASSIGN_OR_RAISE(auto expr, Parse(segment, i++));
     if (expr->Equals(true)) {
       continue;
     }

     expressions.push_back(std::move(expr));
   }

   return and_(std::move(expressions));
 }

 std::shared_ptr<Partitioning> Partitioning::Default() {
   return std::make_shared<DefaultPartitioning>();
 }

 Result<WritePlan> PartitioningFactory::MakeWritePlan(FragmentIterator fragment_it) {
   return Status::NotImplemented("MakeWritePlan from PartitioningFactory of type ",
                                 type_name());
 }

 Result<WritePlan> PartitioningFactory::MakeWritePlan(
     FragmentIterator fragment_it, const std::shared_ptr<Schema>& schema) {
   return Status::NotImplemented("MakeWritePlan from PartitioningFactory of type ",
                                 type_name());
 }

 Result<std::shared_ptr<Expression>> SegmentDictionaryPartitioning::Parse(
     const std::string& segment, int i) const {
   if (static_cast<size_t>(i) < dictionaries_.size()) {
     auto it = dictionaries_[i].find(segment);
     if (it != dictionaries_[i].end()) {
       return it->second;
     }
   }

   return scalar(true);
 }

 Status KeyValuePartitioning::VisitKeys(
     const Expression& expr,
     const std::function<Status(const std::string& name,
                                const std::shared_ptr<Scalar>& value)>& visitor) {
   if (expr.type() == ExpressionType::AND) {
     const auto& and_ = checked_cast<const AndExpression&>(expr);
     RETURN_NOT_OK(VisitKeys(*and_.left_operand(), visitor));
     RETURN_NOT_OK(VisitKeys(*and_.right_operand(), visitor));
     return Status::OK();
   }

   if (expr.type() != ExpressionType::COMPARISON) {
     return Status::OK();
   }

   const auto& cmp = checked_cast<const ComparisonExpression&>(expr);
   if (cmp.op() != compute::CompareOperator::EQUAL) {
     return Status::OK();
   }

   auto lhs = cmp.left_operand().get();
   auto rhs = cmp.right_operand().get();
   if (lhs->type() != ExpressionType::FIELD) std::swap(lhs, rhs);

   if (lhs->type() != ExpressionType::FIELD || rhs->type() != ExpressionType::SCALAR) {
     return Status::OK();
   }

   return visitor(checked_cast<const FieldExpression*>(lhs)->name(),
                  checked_cast<const ScalarExpression*>(rhs)->value());
 }

 Status KeyValuePartitioning::SetDefaultValuesFromKeys(const Expression& expr,
                                                       RecordBatchProjector* projector) {
   return KeyValuePartitioning::VisitKeys(
       expr, [projector](const std::string& name, const std::shared_ptr<Scalar>& value) {
         ARROW_ASSIGN_OR_RAISE(auto match,
                               FieldRef(name).FindOneOrNone(*projector->schema()));
         if (match.indices().empty()) {
           return Status::OK();
         }
         return projector->SetDefaultValue(match, value);
       });
 }

 Result<std::shared_ptr<Expression>> KeyValuePartitioning::ConvertKey(
     const Key& key, const Schema& schema) {
   ARROW_ASSIGN_OR_RAISE(auto field, FieldRef(key.name).GetOneOrNone(schema));
   if (field == nullptr) {
     return scalar(true);
   }

   ARROW_ASSIGN_OR_RAISE(auto converted, Scalar::Parse(field->type(), key.value));
   return equal(field_ref(field->name()), scalar(converted));
 }

 Result<std::shared_ptr<Expression>> KeyValuePartitioning::Parse(
     const std::string& segment, int i) const {
   if (auto key = ParseKey(segment, i)) {
     return ConvertKey(*key, *schema_);
   }

   return scalar(true);
 }

 Result<std::string> KeyValuePartitioning::Format(const Expression& expr, int i) const {
   if (expr.type() != ExpressionType::COMPARISON) {
     return Status::Invalid(expr.ToString(), " is not a comparison expression");
   }

   const auto& cmp = checked_cast<const ComparisonExpression&>(expr);
   if (cmp.op() != compute::CompareOperator::EQUAL) {
     return Status::Invalid(expr.ToString(), " is not an equality comparison expression");
   }

   if (cmp.left_operand()->type() != ExpressionType::FIELD) {
     return Status::Invalid(expr.ToString(), " LHS is not a field");
   }
   const auto& lhs = checked_cast<const FieldExpression&>(*cmp.left_operand());

   if (cmp.right_operand()->type() != ExpressionType::SCALAR) {
     return Status::Invalid(expr.ToString(), " RHS is not a scalar");
   }
   const auto& rhs = checked_cast<const ScalarExpression&>(*cmp.right_operand());

   auto expected_type = schema_->GetFieldByName(lhs.name())->type();
   if (!rhs.value()->type->Equals(expected_type)) {
     return Status::TypeError(expr.ToString(), " expected RHS to have type ",
                              *expected_type);
   }

   return FormatKey({lhs.name(), rhs.value()->ToString()}, i);
 }

 util::optional<KeyValuePartitioning::Key> DirectoryPartitioning::ParseKey(
     const std::string& segment, int i) const {
   if (i >= schema_->num_fields()) {
     return util::nullopt;
   }

   return Key{schema_->field(i)->name(), segment};
 }

 Result<std::string> DirectoryPartitioning::FormatKey(const Key& key, int i) const {
   if (schema_->GetFieldIndex(key.name) != i) {
     return Status::Invalid("field ", key.name, " in unexpected position ", i,
                            " for schema ", *schema_);
   }
   return key.value;
 }

 class KeyValuePartitioningInspectImpl {
  public:
   static Result<std::shared_ptr<DataType>> InferType(
       const std::string& name, const std::vector<std::string>& reprs) {
     if (reprs.empty()) {
       return Status::Invalid("No segments were available for field '", name,
                              "'; couldn't infer type");
     }

     bool all_integral = std::all_of(reprs.begin(), reprs.end(), [](string_view repr) {
       // TODO(bkietz) use ParseUnsigned or so
       return repr.find_first_not_of("0123456789") == string_view::npos;
     });

     if (all_integral) {
       return int32();
     }

     return utf8();
   }

   int GetOrInsertField(const std::string& name) {
     auto name_index =
         name_to_index_.emplace(name, static_cast<int>(name_to_index_.size())).first;

     if (static_cast<size_t>(name_index->second) >= values_.size()) {
       values_.resize(name_index->second + 1);
     }
     return name_index->second;
   }

   void InsertRepr(const std::string& name, std::string repr) {
     InsertRepr(GetOrInsertField(name), std::move(repr));
   }

   void InsertRepr(int index, std::string repr) {
     values_[index].push_back(std::move(repr));
   }

   Result<std::shared_ptr<Schema>> Finish() {
     std::vector<std::shared_ptr<Field>> fields(name_to_index_.size());

     for (const auto& name_index : name_to_index_) {
       const auto& name = name_index.first;
       auto index = name_index.second;
       ARROW_ASSIGN_OR_RAISE(auto type, InferType(name, values_[index]));
       fields[index] = field(name, type);
     }

     return ::arrow::schema(std::move(fields));
   }

  private:
   std::unordered_map<std::string, int> name_to_index_;
   std::vector<std::vector<std::string>> values_;
 };

 class DirectoryPartitioningFactory : public PartitioningFactory {
  public:
   explicit DirectoryPartitioningFactory(std::vector<std::string> field_names)
       : field_names_(std::move(field_names)) {}

   std::string type_name() const override { return "schema"; }

   Result<std::shared_ptr<Schema>> Inspect(
       const std::vector<string_view>& paths) const override {
     KeyValuePartitioningInspectImpl impl;

     for (const auto& name : field_names_) {
       impl.GetOrInsertField(name);
     }

     for (auto path : paths) {
       size_t field_index = 0;
       for (auto&& segment : fs::internal::SplitAbstractPath(path.to_string())) {
         if (field_index == field_names_.size()) break;

         impl.InsertRepr(static_cast<int>(field_index++), std::move(segment));
       }
     }

     return impl.Finish();
   }

   Result<std::shared_ptr<Partitioning>> Finish(
       const std::shared_ptr<Schema>& schema) const override {
     for (FieldRef ref : field_names_) {
       // ensure all of field_names_ are present in schema
       RETURN_NOT_OK(ref.FindOne(*schema).status());
     }

     // drop fields which aren't in field_names_
     auto out_schema = SchemaFromColumnNames(schema, field_names_);

     return std::make_shared<DirectoryPartitioning>(std::move(out_schema));
   }

   struct MakeWritePlanImpl;

   Result<WritePlan> MakeWritePlan(FragmentIterator fragments) override;

   Result<WritePlan> MakeWritePlan(FragmentIterator fragments,
                                   const std::shared_ptr<Schema>& schema) override;

  private:
   std::vector<std::string> field_names_;
 };

 struct DirectoryPartitioningFactory::MakeWritePlanImpl {
   using Indices = std::basic_string<int>;

   MakeWritePlanImpl(DirectoryPartitioningFactory* factory,
                     FragmentVector source_fragments)
       : this_(factory),
         source_fragments_(std::move(source_fragments)),
         right_hand_sides_(source_fragments_.size(), Indices(num_fields(), -1)) {}

   int num_fields() const { return static_cast<int>(this_->field_names_.size()); }

   // For a KeyValuePartitioning, every partition expression will be an equality
   // ComparisonExpression where the left operand is a FieldExpression and the right is a
   // ScalarExpression. Comparing Scalars directly is expensive, so first assemble a
   // dictionary containing the scalars from the right operands of every partition
   // expression. This allows later stages of MakeWritePlan to handle a scalar by its
   // dictionary code, which is both more compact to store and cheap to compare.
   //
   // Scalars are stored such that the dictionary code of a fragment's RHS in the
   // partition expression for a given field is given by
   //     int code = right_hand_sides_[fragment_index][field_index];
   // and the corresponding scalar can be retrieved with
   //     std::shared_ptr<Scalar> scalar = scalar_dict_.code_to_scalar[code];
   Status DictEncodeRightHandSides() {
     if (source_fragments_.empty()) {
       return Status::OK();
     }

     for (size_t fragment_i = 0; fragment_i < source_fragments_.size(); ++fragment_i) {
       const auto& fragment = source_fragments_[fragment_i];

       auto insert_representable_into_dict = [this, fragment_i](
                                                 const std::string& name,
                                                 const std::shared_ptr<Scalar>& value) {
         auto it = std::find(this_->field_names_.begin(), this_->field_names_.end(), name);
         if (it == this_->field_names_.end()) {
           return Status::OK();
         }

         auto field_i = it - this_->field_names_.begin();

         int code = scalar_dict_.GetOrInsert(value);
         right_hand_sides_[fragment_i][field_i] = code;

         return Status::OK();
       };

       RETURN_NOT_OK(KeyValuePartitioning::VisitKeys(*fragment->partition_expression(),
                                                     insert_representable_into_dict));

       auto it = std::find(right_hand_sides_[fragment_i].begin(),
                           right_hand_sides_[fragment_i].end(), -1);
       if (it != right_hand_sides_[fragment_i].end()) {
         // NB: this is an error when writing DirectoryPartitioning but not
         // HivePartitioning (as it will be valid to simply omit segments)
         return Status::Invalid(
             "fragment ", fragment_i, " had no partition expression for field '",
             this_->field_names_.at(it - right_hand_sides_[fragment_i].begin()), "'");
       }
     }

     return Status::OK();
   }

   // Infer the Partitioning schema from partition expressions.
   // For example if one partition expression is "omega"_ == 13
   // we can infer that the field "omega" has type int32
   Result<std::shared_ptr<Schema>> InferPartitioningSchema() const {
     if (source_fragments_.empty()) {
       return Status::Invalid(
           "No fragments were provided so the Partitioning schema could not be "
           "inferred.");
     }

     // NB: under DirectoryPartitioning every fragment has a partition expression for every
     // field, so we can infer the schema by looking only at the first fragment. This will
     // be more complicated for HivePartitioning.
     int fragment_i = 0;

     FieldVector fields(num_fields());
     for (int field_i = 0; field_i < num_fields(); ++field_i) {
       const auto& name = this_->field_names_[field_i];
       const auto& type =
           scalar_dict_.code_to_scalar[right_hand_sides_[fragment_i][field_i]]->type;
       fields[field_i] = field(name, type);
     }

     return schema(std::move(fields));
   }

   // reconstitute fragment_i's partition expression for field_i by reading the right
   // hand side from the scalar dictionary and constructing an equality
   // ComparisonExpression
   std::shared_ptr<Expression> PartitionExpression(size_t fragment_i, int field_i) {
     auto left_hand_side = field_ref(this_->field_names_[field_i]);
     auto right_hand_side =
         scalar(scalar_dict_.code_to_scalar[right_hand_sides_[fragment_i][field_i]]);
     return equal(std::move(left_hand_side), std::move(right_hand_side));
   }

   // create a guid by stringifying the number of milliseconds since the epoch
   std::string Guid() {
     using std::chrono::duration_cast;
     using std::chrono::milliseconds;
     using std::chrono::steady_clock;
     auto milliseconds_since_epoch =
         duration_cast<milliseconds>(steady_clock::now().time_since_epoch()).count();
     return std::to_string(milliseconds_since_epoch);
   }

   // remove fields which will be implicit in the partitioning; writing them to files would
   // be redundant
   Status DropPartitionFields(const std::shared_ptr<Partitioning>& partitioning,
                              Fragment* fragment) {
     auto schema = fragment->schema();
     for (const auto& field : partitioning->schema()->fields()) {
       int field_i = schema->GetFieldIndex(field->name());
       if (field_i != -1) {
         ARROW_ASSIGN_OR_RAISE(schema, schema->RemoveField(field_i));
       }
     }

     // the fragment being scanned to disk will now deselect redundant columns
     fragment->scan_options()->projector = RecordBatchProjector(std::move(schema));
     return Status::OK();
   }

   Result<WritePlan> Finish(std::shared_ptr<Schema> partitioning_schema = nullptr) && {
     WritePlan out;

     RETURN_NOT_OK(DictEncodeRightHandSides());

     if (partitioning_schema == nullptr) {
       ARROW_ASSIGN_OR_RAISE(partitioning_schema, InferPartitioningSchema());
     }
     ARROW_ASSIGN_OR_RAISE(out.partitioning,
                           this_->Finish(std::move(partitioning_schema)));

     auto fragment_schema =
         source_fragments_.empty() ? schema({}) : source_fragments_.front()->schema();
     ARROW_ASSIGN_OR_RAISE(out.schema,
                           UnifySchemas({out.partitioning->schema(), fragment_schema}));

     // Lexicographic ordering WRT right_hand_sides_ ensures that source_fragments_ are in
     // a depth first visitation order WRT their partition expressions. This makes
     // generation of the full directory tree far simpler since a directory's files are
     // grouped.
     auto permutation = internal::ArgSort(right_hand_sides_);
     internal::Permute(permutation, &source_fragments_);
     internal::Permute(permutation, &right_hand_sides_);

     // the basename of out.paths[i] is stored in segments[i] (full paths will be assembled
     // after segments is complete)
     std::vector<std::string> segments;

     // out.paths[parents[i]] is the parent directory of out.paths[i]
     std::vector<int> parents;

     // current_right_hand_sides[field_i] is the RHS dictionary code for the current
     // partition expression corresponding to field_i
     Indices current_right_hand_sides(num_fields(), -1);

     // out.paths[current_parents[field_i]] is the current ancestor directory corresponding
     // to field_i
     Indices current_parents(num_fields() + 1, -1);

     for (size_t fragment_i = 0; fragment_i < source_fragments_.size(); ++fragment_i) {
       RETURN_NOT_OK(
           DropPartitionFields(out.partitioning, source_fragments_[fragment_i].get()));

       int field_i = 0;
       for (; field_i < num_fields(); ++field_i) {
         // these directories have already been created and we're still writing their
         // children
         if (right_hand_sides_[fragment_i][field_i] != current_right_hand_sides[field_i]) {
           break;
         }
       }

       for (; field_i < num_fields(); ++field_i) {
         // push a new directory
         current_parents[field_i + 1] = static_cast<int>(parents.size());
         parents.push_back(current_parents[field_i]);

         auto partition_expression = PartitionExpression(fragment_i, field_i);

         // format segment for partition_expression
         ARROW_ASSIGN_OR_RAISE(auto segment,
                               out.partitioning->Format(*partition_expression, field_i));
         segment.push_back(fs::internal::kSep);
         segments.push_back(std::move(segment));

         // store partition_expression for use in the written Dataset
         out.fragment_or_partition_expressions.emplace_back(
             std::move(partition_expression));

         current_right_hand_sides[field_i] = right_hand_sides_[fragment_i][field_i];
       }

       // push a fragment (not attempting to give files meaningful names)
       parents.push_back(current_parents[field_i]);
       segments.emplace_back(Guid() + "_" + std::to_string(fragment_i));

       // store a fragment for writing to disk
       out.fragment_or_partition_expressions.emplace_back(
           std::move(source_fragments_[fragment_i]));
     }

     // render paths from segments
     for (size_t i = 0; i < segments.size(); ++i) {
       if (parents[i] == -1) {
         out.paths.push_back(segments[i]);
         continue;
       }

       out.paths.push_back(out.paths[parents[i]] + segments[i]);
     }

     return out;
   }

   DirectoryPartitioningFactory* this_;
   FragmentVector source_fragments_;

   struct {
     std::unordered_map<std::shared_ptr<Scalar>, int, Scalar::Hash, Scalar::PtrsEqual>
         scalar_to_code;

     ScalarVector code_to_scalar;

     int GetOrInsert(const std::shared_ptr<Scalar>& scalar) {
       int new_code = static_cast<int>(code_to_scalar.size());

       auto it_inserted = scalar_to_code.emplace(scalar, new_code);
       if (!it_inserted.second) {
         return it_inserted.first->second;
       }

       code_to_scalar.push_back(scalar);
       return new_code;
     }
   } scalar_dict_;
   std::vector<Indices> right_hand_sides_;
 };

 Result<WritePlan> DirectoryPartitioningFactory::MakeWritePlan(
     FragmentIterator fragment_it, const std::shared_ptr<Schema>& schema) {
   ARROW_ASSIGN_OR_RAISE(auto fragments, fragment_it.ToVector());
   return MakeWritePlanImpl(this, std::move(fragments)).Finish(schema);
 }

 Result<WritePlan> DirectoryPartitioningFactory::MakeWritePlan(
     FragmentIterator fragment_it) {
   ARROW_ASSIGN_OR_RAISE(auto fragments, fragment_it.ToVector());
   return MakeWritePlanImpl(this, std::move(fragments)).Finish();
 }

 std::shared_ptr<PartitioningFactory> DirectoryPartitioning::MakeFactory(
     std::vector<std::string> field_names) {
   return std::shared_ptr<PartitioningFactory>(
       new DirectoryPartitioningFactory(std::move(field_names)));
 }

 util::optional<KeyValuePartitioning::Key> HivePartitioning::ParseKey(
     const std::string& segment) {
   auto name_end = string_view(segment).find_first_of('=');
   if (name_end == string_view::npos) {
     return util::nullopt;
   }

   return Key{segment.substr(0, name_end), segment.substr(name_end + 1)};
 }

 Result<std::string> HivePartitioning::FormatKey(const Key& key, int i) const {
   return key.name + "=" + key.value;
 }

 class HivePartitioningFactory : public PartitioningFactory {
  public:
   std::string type_name() const override { return "hive"; }

   Result<std::shared_ptr<Schema>> Inspect(
       const std::vector<string_view>& paths) const override {
     KeyValuePartitioningInspectImpl impl;

     for (auto path : paths) {
       for (auto&& segment : fs::internal::SplitAbstractPath(path.to_string())) {
         if (auto key = HivePartitioning::ParseKey(segment)) {
           impl.InsertRepr(key->name, key->value);
         }
       }
     }

     return impl.Finish();
   }

   Result<std::shared_ptr<Partitioning>> Finish(
       const std::shared_ptr<Schema>& schema) const override {
     return std::shared_ptr<Partitioning>(new HivePartitioning(schema));
   }
 };

 std::shared_ptr<PartitioningFactory> HivePartitioning::MakeFactory() {
   return std::shared_ptr<PartitioningFactory>(new HivePartitioningFactory());
 }

 }  // namespace dataset
 }  // namespace arrow
	// 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.

	#include "arrow/dataset/partition.h"

	#include <algorithm>
	#include <chrono>
	#include <map>
	#include <memory>
	#include <stack>
	#include <utility>
	#include <vector>

	#include "arrow/dataset/dataset_internal.h"
	#include "arrow/dataset/file_base.h"
	#include "arrow/dataset/filter.h"
	#include "arrow/dataset/scanner.h"
	#include "arrow/filesystem/filesystem.h"
	#include "arrow/filesystem/path_util.h"
	#include "arrow/scalar.h"
	#include "arrow/util/iterator.h"
	#include "arrow/util/range.h"
	#include "arrow/util/sort.h"
	#include "arrow/util/string_view.h"

	namespace arrow {
	namespace dataset {

	using util::string_view;

	using internal::checked_cast;

	Result<std::shared_ptr<Expression>> Partitioning::Parse(const std::string& path) const {
	ExpressionVector expressions;
	int i = 0;

	for (auto segment : fs::internal::SplitAbstractPath(path)) {
	ARROW_ASSIGN_OR_RAISE(auto expr, Parse(segment, i++));
	if (expr->Equals(true)) {
	continue;
	}

	expressions.push_back(std::move(expr));
	}

	return and_(std::move(expressions));
	}

	std::shared_ptr<Partitioning> Partitioning::Default() {
	return std::make_shared<DefaultPartitioning>();
	}

	Result<WritePlan> PartitioningFactory::MakeWritePlan(FragmentIterator fragment_it) {
	return Status::NotImplemented("MakeWritePlan from PartitioningFactory of type ",
	type_name());
	}

	Result<WritePlan> PartitioningFactory::MakeWritePlan(
	FragmentIterator fragment_it, const std::shared_ptr<Schema>& schema) {
	return Status::NotImplemented("MakeWritePlan from PartitioningFactory of type ",
	type_name());
	}

	Result<std::shared_ptr<Expression>> SegmentDictionaryPartitioning::Parse(
	const std::string& segment, int i) const {
	if (static_cast<size_t>(i) < dictionaries_.size()) {
	auto it = dictionaries_[i].find(segment);
	if (it != dictionaries_[i].end()) {
	return it->second;
	}
	}

	return scalar(true);
	}

	Status KeyValuePartitioning::VisitKeys(
	const Expression& expr,
	const std::function<Status(const std::string& name,
	const std::shared_ptr<Scalar>& value)>& visitor) {
	if (expr.type() == ExpressionType::AND) {
	const auto& and_ = checked_cast<const AndExpression&>(expr);
	RETURN_NOT_OK(VisitKeys(*and_.left_operand(), visitor));
	RETURN_NOT_OK(VisitKeys(*and_.right_operand(), visitor));
	return Status::OK();
	}

	if (expr.type() != ExpressionType::COMPARISON) {
	return Status::OK();
	}

	const auto& cmp = checked_cast<const ComparisonExpression&>(expr);
	if (cmp.op() != compute::CompareOperator::EQUAL) {
	return Status::OK();
	}

	auto lhs = cmp.left_operand().get();
	auto rhs = cmp.right_operand().get();
	if (lhs->type() != ExpressionType::FIELD) std::swap(lhs, rhs);

	if (lhs->type() != ExpressionType::FIELD \|\| rhs->type() != ExpressionType::SCALAR) {
	return Status::OK();
	}

	return visitor(checked_cast<const FieldExpression*>(lhs)->name(),
	checked_cast<const ScalarExpression*>(rhs)->value());
	}

	Status KeyValuePartitioning::SetDefaultValuesFromKeys(const Expression& expr,
	RecordBatchProjector* projector) {
	return KeyValuePartitioning::VisitKeys(
	expr, [projector](const std::string& name, const std::shared_ptr<Scalar>& value) {
	ARROW_ASSIGN_OR_RAISE(auto match,
	FieldRef(name).FindOneOrNone(*projector->schema()));
	if (match.indices().empty()) {
	return Status::OK();
	}
	return projector->SetDefaultValue(match, value);
	});
	}

	Result<std::shared_ptr<Expression>> KeyValuePartitioning::ConvertKey(
	const Key& key, const Schema& schema) {
	ARROW_ASSIGN_OR_RAISE(auto field, FieldRef(key.name).GetOneOrNone(schema));
	if (field == nullptr) {
	return scalar(true);
	}

	ARROW_ASSIGN_OR_RAISE(auto converted, Scalar::Parse(field->type(), key.value));
	return equal(field_ref(field->name()), scalar(converted));
	}

	Result<std::shared_ptr<Expression>> KeyValuePartitioning::Parse(
	const std::string& segment, int i) const {
	if (auto key = ParseKey(segment, i)) {
	return ConvertKey(key, schema_);
	}

	return scalar(true);
	}

	Result<std::string> KeyValuePartitioning::Format(const Expression& expr, int i) const {
	if (expr.type() != ExpressionType::COMPARISON) {
	return Status::Invalid(expr.ToString(), " is not a comparison expression");
	}

	const auto& cmp = checked_cast<const ComparisonExpression&>(expr);
	if (cmp.op() != compute::CompareOperator::EQUAL) {
	return Status::Invalid(expr.ToString(), " is not an equality comparison expression");
	}

	if (cmp.left_operand()->type() != ExpressionType::FIELD) {
	return Status::Invalid(expr.ToString(), " LHS is not a field");
	}
	const auto& lhs = checked_cast<const FieldExpression&>(*cmp.left_operand());

	if (cmp.right_operand()->type() != ExpressionType::SCALAR) {
	return Status::Invalid(expr.ToString(), " RHS is not a scalar");
	}
	const auto& rhs = checked_cast<const ScalarExpression&>(*cmp.right_operand());

	auto expected_type = schema_->GetFieldByName(lhs.name())->type();
	if (!rhs.value()->type->Equals(expected_type)) {
	return Status::TypeError(expr.ToString(), " expected RHS to have type ",
	*expected_type);
	}

	return FormatKey({lhs.name(), rhs.value()->ToString()}, i);
	}

	util::optional<KeyValuePartitioning::Key> DirectoryPartitioning::ParseKey(
	const std::string& segment, int i) const {
	if (i >= schema_->num_fields()) {
	return util::nullopt;
	}

	return Key{schema_->field(i)->name(), segment};
	}

	Result<std::string> DirectoryPartitioning::FormatKey(const Key& key, int i) const {
	if (schema_->GetFieldIndex(key.name) != i) {
	return Status::Invalid("field ", key.name, " in unexpected position ", i,
	" for schema ", *schema_);
	}
	return key.value;
	}

	class KeyValuePartitioningInspectImpl {
	public:
	static Result<std::shared_ptr<DataType>> InferType(
	const std::string& name, const std::vector<std::string>& reprs) {
	if (reprs.empty()) {
	return Status::Invalid("No segments were available for field '", name,
	"'; couldn't infer type");
	}

	bool all_integral = std::all_of(reprs.begin(), reprs.end(), [](string_view repr) {
	// TODO(bkietz) use ParseUnsigned or so
	return repr.find_first_not_of("0123456789") == string_view::npos;
	});

	if (all_integral) {
	return int32();
	}

	return utf8();
	}

	int GetOrInsertField(const std::string& name) {
	auto name_index =
	name_to_index_.emplace(name, static_cast<int>(name_to_index_.size())).first;

	if (static_cast<size_t>(name_index->second) >= values_.size()) {
	values_.resize(name_index->second + 1);
	}
	return name_index->second;
	}

	void InsertRepr(const std::string& name, std::string repr) {
	InsertRepr(GetOrInsertField(name), std::move(repr));
	}

	void InsertRepr(int index, std::string repr) {
	values_[index].push_back(std::move(repr));
	}

	Result<std::shared_ptr<Schema>> Finish() {
	std::vector<std::shared_ptr<Field>> fields(name_to_index_.size());

	for (const auto& name_index : name_to_index_) {
	const auto& name = name_index.first;
	auto index = name_index.second;
	ARROW_ASSIGN_OR_RAISE(auto type, InferType(name, values_[index]));
	fields[index] = field(name, type);
	}

	return ::arrow::schema(std::move(fields));
	}

	private:
	std::unordered_map<std::string, int> name_to_index_;
	std::vector<std::vector<std::string>> values_;
	};

	class DirectoryPartitioningFactory : public PartitioningFactory {
	public:
	explicit DirectoryPartitioningFactory(std::vector<std::string> field_names)
	: field_names_(std::move(field_names)) {}

	std::string type_name() const override { return "schema"; }

	Result<std::shared_ptr<Schema>> Inspect(
	const std::vector<string_view>& paths) const override {
	KeyValuePartitioningInspectImpl impl;

	for (const auto& name : field_names_) {
	impl.GetOrInsertField(name);
	}

	for (auto path : paths) {
	size_t field_index = 0;
	for (auto&& segment : fs::internal::SplitAbstractPath(path.to_string())) {
	if (field_index == field_names_.size()) break;

	impl.InsertRepr(static_cast<int>(field_index++), std::move(segment));
	}
	}

	return impl.Finish();
	}

	Result<std::shared_ptr<Partitioning>> Finish(
	const std::shared_ptr<Schema>& schema) const override {
	for (FieldRef ref : field_names_) {
	// ensure all of field_names_ are present in schema
	RETURN_NOT_OK(ref.FindOne(*schema).status());
	}

	// drop fields which aren't in field_names_
	auto out_schema = SchemaFromColumnNames(schema, field_names_);

	return std::make_shared<DirectoryPartitioning>(std::move(out_schema));
	}

	struct MakeWritePlanImpl;

	Result<WritePlan> MakeWritePlan(FragmentIterator fragments) override;

	Result<WritePlan> MakeWritePlan(FragmentIterator fragments,
	const std::shared_ptr<Schema>& schema) override;

	private:
	std::vector<std::string> field_names_;
	};

	struct DirectoryPartitioningFactory::MakeWritePlanImpl {
	using Indices = std::basic_string<int>;

	MakeWritePlanImpl(DirectoryPartitioningFactory* factory,
	FragmentVector source_fragments)
	: this_(factory),
	source_fragments_(std::move(source_fragments)),
	right_hand_sides_(source_fragments_.size(), Indices(num_fields(), -1)) {}

	int num_fields() const { return static_cast<int>(this_->field_names_.size()); }

	// For a KeyValuePartitioning, every partition expression will be an equality
	// ComparisonExpression where the left operand is a FieldExpression and the right is a
	// ScalarExpression. Comparing Scalars directly is expensive, so first assemble a
	// dictionary containing the scalars from the right operands of every partition
	// expression. This allows later stages of MakeWritePlan to handle a scalar by its
	// dictionary code, which is both more compact to store and cheap to compare.
	//
	// Scalars are stored such that the dictionary code of a fragment's RHS in the
	// partition expression for a given field is given by
	// int code = right_hand_sides_[fragment_index][field_index];
	// and the corresponding scalar can be retrieved with
	// std::shared_ptr<Scalar> scalar = scalar_dict_.code_to_scalar[code];
	Status DictEncodeRightHandSides() {
	if (source_fragments_.empty()) {
	return Status::OK();
	}

	for (size_t fragment_i = 0; fragment_i < source_fragments_.size(); ++fragment_i) {
	const auto& fragment = source_fragments_[fragment_i];

	auto insert_representable_into_dict = [this, fragment_i](
	const std::string& name,
	const std::shared_ptr<Scalar>& value) {
	auto it = std::find(this_->field_names_.begin(), this_->field_names_.end(), name);
	if (it == this_->field_names_.end()) {
	return Status::OK();
	}

	auto field_i = it - this_->field_names_.begin();

	int code = scalar_dict_.GetOrInsert(value);
	right_hand_sides_[fragment_i][field_i] = code;

	return Status::OK();
	};

	RETURN_NOT_OK(KeyValuePartitioning::VisitKeys(*fragment->partition_expression(),
	insert_representable_into_dict));

	auto it = std::find(right_hand_sides_[fragment_i].begin(),
	right_hand_sides_[fragment_i].end(), -1);
	if (it != right_hand_sides_[fragment_i].end()) {
	// NB: this is an error when writing DirectoryPartitioning but not
	// HivePartitioning (as it will be valid to simply omit segments)
	return Status::Invalid(
	"fragment ", fragment_i, " had no partition expression for field '",
	this_->field_names_.at(it - right_hand_sides_[fragment_i].begin()), "'");
	}
	}

	return Status::OK();
	}

	// Infer the Partitioning schema from partition expressions.
	// For example if one partition expression is "omega"_ == 13
	// we can infer that the field "omega" has type int32
	Result<std::shared_ptr<Schema>> InferPartitioningSchema() const {
	if (source_fragments_.empty()) {
	return Status::Invalid(
	"No fragments were provided so the Partitioning schema could not be "
	"inferred.");
	}

	// NB: under DirectoryPartitioning every fragment has a partition expression for every
	// field, so we can infer the schema by looking only at the first fragment. This will
	// be more complicated for HivePartitioning.
	int fragment_i = 0;

	FieldVector fields(num_fields());
	for (int field_i = 0; field_i < num_fields(); ++field_i) {
	const auto& name = this_->field_names_[field_i];
	const auto& type =
	scalar_dict_.code_to_scalar[right_hand_sides_[fragment_i][field_i]]->type;
	fields[field_i] = field(name, type);
	}

	return schema(std::move(fields));
	}

	// reconstitute fragment_i's partition expression for field_i by reading the right
	// hand side from the scalar dictionary and constructing an equality
	// ComparisonExpression
	std::shared_ptr<Expression> PartitionExpression(size_t fragment_i, int field_i) {
	auto left_hand_side = field_ref(this_->field_names_[field_i]);
	auto right_hand_side =
	scalar(scalar_dict_.code_to_scalar[right_hand_sides_[fragment_i][field_i]]);
	return equal(std::move(left_hand_side), std::move(right_hand_side));
	}

	// create a guid by stringifying the number of milliseconds since the epoch
	std::string Guid() {
	using std::chrono::duration_cast;
	using std::chrono::milliseconds;
	using std::chrono::steady_clock;
	auto milliseconds_since_epoch =
	duration_cast<milliseconds>(steady_clock::now().time_since_epoch()).count();
	return std::to_string(milliseconds_since_epoch);
	}

	// remove fields which will be implicit in the partitioning; writing them to files would
	// be redundant
	Status DropPartitionFields(const std::shared_ptr<Partitioning>& partitioning,
	Fragment* fragment) {
	auto schema = fragment->schema();
	for (const auto& field : partitioning->schema()->fields()) {
	int field_i = schema->GetFieldIndex(field->name());
	if (field_i != -1) {
	ARROW_ASSIGN_OR_RAISE(schema, schema->RemoveField(field_i));
	}
	}

	// the fragment being scanned to disk will now deselect redundant columns
	fragment->scan_options()->projector = RecordBatchProjector(std::move(schema));
	return Status::OK();
	}

	Result<WritePlan> Finish(std::shared_ptr<Schema> partitioning_schema = nullptr) && {
	WritePlan out;

	RETURN_NOT_OK(DictEncodeRightHandSides());

	if (partitioning_schema == nullptr) {
	ARROW_ASSIGN_OR_RAISE(partitioning_schema, InferPartitioningSchema());
	}
	ARROW_ASSIGN_OR_RAISE(out.partitioning,
	this_->Finish(std::move(partitioning_schema)));

	auto fragment_schema =
	source_fragments_.empty() ? schema({}) : source_fragments_.front()->schema();
	ARROW_ASSIGN_OR_RAISE(out.schema,
	UnifySchemas({out.partitioning->schema(), fragment_schema}));

	// Lexicographic ordering WRT right_hand_sides_ ensures that source_fragments_ are in
	// a depth first visitation order WRT their partition expressions. This makes
	// generation of the full directory tree far simpler since a directory's files are
	// grouped.
	auto permutation = internal::ArgSort(right_hand_sides_);
	internal::Permute(permutation, &source_fragments_);
	internal::Permute(permutation, &right_hand_sides_);

	// the basename of out.paths[i] is stored in segments[i] (full paths will be assembled
	// after segments is complete)
	std::vector<std::string> segments;

	// out.paths[parents[i]] is the parent directory of out.paths[i]
	std::vector<int> parents;

	// current_right_hand_sides[field_i] is the RHS dictionary code for the current
	// partition expression corresponding to field_i
	Indices current_right_hand_sides(num_fields(), -1);

	// out.paths[current_parents[field_i]] is the current ancestor directory corresponding
	// to field_i
	Indices current_parents(num_fields() + 1, -1);

	for (size_t fragment_i = 0; fragment_i < source_fragments_.size(); ++fragment_i) {
	RETURN_NOT_OK(
	DropPartitionFields(out.partitioning, source_fragments_[fragment_i].get()));

	int field_i = 0;
	for (; field_i < num_fields(); ++field_i) {
	// these directories have already been created and we're still writing their
	// children
	if (right_hand_sides_[fragment_i][field_i] != current_right_hand_sides[field_i]) {
	break;
	}
	}

	for (; field_i < num_fields(); ++field_i) {
	// push a new directory
	current_parents[field_i + 1] = static_cast<int>(parents.size());
	parents.push_back(current_parents[field_i]);

	auto partition_expression = PartitionExpression(fragment_i, field_i);

	// format segment for partition_expression
	ARROW_ASSIGN_OR_RAISE(auto segment,
	out.partitioning->Format(*partition_expression, field_i));
	segment.push_back(fs::internal::kSep);
	segments.push_back(std::move(segment));

	// store partition_expression for use in the written Dataset
	out.fragment_or_partition_expressions.emplace_back(
	std::move(partition_expression));

	current_right_hand_sides[field_i] = right_hand_sides_[fragment_i][field_i];
	}

	// push a fragment (not attempting to give files meaningful names)
	parents.push_back(current_parents[field_i]);
	segments.emplace_back(Guid() + "_" + std::to_string(fragment_i));

	// store a fragment for writing to disk
	out.fragment_or_partition_expressions.emplace_back(
	std::move(source_fragments_[fragment_i]));
	}

	// render paths from segments
	for (size_t i = 0; i < segments.size(); ++i) {
	if (parents[i] == -1) {
	out.paths.push_back(segments[i]);
	continue;
	}

	out.paths.push_back(out.paths[parents[i]] + segments[i]);
	}

	return out;
	}

	DirectoryPartitioningFactory* this_;
	FragmentVector source_fragments_;

	struct {
	std::unordered_map<std::shared_ptr<Scalar>, int, Scalar::Hash, Scalar::PtrsEqual>
	scalar_to_code;

	ScalarVector code_to_scalar;

	int GetOrInsert(const std::shared_ptr<Scalar>& scalar) {
	int new_code = static_cast<int>(code_to_scalar.size());

	auto it_inserted = scalar_to_code.emplace(scalar, new_code);
	if (!it_inserted.second) {
	return it_inserted.first->second;
	}

	code_to_scalar.push_back(scalar);
	return new_code;
	}
	} scalar_dict_;
	std::vector<Indices> right_hand_sides_;
	};

	Result<WritePlan> DirectoryPartitioningFactory::MakeWritePlan(
	FragmentIterator fragment_it, const std::shared_ptr<Schema>& schema) {
	ARROW_ASSIGN_OR_RAISE(auto fragments, fragment_it.ToVector());
	return MakeWritePlanImpl(this, std::move(fragments)).Finish(schema);
	}

	Result<WritePlan> DirectoryPartitioningFactory::MakeWritePlan(
	FragmentIterator fragment_it) {
	ARROW_ASSIGN_OR_RAISE(auto fragments, fragment_it.ToVector());
	return MakeWritePlanImpl(this, std::move(fragments)).Finish();
	}

	std::shared_ptr<PartitioningFactory> DirectoryPartitioning::MakeFactory(
	std::vector<std::string> field_names) {
	return std::shared_ptr<PartitioningFactory>(
	new DirectoryPartitioningFactory(std::move(field_names)));
	}

	util::optional<KeyValuePartitioning::Key> HivePartitioning::ParseKey(
	const std::string& segment) {
	auto name_end = string_view(segment).find_first_of('=');
	if (name_end == string_view::npos) {
	return util::nullopt;
	}

	return Key{segment.substr(0, name_end), segment.substr(name_end + 1)};
	}

	Result<std::string> HivePartitioning::FormatKey(const Key& key, int i) const {
	return key.name + "=" + key.value;
	}

	class HivePartitioningFactory : public PartitioningFactory {
	public:
	std::string type_name() const override { return "hive"; }

	Result<std::shared_ptr<Schema>> Inspect(
	const std::vector<string_view>& paths) const override {
	KeyValuePartitioningInspectImpl impl;

	for (auto path : paths) {
	for (auto&& segment : fs::internal::SplitAbstractPath(path.to_string())) {
	if (auto key = HivePartitioning::ParseKey(segment)) {
	impl.InsertRepr(key->name, key->value);
	}
	}
	}

	return impl.Finish();
	}

	Result<std::shared_ptr<Partitioning>> Finish(
	const std::shared_ptr<Schema>& schema) const override {
	return std::shared_ptr<Partitioning>(new HivePartitioning(schema));
	}
	};

	std::shared_ptr<PartitioningFactory> HivePartitioning::MakeFactory() {
	return std::shared_ptr<PartitioningFactory>(new HivePartitioningFactory());
	}

	} // namespace dataset
	} // namespace arrow